From 2fde3a6a315c087294a27eab2a905640bf0cdcfc 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/base.py | 6 +++ src/unilab/base/backend/mujoco_backend.py | 6 +++ src/unilab/dr/dr_utils.py | 50 +++++++++++++------ .../envs/locomotion/common/dr_provider.py | 17 ++++++- src/unilab/envs/locomotion/g1/joystick.py | 25 +++++++++- src/unilab/envs/locomotion/g1/joystick_sac.py | 11 ++-- 8 files changed, 101 insertions(+), 21 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/base.py b/src/unilab/base/backend/base.py index 0eac7ef6a..73f4b9071 100644 --- a/src/unilab/base/backend/base.py +++ b/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 # # ------------------------------------------------------------------ # 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..054a9a40c 100644 --- a/src/unilab/dr/dr_utils.py +++ b/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()) diff --git a/src/unilab/envs/locomotion/common/dr_provider.py b/src/unilab/envs/locomotion/common/dr_provider.py index 018d61f17..dc17caa20 100644 --- a/src/unilab/envs/locomotion/common/dr_provider.py +++ b/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 ──────────────────────────── diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 3aa7d9d01..50b8e1849 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,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]: diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index cca619b86..7c79572af 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 @@ -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