feat: G1 per-joint kp/kd domain randomization (issue #129)

conf/offpolicy/task/sac/g1_sac/motrix.yaml

@@ -11,6 +11,10 @@ algo:
   algo_params:
     alpha_init: 0.001
     target_entropy_ratio: 0.0
+env:
+  domain_rand:
+    randomize_kp: false
+    randomize_kd: false
 reward:
   scales:
     tracking_lin_vel: 2.2

conf/ppo/task/g1_joystick/motrix.yaml

@@ -16,6 +16,9 @@ algo:
     entropy_coef: 5.0e-3
 env:
   iterations: 3
+  domain_rand:
+    randomize_kp: false
+    randomize_kd: false
   control_config:
     action_scale: 0.5
   commands:

src/unilab/base/backend/base.py

@@ -170,6 +170,12 @@ def get_physics_state(self) -> np.ndarray:
             f"{self.__class__.__name__} does not support physics-state playback"
         )
 
+    def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]:
+        """Return per-joint (kp, kd) arrays from the backend model."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support reading actuator gains"
+        )
+
     # ------------------------------------------------------------------ #
     # Base kinematics                                                      #
     # ------------------------------------------------------------------ #

src/unilab/base/backend/mujoco_backend.py

@@ -467,6 +467,12 @@ def _translate_reset_randomization(
 
         return translated or None
 
+    def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]:
+        """Return per-joint (kp, kd) arrays read from the current model state."""
+        kp = np.asarray(self._model.actuator_gainprm[:, 0], dtype=np.float64).copy()
+        kd = np.asarray(-self._model.actuator_biasprm[:, 2], dtype=np.float64).copy()
+        return kp, kd
+
     def _apply_position_actuator_gains_to_model(
         self,
         model,

src/unilab/dr/dr_utils.py

@@ -16,7 +16,13 @@ def zero_actions(num_reset: int, num_action: int) -> np.ndarray:
     return np.zeros((num_reset, num_action), dtype=get_global_dtype())
 
 
-def build_common_reset_randomization(env: Any, num_reset: int) -> ResetRandomizationPayload | None:
+def build_common_reset_randomization(
+    env: Any,
+    num_reset: int,
+    *,
+    base_kp: np.ndarray | None = None,
+    base_kd: np.ndarray | None = None,
+) -> ResetRandomizationPayload | None:
     domain_rand = getattr(env.cfg, "domain_rand", None)
     if domain_rand is None:
         return None
@@ -33,25 +39,41 @@ def build_common_reset_randomization(env: Any, num_reset: int) -> ResetRandomiza
         payload.base_com_offset = base_com_offset
 
     num_actuators = getattr(env, "_num_action", None)
-    if num_actuators is not None and getattr(domain_rand, "randomize_kp", False):
-        low, high = domain_rand.kp_multiplier_range
-        kp_multiplier = np.random.uniform(low, high, size=(num_reset, 1))
-        base_kp = float(env.cfg.control_config.Kp)
-        payload.kp = np.broadcast_to(base_kp * kp_multiplier, (num_reset, num_actuators)).copy()
-
-    if num_actuators is not None and getattr(domain_rand, "randomize_kd", False):
-        low, high = domain_rand.kd_multiplier_range
-        kd_multiplier = np.random.uniform(low, high, size=(num_reset, 1))
-        base_kd = float(env.cfg.control_config.Kd)
-        payload.kd = np.broadcast_to(base_kd * kd_multiplier, (num_reset, num_actuators)).copy()
+    need_kp = num_actuators is not None and getattr(domain_rand, "randomize_kp", False)
+    need_kd = num_actuators is not None and getattr(domain_rand, "randomize_kd", False)
+
+    if need_kp or need_kd:
+        assert num_actuators is not None
+
+        if need_kp:
+            kp = (
+                base_kp
+                if base_kp is not None
+                else np.full(num_actuators, float(env.cfg.control_config.Kp))
+            )
+            low, high = domain_rand.kp_multiplier_range
+            payload.kp = (kp * np.random.uniform(low, high, (num_reset, 1))).astype(np.float64)
+
+        if need_kd:
+            kd = (
+                base_kd
+                if base_kd is not None
+                else np.full(num_actuators, float(env.cfg.control_config.Kd))
+            )
+            low, high = domain_rand.kd_multiplier_range
+            payload.kd = (kd * np.random.uniform(low, high, (num_reset, 1))).astype(np.float64)
 
     return None if payload.is_empty() else payload
 
 
 def validate_common_reset_randomization(
-    env: Any, capabilities: DomainRandomizationCapabilities
+    env: Any,
+    capabilities: DomainRandomizationCapabilities,
+    *,
+    base_kp: np.ndarray | None = None,
+    base_kd: np.ndarray | None = None,
 ) -> frozenset[str]:
-    payload = build_common_reset_randomization(env, num_reset=1)
+    payload = build_common_reset_randomization(env, num_reset=1, base_kp=base_kp, base_kd=base_kd)
     if payload is None:
         return frozenset()
     return capabilities.get_unsupported_reset_terms(payload.requested_terms())

src/unilab/envs/locomotion/common/dr_provider.py

@@ -44,8 +44,18 @@ class LocomotionDRProvider(DomainRandomizationProvider):
 
     # ── shared methods ───────────────────────────────────────────────
 
+    def _get_base_actuator_gains(self, env: Any) -> tuple[np.ndarray | None, np.ndarray | None]:
+        """Return (base_kp, base_kd) for per-joint kp/kd domain randomization.
+
+        Override to provide per-joint gains cached at init time.
+        Returns ``(None, None)`` by default, which falls back to scalar
+        ``ControlConfig.Kp`` / ``ControlConfig.Kd``.
+        """
+        return None, None
+
     def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None:
-        validate_common_reset_randomization(env, capabilities)
+        base_kp, base_kd = self._get_base_actuator_gains(env)
+        validate_common_reset_randomization(env, capabilities, base_kp=base_kp, base_kd=base_kd)
         validate_interval_push_support(env, capabilities)
 
     def build_interval_randomization_plan(
@@ -87,12 +97,15 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan:
             "last_actions": zero_actions(num_reset, env._num_action),
         }
         info_updates.update(self._build_extra_info_updates(env, num_reset))
+        base_kp, base_kd = self._get_base_actuator_gains(env)
         return ResetPlan(
             env_ids=env_ids,
             qpos=qpos,
             qvel=qvel,
             info_updates=info_updates,
-            randomization=build_common_reset_randomization(env, num_reset),
+            randomization=build_common_reset_randomization(
+                env, num_reset, base_kp=base_kp, base_kd=base_kd
+            ),
         )
 
     # ── template: build_reset_observation ────────────────────────────

src/unilab/envs/locomotion/g1/joystick.py

@@ -22,6 +22,15 @@
 from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv
 
 
+@dataclass
+class G1DomainRandConfig(DomainRandConfig):
+    randomize_kp: bool = True
+    kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1])
+
+    randomize_kd: bool = True
+    kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1])
+
+
 @dataclass
 class InitState:
     pos = [0.0, 0.0, 0.754]
@@ -107,12 +116,19 @@ class G1JoystickPPOCfg(G1BaseCfg):
     init_state: InitState = field(default_factory=InitState)
     commands: Commands = field(default_factory=Commands)
     reward_config: RewardConfigPPO | None = None
-    domain_rand: DomainRandConfig = field(default_factory=DomainRandConfig)
+    domain_rand: G1DomainRandConfig = field(default_factory=G1DomainRandConfig)
     gait_phase_init_mode: str = "offset_phase"
     reset_base_qvel_limit: float = 0.5
 
 
 class G1JoystickDomainRandomizationProvider(LocomotionDRProvider):
+    def __init__(self, *, base_kp: np.ndarray | None = None, base_kd: np.ndarray | None = None):
+        self._base_kp = base_kp
+        self._base_kd = base_kd
+
+    def _get_base_actuator_gains(self, env: Any) -> tuple[np.ndarray | None, np.ndarray | None]:
+        return self._base_kp, self._base_kd
+
     def _get_qvel_limit(self, env: Any) -> float:
         return float(env.cfg.reset_base_qvel_limit)
 
@@ -173,7 +189,12 @@ def __init__(self, cfg: G1JoystickPPOCfg, num_envs=1, backend_type="mujoco"):
         self._upper_body_pose_weights = build_upper_body_pose_weights(self._reward_cfg.pose_weights)
 
         self._init_reward_functions()
-        self._init_domain_randomization(G1JoystickDomainRandomizationProvider())
+        if cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd:
+            base_kp, base_kd = backend.get_actuator_gains()
+            dr_provider = G1JoystickDomainRandomizationProvider(base_kp=base_kp, base_kd=base_kd)
+        else:
+            dr_provider = G1JoystickDomainRandomizationProvider()
+        self._init_domain_randomization(dr_provider)
 
     @property
     def obs_groups_spec(self) -> dict[str, int]:

src/unilab/envs/locomotion/g1/joystick_sac.py

@@ -14,10 +14,10 @@
 from unilab.base.dtype_config import get_global_dtype
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import Commands
-from unilab.envs.locomotion.common.domain_rand import DomainRandConfig
 from unilab.envs.locomotion.common.rewards import RewardContext
 from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv
 from unilab.envs.locomotion.g1.joystick import (
+    G1DomainRandConfig,
     G1JoystickDomainRandomizationProvider,
     G1JoystickPPO,
     InitState,
@@ -55,7 +55,7 @@ class G1JoystickSACCfg(G1BaseCfg):
     init_state: InitState = field(default_factory=InitState)
     commands: Commands = field(default_factory=Commands)
     control_config: ControlConfigSAC = field(default_factory=ControlConfigSAC)  # type: ignore[assignment]
-    domain_rand: DomainRandConfig = field(default_factory=DomainRandConfig)
+    domain_rand: G1DomainRandConfig = field(default_factory=G1DomainRandConfig)
     gait_phase_init_mode: str = "offset_phase"
     reset_base_qvel_limit: float = 0.5
 
@@ -96,7 +96,12 @@ def __init__(self, cfg: G1JoystickSACCfg, num_envs=1, backend_type="mujoco"):
         )
 
         self._init_reward_functions()
-        self._init_domain_randomization(G1JoystickDomainRandomizationProvider())
+        if cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd:
+            base_kp, base_kd = backend.get_actuator_gains()
+            dr_provider = G1JoystickDomainRandomizationProvider(base_kp=base_kp, base_kd=base_kd)
+        else:
+            dr_provider = G1JoystickDomainRandomizationProvider()
+        self._init_domain_randomization(dr_provider)
 
     def _compute_obs(
         self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel