From 546b291d9bcfc396c5b6743cfda854a29484051c Mon Sep 17 00:00:00 2001 From: limeng Date: Fri, 10 Apr 2026 07:10:23 +0000 Subject: [PATCH] feat: G1 per-joint kp/kd domain randomization (issue #129) --- conf/offpolicy/task/sac/g1_sac/motrix.yaml | 4 +++ conf/ppo/task/g1_joystick/motrix.yaml | 3 ++ src/unilab/base/backend/mujoco_backend.py | 6 ++++ src/unilab/dr/dr_utils.py | 30 ++++++++++++------- src/unilab/envs/locomotion/g1/joystick.py | 11 ++++++- src/unilab/envs/locomotion/g1/joystick_sac.py | 4 +-- 6 files changed, 44 insertions(+), 14 deletions(-) diff --git a/conf/offpolicy/task/sac/g1_sac/motrix.yaml b/conf/offpolicy/task/sac/g1_sac/motrix.yaml index 2a2a7279c..479133db5 100644 --- a/conf/offpolicy/task/sac/g1_sac/motrix.yaml +++ b/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 diff --git a/conf/ppo/task/g1_joystick/motrix.yaml b/conf/ppo/task/g1_joystick/motrix.yaml index 5e6d7b1d5..7ac8fb949 100644 --- a/conf/ppo/task/g1_joystick/motrix.yaml +++ b/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: diff --git a/src/unilab/base/backend/mujoco_backend.py b/src/unilab/base/backend/mujoco_backend.py index 10daf02bc..30d2f1898 100644 --- a/src/unilab/base/backend/mujoco_backend.py +++ b/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, diff --git a/src/unilab/dr/dr_utils.py b/src/unilab/dr/dr_utils.py index a80b38034..34e818ecf 100644 --- a/src/unilab/dr/dr_utils.py +++ b/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 diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 3aa7d9d01..9d0999df3 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/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 diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index cca619b86..662f5248e 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/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