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/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

@@ -33,17 +33,25 @@ 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
+        backend = getattr(env, "_backend", None)
+        if backend is not None and hasattr(backend, "get_actuator_gains"):
+            base_kp, base_kd = backend.get_actuator_gains()
+        else:
+            base_kp = np.full(num_actuators, float(env.cfg.control_config.Kp))
+            base_kd = np.full(num_actuators, float(env.cfg.control_config.Kd))
+
+        if need_kp:
+            low, high = domain_rand.kp_multiplier_range
+            payload.kp = (base_kp * np.random.uniform(low, high, (num_reset, 1))).astype(np.float64)
+
+        if need_kd:
+            low, high = domain_rand.kd_multiplier_range
+            payload.kd = (base_kd * np.random.uniform(low, high, (num_reset, 1))).astype(np.float64)
 
     return None if payload.is_empty() else payload
 

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,7 +116,7 @@ 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
 

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