From 7fe1e0fcec804a672121d96bd05419c2dd00dfe2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=A2=E6=9D=B0?= Date: Wed, 18 Mar 2026 20:12:02 +0800 Subject: [PATCH 1/5] feat: sync go1 motrix ppo/sac and g1 motrix ppo tuning, fix related bugs --- conf/offpolicy/reward/go1_sac_motrix.yaml | 15 +- conf/ppo/reward/g1_ppo_motrix.yaml | 25 +-- conf/ppo/reward/go1_ppo_motrix.yaml | 15 +- scripts/train_offpolicy.py | 114 +++++++++-- scripts/train_rsl_rl.py | 137 ++++++++++++- src/unilab/algos/torch/offpolicy/runner.py | 2 + src/unilab/algos/torch/offpolicy/worker.py | 99 +++++++++- src/unilab/envs/locomotion/g1/joystick.py | 181 ++++++++++++++++-- src/unilab/envs/locomotion/g1/joystick_sac.py | 5 +- src/unilab/envs/locomotion/g1/xml/g1.xml | 19 -- src/unilab/envs/locomotion/go1/joystick.py | 32 +++- tests/envs/test_env_configs.py | 12 ++ tests/scripts/test_train_scripts.py | 48 +++++ 13 files changed, 616 insertions(+), 88 deletions(-) diff --git a/conf/offpolicy/reward/go1_sac_motrix.yaml b/conf/offpolicy/reward/go1_sac_motrix.yaml index 62db4a2ea..3f5dee7ad 100644 --- a/conf/offpolicy/reward/go1_sac_motrix.yaml +++ b/conf/offpolicy/reward/go1_sac_motrix.yaml @@ -1,13 +1,12 @@ # @package _global_ reward: scales: - tracking_lin_vel: 1.2 - tracking_ang_vel: 0.3 - lin_vel_z: -6.0 - ang_vel_xy: -0.15 - base_height: -120.0 - action_rate: -0.006 - similar_to_default: -0.12 - contact: 0.28 + tracking_lin_vel: 1.0 + tracking_ang_vel: 0.2 + lin_vel_z: -5.0 + ang_vel_xy: -0.1 + base_height: -100.0 + action_rate: -0.005 + similar_to_default: -0.1 tracking_sigma: 0.25 base_height_target: 0.3 diff --git a/conf/ppo/reward/g1_ppo_motrix.yaml b/conf/ppo/reward/g1_ppo_motrix.yaml index 829d94e2a..c27836651 100644 --- a/conf/ppo/reward/g1_ppo_motrix.yaml +++ b/conf/ppo/reward/g1_ppo_motrix.yaml @@ -1,19 +1,22 @@ # @package _global_ reward: scales: - tracking_lin_vel: 2.4 - tracking_ang_vel: 0.3 - feet_phase: 1.2 - lin_vel_z: -1.2 - ang_vel_xy: -0.3 - base_height: -600.0 - orientation: -6.0 - action_rate: -0.012 - pose: -0.12 + tracking_lin_vel: 3.0 + tracking_ang_vel: 0.25 + forward_progress: 1.5 + under_speed: -1.0 + upper_body_pose: -0.05 + feet_phase: 0.5 + lin_vel_z: -1.0 + ang_vel_xy: -0.2 + base_height: -120.0 + orientation: -2.5 + action_rate: -0.005 + pose: -0.05 tracking_sigma: 0.25 gait_frequency: 1.5 feet_phase_swing_height: 0.09 feet_phase_tracking_sigma: 0.008 base_height_target: 0.754 - min_base_height: 0.55 - max_tilt_deg: 25.0 + min_base_height: 0.5 + max_tilt_deg: 35.0 diff --git a/conf/ppo/reward/go1_ppo_motrix.yaml b/conf/ppo/reward/go1_ppo_motrix.yaml index 62db4a2ea..3f5dee7ad 100644 --- a/conf/ppo/reward/go1_ppo_motrix.yaml +++ b/conf/ppo/reward/go1_ppo_motrix.yaml @@ -1,13 +1,12 @@ # @package _global_ reward: scales: - tracking_lin_vel: 1.2 - tracking_ang_vel: 0.3 - lin_vel_z: -6.0 - ang_vel_xy: -0.15 - base_height: -120.0 - action_rate: -0.006 - similar_to_default: -0.12 - contact: 0.28 + tracking_lin_vel: 1.0 + tracking_ang_vel: 0.2 + lin_vel_z: -5.0 + ang_vel_xy: -0.1 + base_height: -100.0 + action_rate: -0.005 + similar_to_default: -0.1 tracking_sigma: 0.25 base_height_target: 0.3 diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 279ede2df..d56f631b9 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -78,12 +78,95 @@ def resolve_checkpoint_path( return load_path, load_path_dir -def build_runner(algo_name: str, cfg: DictConfig): - """Build algorithm runner from unified Hydra config.""" +def extract_reset_obs(reset_result): + """Extract obs_dict from env.reset(...) across old/new return conventions.""" + if isinstance(reset_result, tuple): + if len(reset_result) == 2: + obs_out, _ = reset_result + return obs_out + if len(reset_result) == 3: + _, obs_out, _ = reset_result + return obs_out + raise ValueError(f"Unexpected env.reset return format: {type(reset_result)!r}") + + +def resolve_play_obs_dim(obs_groups_spec: dict[str, int]) -> int: + from unilab.utils.obs_utils import get_obs_dims + + obs_dim, _ = get_obs_dims(obs_groups_spec) + return obs_dim + + +def extract_play_obs(obs_dict): + from unilab.utils.obs_utils import split_obs_dict + + obs_out, _ = split_obs_dict(obs_dict) + return obs_out + + +_GO1_LEGACY_COMMAND_VEL_LIMIT = [[0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] +_GO1_LEGACY_REWARD_CONFIG = { + "scales": { + "tracking_lin_vel": 1.0, + "tracking_ang_vel": 0.2, + "lin_vel_z": -5.0, + "ang_vel_xy": -0.1, + "base_height": -100.0, + "action_rate": -0.005, + "similar_to_default": -0.1, + }, + "tracking_sigma": 0.25, + "base_height_target": 0.3, +} + + +def get_go1_motrix_legacy_offpolicy_profile( + algo_name: str, task_name: str, sim_backend: str +) -> dict | None: + if algo_name != "sac" or task_name != "Go1JoystickFlatTerrain" or sim_backend != "motrix": + return None + return { + "algo_overrides": { + "num_envs": 4096, + "max_iterations": 2000, + }, + "env_cfg_override": { + "legacy_motrix_profile": { + "enabled": True, + "command_vel_limit": _GO1_LEGACY_COMMAND_VEL_LIMIT, + }, + "reward_config": _GO1_LEGACY_REWARD_CONFIG, + }, + } + + +def build_go1_motrix_legacy_offpolicy_env_cfg_override( + algo_name: str, cfg: DictConfig +) -> dict | None: from unilab.utils.reward_utils import extract_reward_config env_cfg_override = extract_reward_config(cfg) + profile = get_go1_motrix_legacy_offpolicy_profile( + algo_name, cfg.training.task_name, cfg.training.sim_backend + ) + if profile is not None: + cfg.algo.num_envs = profile["algo_overrides"]["num_envs"] + cfg.algo.max_iterations = profile["algo_overrides"]["max_iterations"] + env_cfg_override.update(profile["env_cfg_override"]) + + return env_cfg_override + + +def resolve_sac_use_symmetry(cfg: DictConfig) -> bool: + """Symmetry augmentation currently depends on MuJoCo model APIs.""" + return bool(cfg.algo.use_symmetry and cfg.training.sim_backend == "mujoco") + + +def build_runner(algo_name: str, cfg: DictConfig): + """Build algorithm runner from unified Hydra config.""" + env_cfg_override = build_go1_motrix_legacy_offpolicy_env_cfg_override(algo_name, cfg) + if algo_name == "sac": from unilab.algos.torch.fast_sac.learner import FastSACLearner from unilab.algos.torch.fast_sac.runner import FastSACRunner @@ -94,17 +177,18 @@ def build_runner(algo_name: str, cfg: DictConfig): from unilab.algos.torch.offpolicy.multi_gpu_runner import MultiGPUOffPolicyRunner from unilab.base import registry from unilab.utils.algo_utils import ensure_registries - from unilab.utils.reward_utils import extract_reward_config ensure_registries() device = cfg.training.device or get_default_device() - env_cfg_override = extract_reward_config(cfg) env = registry.make( - cfg.training.task_name, num_envs=1, sim_backend=cfg.training.sim_backend, - env_cfg_override=env_cfg_override + cfg.training.task_name, + num_envs=1, + sim_backend=cfg.training.sim_backend, + env_cfg_override=env_cfg_override, ) assert env.action_space.shape from unilab.utils.obs_utils import get_obs_dims + obs_dim, privileged_dim = get_obs_dims(env.obs_groups_spec) action_dim = env.action_space.shape[0] env.close() @@ -180,7 +264,7 @@ def build_runner(algo_name: str, cfg: DictConfig): sync_collection=not cfg.training.no_sync_collection, env_steps_per_sync=cfg.training.env_steps_per_sync, sim_backend=cfg.training.sim_backend, - use_symmetry=cfg.algo.use_symmetry, + use_symmetry=resolve_sac_use_symmetry(cfg), ) if algo_name == "td3": @@ -231,22 +315,18 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> None: from unilab.utils import render_many from unilab.utils.algo_utils import build_actor - from unilab.utils.reward_utils import extract_reward_config - - env_cfg_override = extract_reward_config(cfg) + env_cfg_override = build_go1_motrix_legacy_offpolicy_env_cfg_override(algo_name, cfg) device = default_device(torch, cfg.training.device) print(f"Using device for play: {device}") - from unilab.utils.obs_utils import flatten_obs_dict - env = registry.make( cfg.training.task_name, num_envs=cfg.training.play_env_num, sim_backend=cfg.training.sim_backend, env_cfg_override=env_cfg_override, ) - obs_dim = sum(env.obs_groups_spec.values()) + obs_dim = resolve_play_obs_dim(env.obs_groups_spec) action_dim = env.action_space.shape[0] normalizer = None @@ -305,8 +385,8 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> None: if env.state is None: env.init_state() env_indices = np.arange(cfg.training.play_env_num, dtype=np.int32) - _, obs_out, _ = env.reset(env_indices) - obs_np = np.asarray(flatten_obs_dict(obs_out), dtype=np.float32) + obs_out = extract_reset_obs(env.reset(env_indices)) + obs_np = np.asarray(extract_play_obs(obs_out), dtype=np.float32) # Use Motrix native rendering if cfg.training.sim_backend == "motrix": @@ -331,7 +411,7 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> None: else actor(obs_torch).cpu().numpy() ) state = env.step(actions_np) - obs_np = np.asarray(flatten_obs_dict(state.obs), dtype=np.float32) + obs_np = np.asarray(extract_play_obs(state.obs), dtype=np.float32) current_time = time.perf_counter() elapsed = current_time - last_render_time @@ -363,7 +443,7 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> None: else actor(obs_torch).cpu().numpy() ) state = env.step(actions_np) - obs_np = np.asarray(flatten_obs_dict(state.obs), dtype=np.float32) + obs_np = np.asarray(extract_play_obs(state.obs), dtype=np.float32) state_list.append(np.asarray(env._backend.get_physics_state(), dtype=np.float32).copy()) print("Rendering frames...") diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index a1ab2c844..53ee1bce5 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -47,6 +47,129 @@ def ensure_registries(): pass +_GO1_LEGACY_COMMAND_VEL_LIMIT = [[0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] +_GO1_LEGACY_REWARD_CONFIG = { + "scales": { + "tracking_lin_vel": 1.0, + "tracking_ang_vel": 0.2, + "lin_vel_z": -5.0, + "ang_vel_xy": -0.1, + "base_height": -100.0, + "action_rate": -0.005, + "similar_to_default": -0.1, + }, + "tracking_sigma": 0.25, + "base_height_target": 0.3, +} + +_G1_MOTRIX_LEGACY_PPO_REWARD_CONFIG = { + "scales": { + "tracking_lin_vel": 3.0, + "tracking_ang_vel": 0.25, + "forward_progress": 1.5, + "under_speed": -1.0, + "upper_body_pose": -0.05, + "feet_phase": 0.5, + "lin_vel_z": -1.0, + "ang_vel_xy": -0.2, + "base_height": -120.0, + "orientation": -2.5, + "action_rate": -0.005, + "pose": -0.05, + }, + "tracking_sigma": 0.25, + "gait_frequency": 1.5, + "feet_phase_swing_height": 0.09, + "feet_phase_tracking_sigma": 0.008, + "base_height_target": 0.754, + "min_base_height": 0.50, + "max_tilt_deg": 35.0, +} + + +def get_go1_motrix_legacy_ppo_profile(task_name: str, sim_backend: str) -> dict | None: + if task_name != "Go1JoystickFlatTerrain" or sim_backend != "motrix": + return None + return { + "algo_overrides": { + "max_iterations": 151, + "empirical_normalization": True, + "policy": {"init_noise_std": 0.5}, + "algorithm": {"learning_rate": 3.0e-4, "entropy_coef": 1.0e-3}, + }, + "env_cfg_override": { + "legacy_motrix_profile": { + "enabled": True, + "command_vel_limit": _GO1_LEGACY_COMMAND_VEL_LIMIT, + }, + "reward_config": _GO1_LEGACY_REWARD_CONFIG, + }, + } + + +def get_g1_motrix_legacy_ppo_profile(task_name: str, sim_backend: str) -> dict | None: + if task_name != "G1JoystickFlatTerrain" or sim_backend != "motrix": + return None + return { + "algo_overrides": { + "max_iterations": 151, + "empirical_normalization": True, + "policy": {"init_noise_std": 0.5}, + "algorithm": {"learning_rate": 3.0e-4, "entropy_coef": 5.0e-3}, + }, + "env_cfg_override": { + "reward_config": _G1_MOTRIX_LEGACY_PPO_REWARD_CONFIG, + "backend_overrides": { + "enabled": True, + "control_action_scale": 0.5, + "command_vel_limit": [[0.45, 0.0, 0.0], [0.55, 0.0, 0.0]], + "gait_phase_init_mode": "independent", + "reset_base_qvel_limit": 0.05, + "reward_scale_overrides": { + "under_speed": -1.0, + "upper_body_pose": -0.05, + "tracking_ang_vel": 0.25, + }, + }, + }, + } + + +def build_go1_motrix_legacy_ppo_env_cfg_override(cfg: DictConfig) -> dict: + env_cfg_override = {} + if hasattr(cfg, "reward") and cfg.reward: + env_cfg_override["reward_config"] = OmegaConf.to_container(cfg.reward, resolve=True) + + profile = get_go1_motrix_legacy_ppo_profile(cfg.training.task_name, cfg.training.sim_backend) + if profile is not None: + cfg.algo.max_iterations = profile["algo_overrides"]["max_iterations"] + cfg.algo.empirical_normalization = profile["algo_overrides"]["empirical_normalization"] + cfg.algo.policy.init_noise_std = profile["algo_overrides"]["policy"]["init_noise_std"] + cfg.algo.algorithm.learning_rate = profile["algo_overrides"]["algorithm"][ + "learning_rate" + ] + cfg.algo.algorithm.entropy_coef = profile["algo_overrides"]["algorithm"]["entropy_coef"] + env_cfg_override.update(profile["env_cfg_override"]) + + return env_cfg_override + + +def build_g1_motrix_legacy_ppo_env_cfg_override(cfg: DictConfig) -> dict: + env_cfg_override = {} + profile = get_g1_motrix_legacy_ppo_profile(cfg.training.task_name, cfg.training.sim_backend) + if profile is not None: + cfg.algo.max_iterations = profile["algo_overrides"]["max_iterations"] + cfg.algo.empirical_normalization = profile["algo_overrides"]["empirical_normalization"] + cfg.algo.policy.init_noise_std = profile["algo_overrides"]["policy"]["init_noise_std"] + cfg.algo.algorithm.learning_rate = profile["algo_overrides"]["algorithm"][ + "learning_rate" + ] + cfg.algo.algorithm.entropy_coef = profile["algo_overrides"]["algorithm"]["entropy_coef"] + cfg.algo.obs_groups["actor"] = ["policy"] + env_cfg_override.update(profile["env_cfg_override"]) + return env_cfg_override + + class RslRlVecEnvWrapper: """Wrapper to adapt NpEnv to RSL-RL OnPolicyRunner interface.""" @@ -74,7 +197,11 @@ def _obs_to_tensordict(self, obs: dict[str, np.ndarray]) -> TensorDict: td = {"actor": actor} if "privileged" in obs: td["privileged"] = to_torch(obs["privileged"], self.device) - td["policy"] = to_torch(flatten_obs_dict(obs), self.device) + if hasattr(self.env, "get_policy_obs"): + policy_obs = self.env.get_policy_obs(obs) + else: + policy_obs = flatten_obs_dict(obs) + td["policy"] = to_torch(policy_obs, self.device) else: td["policy"] = actor return TensorDict(td, batch_size=self.num_envs, device=self.device) @@ -131,9 +258,9 @@ def play_rsl_rl(cfg: DictConfig, device: str): """Play mode for RSL-RL.""" from unilab.base import registry - from unilab.utils.reward_utils import extract_reward_config - env_cfg_override = extract_reward_config(cfg) + env_cfg_override = build_go1_motrix_legacy_ppo_env_cfg_override(cfg) + env_cfg_override.update(build_g1_motrix_legacy_ppo_env_cfg_override(cfg)) env = registry.make( cfg.training.task_name, @@ -263,9 +390,9 @@ def main(cfg: DictConfig) -> None: from omegaconf import OmegaConf from unilab.base import registry - from unilab.utils.reward_utils import extract_reward_config - env_cfg_override = extract_reward_config(cfg) + env_cfg_override = build_go1_motrix_legacy_ppo_env_cfg_override(cfg) + env_cfg_override.update(build_g1_motrix_legacy_ppo_env_cfg_override(cfg)) if torch.cuda.is_available(): device = "cuda" diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/torch/offpolicy/runner.py index 65fff3042..22846d554 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/torch/offpolicy/runner.py @@ -137,6 +137,8 @@ def learn( "obs_normalization": self.obs_normalization, "shared_obs_normalizer_stats": shared_obs_normalizer_stats, "env_cfg_override": self.env_cfg_override, + "obs_dim": self.obs_dim, + "action_dim": self.action_dim, } self._start_collector( target_fn=off_policy_collector_fn, diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index 1805e2dfd..322624c9a 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -14,6 +14,66 @@ from unilab.utils.obs_utils import split_obs_dict +def resolve_collector_actor_dims( + env, + obs_dim: int | None = None, + action_dim: int | None = None, +) -> tuple[int, int]: + """Resolve actor dims for the collector. + + Prefer explicit dims from the parent process so learner and collector + build identical actor shapes on legacy / override-heavy env paths. + """ + if obs_dim is None: + from unilab.utils.obs_utils import get_obs_dims + + obs_dim, _ = get_obs_dims(env.obs_groups_spec) + + if action_dim is None: + assert env.action_space.shape is not None + action_dim = env.action_space.shape[0] + + return obs_dim, action_dim + + +def maybe_adapt_legacy_go1_motrix_obs( + *, + env_name: str, + env_cfg_override: dict | None, + obs: np.ndarray, + privileged: np.ndarray | None, + expected_obs_dim: int, +) -> tuple[np.ndarray, np.ndarray | None]: + """Adapt legacy Go1 Motrix collector obs when the env still emits the default layout. + + Default Go1 obs is 49D and privileged linvel is 3D: + gyro(3) + gravity(3) + diff(12) + dof_vel(12) + last_actions(12) + command(3) + feet_phase(4) + + Legacy Motrix Go1 obs is 48D: + linvel(3) + gyro(3) + gravity(3) + diff(12) + dof_vel(12) + last_actions(12) + command(3) + + If the collector unexpectedly sees the default 49D obs while the runner + expects the legacy 48D layout, rebuild the exact legacy obs from + privileged linvel plus the first 45 default obs dims. + """ + if obs.shape[1] == expected_obs_dim: + return obs, privileged + + legacy_profile = (env_cfg_override or {}).get("legacy_motrix_profile", {}) + if ( + env_name == "Go1JoystickFlatTerrain" + and legacy_profile.get("enabled", False) + and expected_obs_dim == 48 + and obs.shape[1] == 49 + and privileged is not None + and privileged.shape[1] == 3 + ): + adapted_obs = np.concatenate([privileged, obs[:, :45]], axis=1).astype(np.float32) + return adapted_obs, None + + return obs, privileged + + def off_policy_collector_fn( stop_event, env_name: str, @@ -35,6 +95,8 @@ def off_policy_collector_fn( shared_obs_normalizer_stats=None, sim_backend: str = "mujoco", env_cfg_override: dict | None = None, + obs_dim: int | None = None, + action_dim: int | None = None, **kwargs, ): """Entry point for the off-policy collector subprocess.""" @@ -64,6 +126,8 @@ def off_policy_collector_fn( shared_obs_normalizer_stats=shared_obs_normalizer_stats, sim_backend=sim_backend, env_cfg_override=env_cfg_override, + obs_dim=obs_dim, + action_dim=action_dim, ) except Exception as e: print(f"[Collector] Exception: {e}", file=sys.stderr, flush=True) @@ -96,6 +160,8 @@ def _run_collector( shared_obs_normalizer_stats, sim_backend, env_cfg_override, + obs_dim, + action_dim, ): from unilab.base import registry from unilab.ipc import SharedWeightSync @@ -115,10 +181,11 @@ def _run_collector( ) # Build actor (always on CPU for env interaction) - from unilab.utils.obs_utils import get_obs_dims - obs_dim, _ = get_obs_dims(env.obs_groups_spec) - assert env.action_space.shape is not None - action_dim = env.action_space.shape[0] + obs_dim, action_dim = resolve_collector_actor_dims( + env, + obs_dim=obs_dim, + action_dim=action_dim, + ) actor = build_actor( algo_type, obs_dim, action_dim, actor_hidden_dim, use_layer_norm, "cpu", num_envs ) @@ -151,6 +218,13 @@ def _run_collector( obs_np = np.asarray(obs_np, dtype=np.float32) if priv_np is not None: priv_np = np.asarray(priv_np, dtype=np.float32) + obs_np, priv_np = maybe_adapt_legacy_go1_motrix_obs( + env_name=env_name, + env_cfg_override=env_cfg_override, + obs=obs_np, + privileged=priv_np, + expected_obs_dim=obs_dim, + ) max_episode_steps = getattr(getattr(env, "cfg", None), "max_episode_steps", None) if max_episode_steps is not None and int(max_episode_steps) > 0: step_offsets = np.random.randint( @@ -227,6 +301,13 @@ def _run_collector( next_obs_np = np.asarray(next_obs_np, dtype=np.float32) if next_priv_np is not None: next_priv_np = np.asarray(next_priv_np, dtype=np.float32) + next_obs_np, next_priv_np = maybe_adapt_legacy_go1_motrix_obs( + env_name=env_name, + env_cfg_override=env_cfg_override, + obs=next_obs_np, + privileged=next_priv_np, + expected_obs_dim=obs_dim, + ) rewards_np = np.asarray(state.reward, dtype=np.float32).ravel() terminated_np = ( @@ -255,9 +336,17 @@ def _run_collector( if np.any(has_final_np): final_obs_np, final_priv_np = split_obs_dict(state.info["final_observation"]) final_obs_np = np.asarray(final_obs_np, dtype=np.float32) + if final_priv_np is not None: + final_priv_np = np.asarray(final_priv_np, dtype=np.float32) + final_obs_np, final_priv_np = maybe_adapt_legacy_go1_motrix_obs( + env_name=env_name, + env_cfg_override=env_cfg_override, + obs=final_obs_np, + privileged=final_priv_np, + expected_obs_dim=obs_dim, + ) next_obs_np[has_final_np] = final_obs_np[has_final_np] if final_priv_np is not None and next_priv_np is not None: - final_priv_np = np.asarray(final_priv_np, dtype=np.float32) next_priv_np[has_final_np] = final_priv_np[has_final_np] # Write to replay buffer diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 2bc9843d0..c2ccda4ea 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -29,6 +29,65 @@ class Commands: ] +@dataclass +class Domain_Rand: + randomize_base_mass: bool = False + added_mass_range: list[float] = field(default_factory=lambda: [-1.5, 1.5]) + + random_com: bool = False + com_offset_x: list[float] = field(default_factory=lambda: [-0.05, 0.05]) + + push_robots: bool = False + push_interval: int = 750 + max_force: list[float] = field(default_factory=lambda: [1.0, 1.0, 0.5]) + + +def sample_velocity_commands( + rng, num_samples: int, low: np.ndarray, high: np.ndarray +) -> np.ndarray: + return rng.uniform(low=low, high=high, size=(num_samples, 3)).astype(get_global_dtype()) + + +def sample_gait_phase_pairs(rng, num_samples: int, mode: str) -> np.ndarray: + if mode == "independent": + return np.column_stack( + [ + rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)), + rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)), + ] + ).astype(get_global_dtype()) + + phase = rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)) + return np.column_stack([phase, phase + np.pi]).astype(get_global_dtype()) + + +def sample_reset_base_qvel(rng, num_samples: int, limit: float) -> np.ndarray: + return rng.uniform(-limit, limit, size=(num_samples, 6)).astype(get_global_dtype()) + + +def build_upper_body_pose_weights(pose_weights: list[float]) -> np.ndarray: + weights = np.asarray(pose_weights, dtype=get_global_dtype()).copy() + weights[:12] = 0.0 + return weights + + +def compute_weighted_pose_penalty(diff: np.ndarray, weights: np.ndarray) -> np.ndarray: + return np.sum(weights * np.square(diff), axis=1) + + +def compute_forward_progress_reward(commands: np.ndarray, linvel: np.ndarray) -> np.ndarray: + commanded_speed = np.maximum(commands[:, 0], 1e-6) + forward_speed = np.maximum(linvel[:, 0], 0.0) + return np.minimum(forward_speed / commanded_speed, 1.0).astype(get_global_dtype()) + + +def compute_under_speed_penalty(commands: np.ndarray, linvel: np.ndarray) -> np.ndarray: + commanded_speed = np.maximum(commands[:, 0], 1e-6) + forward_speed = np.maximum(linvel[:, 0], 0.0) + under_speed = np.maximum(commands[:, 0] - forward_speed, 0.0) + return (under_speed / commanded_speed).astype(get_global_dtype()) + + @dataclass class RewardConfigPPO: scales: dict[str, float] @@ -39,7 +98,39 @@ class RewardConfigPPO: base_height_target: float min_base_height: float max_tilt_deg: float - pose_weights: list[float] + pose_weights: list[float] = field( + default_factory=lambda: [ + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + ] + ) # PPO Environment @@ -51,6 +142,38 @@ class G1JoystickPPOCfg(G1BaseCfg): init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfigPPO | None = None + domain_rand: Domain_Rand = field(default_factory=Domain_Rand) + gait_phase_init_mode: str = "offset_phase" + reset_base_qvel_limit: float = 0.5 + backend_overrides: dict | None = None + + def apply_backend_overrides(self, backend_type: str) -> None: + if not self.backend_overrides or not self.backend_overrides.get("enabled", False): + return + if backend_type != "motrix": + return + if self.reward_config is None: + raise ValueError("reward_config must be provided before applying backend overrides") + + action_scale = self.backend_overrides.get("control_action_scale") + if action_scale is not None: + self.control_config.action_scale = float(action_scale) + + command_vel_limit = self.backend_overrides.get("command_vel_limit") + if command_vel_limit is not None: + self.commands.vel_limit = [list(command_vel_limit[0]), list(command_vel_limit[1])] + + reward_scale_overrides = self.backend_overrides.get("reward_scale_overrides") + if reward_scale_overrides: + self.reward_config.scales.update(reward_scale_overrides) + + gait_phase_init_mode = self.backend_overrides.get("gait_phase_init_mode") + if gait_phase_init_mode is not None: + self.gait_phase_init_mode = str(gait_phase_init_mode) + + reset_base_qvel_limit = self.backend_overrides.get("reset_base_qvel_limit") + if reset_base_qvel_limit is not None: + self.reset_base_qvel_limit = float(reset_base_qvel_limit) @registry.env("G1JoystickFlatTerrain", sim_backend="mujoco") @@ -61,6 +184,7 @@ class G1JoystickPPO(G1BaseEnv): def __init__(self, cfg: G1JoystickPPOCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: raise ValueError("reward_config must be provided via Hydra configuration") + cfg.apply_backend_overrides(backend_type) backend = create_backend( backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name ) @@ -74,6 +198,7 @@ def __init__(self, cfg: G1JoystickPPOCfg, num_envs=1, backend_type="mujoco"): self._pose_weights = np.array(self._reward_cfg.pose_weights, dtype=get_global_dtype()) if self._pose_weights.shape[0] != self._num_action: raise ValueError("pose_weights length mismatch") + self._upper_body_pose_weights = build_upper_body_pose_weights(self._reward_cfg.pose_weights) self._init_reward_functions() @@ -86,6 +211,9 @@ def _init_reward_functions(self): self._reward_fns = { "tracking_lin_vel": self._reward_tracking_lin_vel, "tracking_ang_vel": self._reward_tracking_ang_vel, + "forward_progress": self._reward_forward_progress, + "under_speed": self._reward_under_speed, + "upper_body_pose": self._reward_upper_body_pose, "feet_phase": self._reward_feet_phase, "lin_vel_z": self._reward_lin_vel_z, "orientation": self._reward_orientation, @@ -95,6 +223,14 @@ def _init_reward_functions(self): "pose": self._reward_pose, } + def _use_legacy_motrix_baseline(self) -> bool: + backend_type = getattr(getattr(self, "_backend", None), "backend_type", None) + return bool( + backend_type == "motrix" + and self._cfg.backend_overrides + and self._cfg.backend_overrides.get("enabled", False) + ) + def update_state(self, state: NpEnvState) -> NpEnvState: linvel = self.get_local_linvel() gyro = self.get_gyro() @@ -127,6 +263,11 @@ def _compute_obs( ) return {"obs": actor, "privileged": linvel} + def get_policy_obs(self, obs: dict[str, np.ndarray]) -> np.ndarray: + if self._use_legacy_motrix_baseline() and "privileged" in obs: + return np.concatenate([obs["privileged"], obs["obs"]], axis=1) + return np.concatenate(list(obs.values()), axis=1) + def get_obs_structure(self) -> dict: """Return observation structure for symmetry augmentation.""" return { @@ -171,6 +312,12 @@ def _reward_tracking_ang_vel(self, info, linvel, gyro, gravity, dof_pos, dof_vel ang_vel_error = np.square(commands[:, 2] - gyro[:, 2]) return np.exp(-ang_vel_error / self._reward_cfg.tracking_sigma) + def _reward_forward_progress(self, info, linvel, gyro, gravity, dof_pos, dof_vel): + return compute_forward_progress_reward(info["commands"], linvel) + + def _reward_under_speed(self, info, linvel, gyro, gravity, dof_pos, dof_vel): + return compute_under_speed_penalty(info["commands"], linvel) + def _reward_feet_phase(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """步态相位奖励：鼓励正确的摆动腿高度""" left_foot = self._backend.get_sensor_data("left_foot_pos") @@ -220,7 +367,11 @@ def _reward_action_rate(self, info, linvel, gyro, gravity, dof_pos, dof_vel): def _reward_pose(self, info, linvel, gyro, gravity, dof_pos, dof_vel): diff = dof_pos - self.default_angles - return np.sum(self._pose_weights * np.square(diff), axis=1) + return compute_weighted_pose_penalty(diff, self._pose_weights) + + def _reward_upper_body_pose(self, info, linvel, gyro, gravity, dof_pos, dof_vel): + diff = dof_pos - self.default_angles + return compute_weighted_pose_penalty(diff, self._upper_body_pose_weights) def reset(self, env_indices: np.ndarray): dtype = get_global_dtype() @@ -233,26 +384,30 @@ def reset(self, env_indices: np.ndarray): yaw = np.random.uniform(-np.pi, np.pi, (num_reset,)) quat_yaw = np_yaw_to_quat(yaw) qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], quat_yaw) - qvel[:, 0:6] = np.random.uniform(-0.5, 0.5, (num_reset, 6)) + qvel[:, 0:6] = sample_reset_base_qvel( + rng=np.random, + num_samples=num_reset, + limit=self._cfg.reset_base_qvel_limit, + ) self._backend.set_state(env_indices, qpos, qvel) - commands = np.random.uniform( - low=self._cfg.commands.vel_limit[0], - high=self._cfg.commands.vel_limit[1], - size=(num_reset, 3), + commands = sample_velocity_commands( + rng=np.random, + num_samples=num_reset, + low=np.asarray(self._cfg.commands.vel_limit[0], dtype=get_global_dtype()), + high=np.asarray(self._cfg.commands.vel_limit[1], dtype=get_global_dtype()), ) info = { "commands": commands, "current_actions": np.zeros((num_reset, self._num_action), dtype=dtype), "last_actions": np.zeros((num_reset, self._num_action), dtype=dtype), - "gait_phase": np.column_stack( - [ - np.random.uniform(0, 2 * np.pi, num_reset), - np.random.uniform(0, 2 * np.pi, num_reset) + np.pi, - ] - ).astype(dtype), + "gait_phase": sample_gait_phase_pairs( + rng=np.random, + num_samples=num_reset, + mode=self._cfg.gait_phase_init_mode, + ), } linvel = self.get_local_linvel()[env_indices] diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index 984444c3a..404fd9a54 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/src/unilab/envs/locomotion/g1/joystick_sac.py @@ -12,7 +12,7 @@ from unilab.base.curriculum import EpisodeLengthTracker, PenaltyCurriculum from unilab.base.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv -from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, InitState +from unilab.envs.locomotion.g1.joystick import Domain_Rand, G1JoystickPPO, InitState @dataclass @@ -56,6 +56,9 @@ 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: Domain_Rand = field(default_factory=Domain_Rand) + gait_phase_init_mode: str = "offset_phase" + reset_base_qvel_limit: float = 0.5 @registry.env("G1WalkTaskMjSAC", sim_backend="mujoco") diff --git a/src/unilab/envs/locomotion/g1/xml/g1.xml b/src/unilab/envs/locomotion/g1/xml/g1.xml index 80c298e98..fc14205ca 100644 --- a/src/unilab/envs/locomotion/g1/xml/g1.xml +++ b/src/unilab/envs/locomotion/g1/xml/g1.xml @@ -178,9 +178,7 @@ - - @@ -188,47 +186,39 @@ diaginertia="0.000465864 0.000432842 0.000406394"/> - - - - - - - @@ -243,48 +233,39 @@ diaginertia="0.000465864 0.000432842 0.000406394"/> - - - - - - - diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index fac400c21..137cdebdf 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -33,6 +33,12 @@ class RewardConfig: base_height_target: float +@dataclass +class LegacyMotrixProfile: + enabled: bool = False + command_vel_limit: list[list[float]] | None = None + + @dataclass class JoystickSensor: local_linvel = "local_linvel" @@ -63,9 +69,17 @@ class Go1JoystickCfg(Go1BaseCfg): init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfig | None = None + legacy_motrix_profile: LegacyMotrixProfile = field(default_factory=LegacyMotrixProfile) sensor: JoystickSensor = field(default_factory=JoystickSensor) # type: ignore[assignment] domain_rand: Domain_Rand = field(default_factory=Domain_Rand) + def apply_legacy_motrix_profile(self) -> None: + profile = self.legacy_motrix_profile + if not profile.enabled: + return + if profile.command_vel_limit is not None: + self.commands.vel_limit = [list(v) for v in profile.command_vel_limit] + @registry.env("Go1JoystickFlatTerrain", sim_backend="mujoco") @registry.env("Go1JoystickFlatTerrain", sim_backend="motrix") @@ -75,6 +89,7 @@ class Go1WalkTask(Go1BaseEnv): def __init__(self, cfg: Go1JoystickCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: raise ValueError("reward_config must be provided via Hydra configuration") + cfg.apply_legacy_motrix_profile() backend = create_backend( backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name ) @@ -89,9 +104,16 @@ def __init__(self, cfg: Go1JoystickCfg, num_envs=1, backend_type="mujoco"): @property def obs_groups_spec(self) -> dict[str, int]: + if self._use_legacy_motrix_profile(): + return {"obs": 48} # gyro(3) + gravity(3) + diff(12) + dof_vel(12) + action(12) + cmd(3) + phase(4) = 49 return {"obs": 49, "privileged": 3} + def _use_legacy_motrix_profile(self) -> bool: + profile = getattr(self._cfg, "legacy_motrix_profile", None) + backend_type = getattr(getattr(self, "_backend", None), "backend_type", None) + return bool(profile and profile.enabled and backend_type == "motrix") + def _init_reward_functions(self): self._reward_fns = { "tracking_lin_vel": self._reward_tracking_lin_vel, @@ -101,8 +123,9 @@ def _init_reward_functions(self): "base_height": self._reward_base_height, "action_rate": self._reward_action_rate, "similar_to_default": self._reward_similar_to_default, - "contact": self._reward_contact, } + if not self._use_legacy_motrix_profile(): + self._reward_fns["contact"] = self._reward_contact def update_state(self, state: NpEnvState) -> NpEnvState: self.phase = np.fmod(self.phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0) @@ -133,6 +156,13 @@ def _compute_obs( diff = dof_pos - self.default_angles command = info["commands"] last_actions = info.get("current_actions", np.zeros_like(diff)) + if self._use_legacy_motrix_profile(): + obs = np.concatenate( + [linvel, gyro, -gravity, diff, dof_vel, last_actions, command], + axis=1, + dtype=get_global_dtype(), + ) + return {"obs": obs} obs = np.concatenate( [gyro, -gravity, diff, dof_vel, last_actions, command, feet_phase], axis=1, diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 72fc3b4ce..94c74d9ed 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -52,6 +52,18 @@ def test_g1_joystick_sac_cfg_no_obs_config(): ) +def test_g1_joystick_sac_cfg_has_domain_rand_for_motrix(): + from unilab.envs.locomotion.g1.joystick_sac import G1JoystickSACCfg + + cfg = G1JoystickSACCfg() + assert hasattr(cfg, "domain_rand") + assert hasattr(cfg, "gait_phase_init_mode") + assert hasattr(cfg, "reset_base_qvel_limit") + assert cfg.domain_rand.randomize_base_mass is False + assert cfg.domain_rand.random_com is False + assert cfg.domain_rand.push_robots is False + + def test_g1_joystick_ppo_obs_groups_spec_dims(): """obs_groups_spec total dim must match what _compute_obs actually produces. diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 496b90646..0b49695b6 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -94,6 +94,18 @@ def test_offpolicy_hydra_algo_td3(): assert cfg.algo.algo == "td3" +def test_offpolicy_resolve_sac_use_symmetry_keeps_mujoco_setting(): + cfg = _offpolicy_cfg(["task=g1_sac", "training.sim_backend=mujoco"]) + + assert _offpolicy().resolve_sac_use_symmetry(cfg) is True + + +def test_offpolicy_resolve_sac_use_symmetry_disables_motrix(): + cfg = _offpolicy_cfg(["task=g1_sac", "training.sim_backend=motrix"]) + + assert _offpolicy().resolve_sac_use_symmetry(cfg) is False + + # --------------------------------------------------------------------------- # train_offpolicy.py — default_device() # --------------------------------------------------------------------------- @@ -199,6 +211,42 @@ def test_resolve_checkpoint_empty_run_dir(tmp_path): assert path is None +def test_offpolicy_extract_reset_obs_handles_two_tuple(): + obs = {"obs": "value"} + + result = _offpolicy().extract_reset_obs((obs, {"info": 1})) + + assert result is obs + + +def test_offpolicy_extract_reset_obs_handles_three_tuple(): + obs = {"obs": "value"} + + result = _offpolicy().extract_reset_obs(("ignored", obs, {"info": 1})) + + assert result is obs + + +def test_offpolicy_resolve_play_obs_dim_ignores_privileged(): + obs_dim = _offpolicy().resolve_play_obs_dim({"obs": 98, "privileged": 3}) + + assert obs_dim == 98 + + +def test_offpolicy_extract_play_obs_uses_obs_group_only(): + import numpy as np + + obs = { + "obs": np.ones((2, 98), dtype=np.float32), + "privileged": np.full((2, 3), 2.0, dtype=np.float32), + } + + play_obs = _offpolicy().extract_play_obs(obs) + + assert play_obs.shape == (2, 98) + assert np.allclose(play_obs, 1.0) + + # --------------------------------------------------------------------------- # train_mlx_ppo.py — get_latest_run() / get_latest_checkpoint() # --------------------------------------------------------------------------- From e07f37c67f3396579cc740158bc06f52b1e6c897 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=A2=E6=9D=B0?= Date: Wed, 18 Mar 2026 23:28:15 +0800 Subject: [PATCH 2/5] Ignore project-local worktrees --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index dcd6a3627..e0016ce15 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,7 @@ dist/ # Project specific logs/ +.worktrees/ *.log .DS_Store From ba2ab7cc1509ae40713438beaa7c1e35c8baaf74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=A2=E6=9D=B0?= Date: Thu, 19 Mar 2026 11:33:12 +0800 Subject: [PATCH 3/5] feat: sync motrix task overrides and pass make check --- conf/offpolicy/task/g1_sac.yaml | 4 + conf/offpolicy/task/go1_joystick.yaml | 24 +++ conf/ppo/task/g1_joystick.yaml | 47 +++++ conf/ppo/task/go1_joystick.yaml | 27 +++ pyproject.toml | 1 + scripts/train_appo.py | 3 +- scripts/train_offpolicy.py | 130 +++++++----- scripts/train_rsl_rl.py | 189 +++++++----------- src/unilab/algos/mlx/common/activations.py | 4 +- src/unilab/algos/torch/appo/runner.py | 3 +- src/unilab/algos/torch/fast_sac/learner.py | 1 + src/unilab/algos/torch/fast_sac/runner.py | 1 + src/unilab/algos/torch/fast_td3/learner.py | 1 + src/unilab/algos/torch/fast_td3/runner.py | 2 +- .../algos/torch/offpolicy/multi_gpu_runner.py | 4 +- src/unilab/algos/torch/offpolicy/worker.py | 2 + src/unilab/base/registry.py | 2 +- src/unilab/config/structured_configs.py | 6 +- src/unilab/envs/locomotion/g1/joystick.py | 34 ++-- .../g1/motion_tracking/math_utils.py | 14 +- .../g1/motion_tracking/motion_loader.py | 6 +- .../locomotion/g1/motion_tracking/tracking.py | 31 ++- .../inhand_rot_allegro/rotation.py | 4 +- .../inhand_rot_allegro/rotation_sac.py | 1 + src/unilab/ipc/replay_buffer.py | 3 + src/unilab/utils/device_utils.py | 4 +- src/unilab/utils/reward_utils.py | 4 +- src/unilab/utils/xml_utils.py | 6 +- tests/algos/test_mlx_ppo.py | 7 +- tests/algos/test_rsl_rl_runner.py | 8 +- tests/config/test_config_system.py | 7 +- tests/conftest.py | 65 +++++- tests/envs/test_env_configs.py | 4 +- .../test_reward_injection_integration.py | 2 +- tests/scripts/test_train_scripts.py | 138 ++++++++++++- 35 files changed, 555 insertions(+), 234 deletions(-) diff --git a/conf/offpolicy/task/g1_sac.yaml b/conf/offpolicy/task/g1_sac.yaml index 41374537c..d55693745 100644 --- a/conf/offpolicy/task/g1_sac.yaml +++ b/conf/offpolicy/task/g1_sac.yaml @@ -13,3 +13,7 @@ algo: algo_params: alpha_init: 0.01 target_entropy_ratio: 0.0 +motrix_legacy: + enabled: true + algo_overrides: + use_symmetry: false diff --git a/conf/offpolicy/task/go1_joystick.yaml b/conf/offpolicy/task/go1_joystick.yaml index 59cb32025..ca9422c75 100644 --- a/conf/offpolicy/task/go1_joystick.yaml +++ b/conf/offpolicy/task/go1_joystick.yaml @@ -7,3 +7,27 @@ training: algo: num_envs: 2048 max_iterations: 3000 +motrix_legacy: + enabled: true + applies_to: + algo: sac + algo_overrides: + num_envs: 4096 + max_iterations: 2000 + env_cfg_override: + legacy_motrix_profile: + enabled: true + command_vel_limit: + - [0.5, 0.0, 0.0] + - [0.5, 0.0, 0.0] + reward_config: + scales: + tracking_lin_vel: 1.0 + tracking_ang_vel: 0.2 + lin_vel_z: -5.0 + ang_vel_xy: -0.1 + base_height: -100.0 + action_rate: -0.005 + similar_to_default: -0.1 + tracking_sigma: 0.25 + base_height_target: 0.3 diff --git a/conf/ppo/task/g1_joystick.yaml b/conf/ppo/task/g1_joystick.yaml index bfcae846d..db3f8ede4 100644 --- a/conf/ppo/task/g1_joystick.yaml +++ b/conf/ppo/task/g1_joystick.yaml @@ -12,3 +12,50 @@ algo: - policy actor: - actor +motrix_legacy: + enabled: true + algo_overrides: + max_iterations: 151 + empirical_normalization: true + obs_groups: + actor: + - policy + policy: + init_noise_std: 0.5 + algorithm: + learning_rate: 3.0e-4 + entropy_coef: 5.0e-3 + env_cfg_override: + reward_config: + scales: + tracking_lin_vel: 3.0 + tracking_ang_vel: 0.25 + forward_progress: 1.5 + under_speed: -1.0 + upper_body_pose: -0.05 + feet_phase: 0.5 + lin_vel_z: -1.0 + ang_vel_xy: -0.2 + base_height: -120.0 + orientation: -2.5 + action_rate: -0.005 + pose: -0.05 + tracking_sigma: 0.25 + gait_frequency: 1.5 + feet_phase_swing_height: 0.09 + feet_phase_tracking_sigma: 0.008 + base_height_target: 0.754 + min_base_height: 0.5 + max_tilt_deg: 35.0 + backend_overrides: + enabled: true + control_action_scale: 0.5 + command_vel_limit: + - [0.45, 0.0, 0.0] + - [0.55, 0.0, 0.0] + gait_phase_init_mode: independent + reset_base_qvel_limit: 0.05 + reward_scale_overrides: + under_speed: -1.0 + upper_body_pose: -0.05 + tracking_ang_vel: 0.25 diff --git a/conf/ppo/task/go1_joystick.yaml b/conf/ppo/task/go1_joystick.yaml index 87c0a4557..085158c4e 100644 --- a/conf/ppo/task/go1_joystick.yaml +++ b/conf/ppo/task/go1_joystick.yaml @@ -11,3 +11,30 @@ algo: - policy actor: - actor +motrix_legacy: + enabled: true + algo_overrides: + max_iterations: 151 + empirical_normalization: true + policy: + init_noise_std: 0.5 + algorithm: + learning_rate: 3.0e-4 + entropy_coef: 1.0e-3 + env_cfg_override: + legacy_motrix_profile: + enabled: true + command_vel_limit: + - [0.5, 0.0, 0.0] + - [0.5, 0.0, 0.0] + reward_config: + scales: + tracking_lin_vel: 1.0 + tracking_ang_vel: 0.2 + lin_vel_z: -5.0 + ang_vel_xy: -0.1 + base_height: -100.0 + action_rate: -0.005 + similar_to_default: -0.1 + tracking_sigma: 0.25 + base_height_target: 0.3 diff --git a/pyproject.toml b/pyproject.toml index 78389c6f4..4c31584c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ mujoco-uni = { index = "test-pypi" } line-length = 100 target-version = "py310" src = ["src"] +exclude = [".history"] [tool.ruff.lint] select = ["E", "F", "I", "N", "W"] diff --git a/scripts/train_appo.py b/scripts/train_appo.py index 94428f806..270a15621 100644 --- a/scripts/train_appo.py +++ b/scripts/train_appo.py @@ -54,8 +54,8 @@ def play_appo(cfg: DictConfig, rl_cfg: dict): from unilab.base import registry from unilab.utils import render_many - from unilab.utils.rsl_rl_compat import convert_config_v3_to_v4, is_rsl_rl_v4, is_rsl_rl_v5 from unilab.utils.reward_utils import extract_reward_config + from unilab.utils.rsl_rl_compat import convert_config_v3_to_v4, is_rsl_rl_v4, is_rsl_rl_v5 env_cfg_override = extract_reward_config(cfg) @@ -75,6 +75,7 @@ def play_appo(cfg: DictConfig, rl_cfg: dict): env_cfg_override=env_cfg_override, ) from unilab.utils.obs_utils import get_obs_dims + obs_dim, privileged_dim = get_obs_dims(env.obs_groups_spec) action_dim = env.action_space.shape[0] diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index d56f631b9..db6cc6df3 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -79,14 +79,11 @@ def resolve_checkpoint_path( def extract_reset_obs(reset_result): - """Extract obs_dict from env.reset(...) across old/new return conventions.""" + """Extract obs_dict from env.reset(...) using the current (obs_dict, info_dict) contract.""" if isinstance(reset_result, tuple): if len(reset_result) == 2: obs_out, _ = reset_result return obs_out - if len(reset_result) == 3: - _, obs_out, _ = reset_result - return obs_out raise ValueError(f"Unexpected env.reset return format: {type(reset_result)!r}") @@ -104,56 +101,85 @@ def extract_play_obs(obs_dict): return obs_out -_GO1_LEGACY_COMMAND_VEL_LIMIT = [[0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] -_GO1_LEGACY_REWARD_CONFIG = { - "scales": { - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.2, - "lin_vel_z": -5.0, - "ang_vel_xy": -0.1, - "base_height": -100.0, - "action_rate": -0.005, - "similar_to_default": -0.1, - }, - "tracking_sigma": 0.25, - "base_height_target": 0.3, -} - - -def get_go1_motrix_legacy_offpolicy_profile( - algo_name: str, task_name: str, sim_backend: str -) -> dict | None: - if algo_name != "sac" or task_name != "Go1JoystickFlatTerrain" or sim_backend != "motrix": - return None - return { - "algo_overrides": { - "num_envs": 4096, - "max_iterations": 2000, - }, - "env_cfg_override": { - "legacy_motrix_profile": { - "enabled": True, - "command_vel_limit": _GO1_LEGACY_COMMAND_VEL_LIMIT, - }, - "reward_config": _GO1_LEGACY_REWARD_CONFIG, - }, - } - - -def build_go1_motrix_legacy_offpolicy_env_cfg_override( - algo_name: str, cfg: DictConfig -) -> dict | None: +def _explicit_cli_override_keys() -> set[str]: + return {arg.split("=", 1)[0] for arg in sys.argv[1:] if "=" in arg} + + +def _apply_task_algo_overrides( + target, overrides, *, base_path: str, explicit_keys: set[str] +) -> None: + if OmegaConf.is_config(overrides): + override_dict = OmegaConf.to_container(overrides, resolve=True) + elif isinstance(overrides, dict): + override_dict = overrides + else: + return + + for key, value in override_dict.items(): + path = f"{base_path}.{key}" + if isinstance(value, dict): + if path in explicit_keys: + continue + if key not in target or target[key] is None: + target[key] = {} + _apply_task_algo_overrides( + target[key], value, base_path=path, explicit_keys=explicit_keys + ) + continue + if path not in explicit_keys: + target[key] = value + + +def _apply_task_env_profile( + env_cfg_override: dict, env_profile, *, explicit_keys: set[str] +) -> None: + if OmegaConf.is_config(env_profile): + env_profile_dict = OmegaConf.to_container(env_profile, resolve=True) + elif isinstance(env_profile, dict): + env_profile_dict = env_profile + else: + return + + reward_explicit = "reward" in explicit_keys or any( + k.startswith("reward.") for k in explicit_keys + ) + for key, value in env_profile_dict.items(): + if key == "reward_config" and reward_explicit: + continue + env_cfg_override[key] = value + + +def build_task_motrix_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict | None: from unilab.utils.reward_utils import extract_reward_config env_cfg_override = extract_reward_config(cfg) + explicit_keys = _explicit_cli_override_keys() + motrix_legacy = getattr(cfg, "motrix_legacy", None) + if motrix_legacy is None or not motrix_legacy.enabled or cfg.training.sim_backend != "motrix": + return env_cfg_override + + applies_to = getattr(motrix_legacy, "applies_to", None) + if applies_to is not None: + target_algo = getattr(applies_to, "algo", None) + if target_algo is not None and target_algo != algo_name: + return env_cfg_override + + algo_overrides = getattr(motrix_legacy, "algo_overrides", None) + if algo_overrides is not None: + _apply_task_algo_overrides( + cfg.algo, + algo_overrides, + base_path="algo", + explicit_keys=explicit_keys, + ) - profile = get_go1_motrix_legacy_offpolicy_profile( - algo_name, cfg.training.task_name, cfg.training.sim_backend - ) - if profile is not None: - cfg.algo.num_envs = profile["algo_overrides"]["num_envs"] - cfg.algo.max_iterations = profile["algo_overrides"]["max_iterations"] - env_cfg_override.update(profile["env_cfg_override"]) + env_profile = getattr(motrix_legacy, "env_cfg_override", None) + if env_profile is not None: + _apply_task_env_profile( + env_cfg_override, + env_profile, + explicit_keys=explicit_keys, + ) return env_cfg_override @@ -165,7 +191,7 @@ def resolve_sac_use_symmetry(cfg: DictConfig) -> bool: def build_runner(algo_name: str, cfg: DictConfig): """Build algorithm runner from unified Hydra config.""" - env_cfg_override = build_go1_motrix_legacy_offpolicy_env_cfg_override(algo_name, cfg) + env_cfg_override = build_task_motrix_offpolicy_env_cfg_override(algo_name, cfg) if algo_name == "sac": from unilab.algos.torch.fast_sac.learner import FastSACLearner @@ -315,7 +341,7 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> None: from unilab.utils import render_many from unilab.utils.algo_utils import build_actor - env_cfg_override = build_go1_motrix_legacy_offpolicy_env_cfg_override(algo_name, cfg) + env_cfg_override = build_task_motrix_offpolicy_env_cfg_override(algo_name, cfg) device = default_device(torch, cfg.training.device) print(f"Using device for play: {device}") diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index 53ee1bce5..219a6ab4f 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -47,126 +47,79 @@ def ensure_registries(): pass -_GO1_LEGACY_COMMAND_VEL_LIMIT = [[0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] -_GO1_LEGACY_REWARD_CONFIG = { - "scales": { - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.2, - "lin_vel_z": -5.0, - "ang_vel_xy": -0.1, - "base_height": -100.0, - "action_rate": -0.005, - "similar_to_default": -0.1, - }, - "tracking_sigma": 0.25, - "base_height_target": 0.3, -} - -_G1_MOTRIX_LEGACY_PPO_REWARD_CONFIG = { - "scales": { - "tracking_lin_vel": 3.0, - "tracking_ang_vel": 0.25, - "forward_progress": 1.5, - "under_speed": -1.0, - "upper_body_pose": -0.05, - "feet_phase": 0.5, - "lin_vel_z": -1.0, - "ang_vel_xy": -0.2, - "base_height": -120.0, - "orientation": -2.5, - "action_rate": -0.005, - "pose": -0.05, - }, - "tracking_sigma": 0.25, - "gait_frequency": 1.5, - "feet_phase_swing_height": 0.09, - "feet_phase_tracking_sigma": 0.008, - "base_height_target": 0.754, - "min_base_height": 0.50, - "max_tilt_deg": 35.0, -} - - -def get_go1_motrix_legacy_ppo_profile(task_name: str, sim_backend: str) -> dict | None: - if task_name != "Go1JoystickFlatTerrain" or sim_backend != "motrix": - return None - return { - "algo_overrides": { - "max_iterations": 151, - "empirical_normalization": True, - "policy": {"init_noise_std": 0.5}, - "algorithm": {"learning_rate": 3.0e-4, "entropy_coef": 1.0e-3}, - }, - "env_cfg_override": { - "legacy_motrix_profile": { - "enabled": True, - "command_vel_limit": _GO1_LEGACY_COMMAND_VEL_LIMIT, - }, - "reward_config": _GO1_LEGACY_REWARD_CONFIG, - }, - } - - -def get_g1_motrix_legacy_ppo_profile(task_name: str, sim_backend: str) -> dict | None: - if task_name != "G1JoystickFlatTerrain" or sim_backend != "motrix": - return None - return { - "algo_overrides": { - "max_iterations": 151, - "empirical_normalization": True, - "policy": {"init_noise_std": 0.5}, - "algorithm": {"learning_rate": 3.0e-4, "entropy_coef": 5.0e-3}, - }, - "env_cfg_override": { - "reward_config": _G1_MOTRIX_LEGACY_PPO_REWARD_CONFIG, - "backend_overrides": { - "enabled": True, - "control_action_scale": 0.5, - "command_vel_limit": [[0.45, 0.0, 0.0], [0.55, 0.0, 0.0]], - "gait_phase_init_mode": "independent", - "reset_base_qvel_limit": 0.05, - "reward_scale_overrides": { - "under_speed": -1.0, - "upper_body_pose": -0.05, - "tracking_ang_vel": 0.25, - }, - }, - }, - } - - -def build_go1_motrix_legacy_ppo_env_cfg_override(cfg: DictConfig) -> dict: - env_cfg_override = {} - if hasattr(cfg, "reward") and cfg.reward: - env_cfg_override["reward_config"] = OmegaConf.to_container(cfg.reward, resolve=True) +def _explicit_cli_override_keys() -> set[str]: + return {arg.split("=", 1)[0] for arg in sys.argv[1:] if "=" in arg} - profile = get_go1_motrix_legacy_ppo_profile(cfg.training.task_name, cfg.training.sim_backend) - if profile is not None: - cfg.algo.max_iterations = profile["algo_overrides"]["max_iterations"] - cfg.algo.empirical_normalization = profile["algo_overrides"]["empirical_normalization"] - cfg.algo.policy.init_noise_std = profile["algo_overrides"]["policy"]["init_noise_std"] - cfg.algo.algorithm.learning_rate = profile["algo_overrides"]["algorithm"][ - "learning_rate" - ] - cfg.algo.algorithm.entropy_coef = profile["algo_overrides"]["algorithm"]["entropy_coef"] - env_cfg_override.update(profile["env_cfg_override"]) - return env_cfg_override +def _apply_task_algo_overrides( + target, overrides, *, base_path: str, explicit_keys: set[str] +) -> None: + if OmegaConf.is_config(overrides): + override_dict = OmegaConf.to_container(overrides, resolve=True) + elif isinstance(overrides, dict): + override_dict = overrides + else: + return + + for key, value in override_dict.items(): + path = f"{base_path}.{key}" + if isinstance(value, dict): + if path in explicit_keys: + continue + if key not in target or target[key] is None: + target[key] = {} + _apply_task_algo_overrides( + target[key], value, base_path=path, explicit_keys=explicit_keys + ) + continue + if path not in explicit_keys: + target[key] = value + + +def _apply_task_env_profile( + env_cfg_override: dict, env_profile, *, explicit_keys: set[str] +) -> None: + if OmegaConf.is_config(env_profile): + env_profile_dict = OmegaConf.to_container(env_profile, resolve=True) + elif isinstance(env_profile, dict): + env_profile_dict = env_profile + else: + return + + reward_explicit = "reward" in explicit_keys or any( + k.startswith("reward.") for k in explicit_keys + ) + for key, value in env_profile_dict.items(): + if key == "reward_config" and reward_explicit: + continue + env_cfg_override[key] = value -def build_g1_motrix_legacy_ppo_env_cfg_override(cfg: DictConfig) -> dict: +def build_task_motrix_ppo_env_cfg_override(cfg: DictConfig) -> dict: env_cfg_override = {} - profile = get_g1_motrix_legacy_ppo_profile(cfg.training.task_name, cfg.training.sim_backend) - if profile is not None: - cfg.algo.max_iterations = profile["algo_overrides"]["max_iterations"] - cfg.algo.empirical_normalization = profile["algo_overrides"]["empirical_normalization"] - cfg.algo.policy.init_noise_std = profile["algo_overrides"]["policy"]["init_noise_std"] - cfg.algo.algorithm.learning_rate = profile["algo_overrides"]["algorithm"][ - "learning_rate" - ] - cfg.algo.algorithm.entropy_coef = profile["algo_overrides"]["algorithm"]["entropy_coef"] - cfg.algo.obs_groups["actor"] = ["policy"] - env_cfg_override.update(profile["env_cfg_override"]) + motrix_legacy = getattr(cfg, "motrix_legacy", None) + explicit_keys = _explicit_cli_override_keys() + if hasattr(cfg, "reward") and cfg.reward: + env_cfg_override["reward_config"] = OmegaConf.to_container(cfg.reward, resolve=True) + if motrix_legacy is None or not motrix_legacy.enabled or cfg.training.sim_backend != "motrix": + return env_cfg_override + + algo_overrides = getattr(motrix_legacy, "algo_overrides", None) + if algo_overrides is not None: + _apply_task_algo_overrides( + cfg.algo, + algo_overrides, + base_path="algo", + explicit_keys=explicit_keys, + ) + + env_profile = getattr(motrix_legacy, "env_cfg_override", None) + if env_profile is not None: + _apply_task_env_profile( + env_cfg_override, + env_profile, + explicit_keys=explicit_keys, + ) return env_cfg_override @@ -259,8 +212,7 @@ def play_rsl_rl(cfg: DictConfig, device: str): from unilab.base import registry - env_cfg_override = build_go1_motrix_legacy_ppo_env_cfg_override(cfg) - env_cfg_override.update(build_g1_motrix_legacy_ppo_env_cfg_override(cfg)) + env_cfg_override = build_task_motrix_ppo_env_cfg_override(cfg) env = registry.make( cfg.training.task_name, @@ -391,8 +343,7 @@ def main(cfg: DictConfig) -> None: from unilab.base import registry - env_cfg_override = build_go1_motrix_legacy_ppo_env_cfg_override(cfg) - env_cfg_override.update(build_g1_motrix_legacy_ppo_env_cfg_override(cfg)) + env_cfg_override = build_task_motrix_ppo_env_cfg_override(cfg) if torch.cuda.is_available(): device = "cuda" diff --git a/src/unilab/algos/mlx/common/activations.py b/src/unilab/algos/mlx/common/activations.py index 2b8320a41..2c8ca019a 100644 --- a/src/unilab/algos/mlx/common/activations.py +++ b/src/unilab/algos/mlx/common/activations.py @@ -18,9 +18,9 @@ def get_activation(name: str | None) -> Callable[[mx.array], mx.array]: if name == "elu": return lambda x: mx.where(x > 0.0, x, mx.exp(x) - 1.0) if name == "tanh": - return mx.tanh + return lambda x: mx.tanh(x) if name == "sigmoid": - return mx.sigmoid + return lambda x: mx.sigmoid(x) if name == "swish": return lambda x: x * mx.sigmoid(x) if name == "identity": diff --git a/src/unilab/algos/torch/appo/runner.py b/src/unilab/algos/torch/appo/runner.py index a67d1ac36..19183a44c 100644 --- a/src/unilab/algos/torch/appo/runner.py +++ b/src/unilab/algos/torch/appo/runner.py @@ -12,6 +12,7 @@ import time from collections import deque from copy import deepcopy +from typing import Any import torch from rsl_rl.utils import resolve_callable @@ -122,7 +123,7 @@ def _build_learner(self): actor = actor_cls(td_example, cfg["obs_groups"], "actor", self.action_dim, **actor_cfg) # Build critic (deterministic MLPModel, no distribution) - critic_cfg = deepcopy(cfg.get("critic", cfg.get("actor", ))) + critic_cfg: dict[str, Any] = deepcopy(cfg.get("critic") or cfg.get("actor") or {}) critic_cls = resolve_callable(critic_cfg.pop("class_name", "rsl_rl.models.MLPModel")) critic_cfg.pop("num_actions", None) critic_cfg.pop("distribution_cfg", None) # critic is deterministic diff --git a/src/unilab/algos/torch/fast_sac/learner.py b/src/unilab/algos/torch/fast_sac/learner.py index 8a98a7f65..ed83bad07 100644 --- a/src/unilab/algos/torch/fast_sac/learner.py +++ b/src/unilab/algos/torch/fast_sac/learner.py @@ -495,6 +495,7 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: # Critic input: obs + privileged if privileged is not None: + assert next_privileged is not None critic_obs = torch.cat([obs, privileged], dim=-1) critic_next_obs = torch.cat([next_obs, next_privileged], dim=-1) else: diff --git a/src/unilab/algos/torch/fast_sac/runner.py b/src/unilab/algos/torch/fast_sac/runner.py index 554278d9d..32c7cd98e 100644 --- a/src/unilab/algos/torch/fast_sac/runner.py +++ b/src/unilab/algos/torch/fast_sac/runner.py @@ -49,6 +49,7 @@ def __init__( env_name, num_envs=1, sim_backend=sim_backend, env_cfg_override=env_cfg_override ) from unilab.utils.obs_utils import get_obs_dims + obs_dim, privileged_dim = get_obs_dims(env.obs_groups_spec) act_space_shape = env.action_space.shape assert act_space_shape is not None diff --git a/src/unilab/algos/torch/fast_td3/learner.py b/src/unilab/algos/torch/fast_td3/learner.py index f8aa0cfdd..68d0bf223 100644 --- a/src/unilab/algos/torch/fast_td3/learner.py +++ b/src/unilab/algos/torch/fast_td3/learner.py @@ -247,6 +247,7 @@ def update_critic(self, data: Dict[str, torch.Tensor]) -> Dict[str, float]: # Critic input: obs + privileged if privileged is not None: + assert next_privileged is not None critic_obs = torch.cat([observations, privileged], dim=-1) critic_next_obs = torch.cat([next_observations, next_privileged], dim=-1) else: diff --git a/src/unilab/algos/torch/fast_td3/runner.py b/src/unilab/algos/torch/fast_td3/runner.py index 5f0dd497f..16e05b16c 100644 --- a/src/unilab/algos/torch/fast_td3/runner.py +++ b/src/unilab/algos/torch/fast_td3/runner.py @@ -42,7 +42,7 @@ def __init__( obs_normalization: bool = True, sim_backend: str = "mujoco", ): - obs_dim, action_dim = get_env_dims(env_name, sim_backend) + obs_dim, action_dim, _ = get_env_dims(env_name, sim_backend) learner = FastTD3Learner( obs_dim=obs_dim, action_dim=action_dim, diff --git a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py index 0d3e59654..6d3579a66 100644 --- a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -35,7 +35,7 @@ def _find_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) - return s.getsockname()[1] + return int(s.getsockname()[1]) def _broadcast_initial_params(learner: FastSACLearner, rank: int) -> None: @@ -400,7 +400,7 @@ def _learn_multi_gpu( } try: - tmp.spawn( + tmp.spawn( # pyright: ignore[reportPrivateImportUsage] _learner_worker, args=( self.num_gpus, diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index 322624c9a..c8f122a76 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -33,6 +33,8 @@ def resolve_collector_actor_dims( assert env.action_space.shape is not None action_dim = env.action_space.shape[0] + assert obs_dim is not None + assert action_dim is not None return obs_dim, action_dim diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index 45bde7ae0..36699a534 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -128,7 +128,7 @@ def make( # Create environment config env_cfg = meta.env_cfg_cls() if env_cfg_override is not None: - from typing import get_type_hints, get_args, get_origin + from typing import get_args, get_origin, get_type_hints # Get type hints for the config class type_hints = get_type_hints(env_cfg.__class__) diff --git a/src/unilab/config/structured_configs.py b/src/unilab/config/structured_configs.py index 5b7ac2124..f30e66aa5 100644 --- a/src/unilab/config/structured_configs.py +++ b/src/unilab/config/structured_configs.py @@ -8,12 +8,12 @@ import dataclasses from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Optional, cast class BaseConfig: - def to_dict(self) -> dict: - return dataclasses.asdict(self) # recursive, keeps test compatibility + def to_dict(self) -> dict[str, Any]: + return cast(dict[str, Any], dataclasses.asdict(cast(Any, self))) # ── Off-policy: SAC ────────────────────────────────────────────────────────── diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index c2ccda4ea..836809cf3 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -4,6 +4,7 @@ import math from dataclasses import dataclass, field +from typing import TYPE_CHECKING import numpy as np from etils import epath @@ -15,6 +16,9 @@ from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat +if TYPE_CHECKING: + from .joystick_sac import RewardConfigSAC + @dataclass class InitState: @@ -45,24 +49,29 @@ class Domain_Rand: def sample_velocity_commands( rng, num_samples: int, low: np.ndarray, high: np.ndarray ) -> np.ndarray: - return rng.uniform(low=low, high=high, size=(num_samples, 3)).astype(get_global_dtype()) + return np.asarray( + rng.uniform(low=low, high=high, size=(num_samples, 3)), dtype=get_global_dtype() + ) def sample_gait_phase_pairs(rng, num_samples: int, mode: str) -> np.ndarray: if mode == "independent": - return np.column_stack( - [ - rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)), - rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)), - ] - ).astype(get_global_dtype()) + return np.asarray( + np.column_stack( + [ + rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)), + rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)), + ] + ), + dtype=get_global_dtype(), + ) phase = rng.uniform(0.0, 2.0 * np.pi, size=(num_samples,)) - return np.column_stack([phase, phase + np.pi]).astype(get_global_dtype()) + return np.asarray(np.column_stack([phase, phase + np.pi]), dtype=get_global_dtype()) def sample_reset_base_qvel(rng, num_samples: int, limit: float) -> np.ndarray: - return rng.uniform(-limit, limit, size=(num_samples, 6)).astype(get_global_dtype()) + return np.asarray(rng.uniform(-limit, limit, size=(num_samples, 6)), dtype=get_global_dtype()) def build_upper_body_pose_weights(pose_weights: list[float]) -> np.ndarray: @@ -72,20 +81,20 @@ def build_upper_body_pose_weights(pose_weights: list[float]) -> np.ndarray: def compute_weighted_pose_penalty(diff: np.ndarray, weights: np.ndarray) -> np.ndarray: - return np.sum(weights * np.square(diff), axis=1) + return np.asarray(np.sum(weights * np.square(diff), axis=1), dtype=get_global_dtype()) def compute_forward_progress_reward(commands: np.ndarray, linvel: np.ndarray) -> np.ndarray: commanded_speed = np.maximum(commands[:, 0], 1e-6) forward_speed = np.maximum(linvel[:, 0], 0.0) - return np.minimum(forward_speed / commanded_speed, 1.0).astype(get_global_dtype()) + return np.asarray(np.minimum(forward_speed / commanded_speed, 1.0), dtype=get_global_dtype()) def compute_under_speed_penalty(commands: np.ndarray, linvel: np.ndarray) -> np.ndarray: commanded_speed = np.maximum(commands[:, 0], 1e-6) forward_speed = np.maximum(linvel[:, 0], 0.0) under_speed = np.maximum(commands[:, 0] - forward_speed, 0.0) - return (under_speed / commanded_speed).astype(get_global_dtype()) + return np.asarray(under_speed / commanded_speed, dtype=get_global_dtype()) @dataclass @@ -180,6 +189,7 @@ def apply_backend_overrides(self, backend_type: str) -> None: @registry.env("G1JoystickFlatTerrain", sim_backend="motrix") class G1JoystickPPO(G1BaseEnv): _cfg: G1JoystickPPOCfg + _reward_cfg: RewardConfigPPO | RewardConfigSAC def __init__(self, cfg: G1JoystickPPOCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: diff --git a/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py b/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py index 88e552939..a776a188b 100644 --- a/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py +++ b/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import cast + import numpy as np @@ -35,7 +37,7 @@ def quat_mul(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: axis=1, ) - return result[0] if squeeze else result + return cast(np.ndarray, np.asarray(result[0] if squeeze else result)) def quat_conjugate(q: np.ndarray) -> np.ndarray: @@ -112,7 +114,7 @@ def quat_apply(q: np.ndarray, v: np.ndarray) -> np.ndarray: axis=1, ) - return result[0] if squeeze else result + return cast(np.ndarray, np.asarray(result[0] if squeeze else result)) def quat_apply_inverse(q: np.ndarray, v: np.ndarray) -> np.ndarray: @@ -152,7 +154,7 @@ def quat_error_magnitude(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: xyz_norm = np.linalg.norm(q_rel[:, 1:], axis=1) error = 2 * np.arcsin(np.clip(xyz_norm, -1, 1)) - return error[0] if squeeze else error + return cast(np.ndarray, np.asarray(error[0] if squeeze else error)) def quat_from_euler_xyz(roll: np.ndarray, pitch: np.ndarray, yaw: np.ndarray) -> np.ndarray: @@ -184,7 +186,7 @@ def quat_from_euler_xyz(roll: np.ndarray, pitch: np.ndarray, yaw: np.ndarray) -> z = cr * cp * sy - sr * sp * cy result = np.stack([w, x, y, z], axis=1) - return result[0] if squeeze else result + return cast(np.ndarray, np.asarray(result[0] if squeeze else result)) def yaw_quat(q: np.ndarray) -> np.ndarray: @@ -218,7 +220,7 @@ def yaw_quat(q: np.ndarray) -> np.ndarray: axis=1, ) - return result[0] if squeeze else result + return cast(np.ndarray, np.asarray(result[0] if squeeze else result)) def matrix_from_quat(q: np.ndarray) -> np.ndarray: @@ -258,7 +260,7 @@ def matrix_from_quat(q: np.ndarray) -> np.ndarray: axis=1, ) - return result[0] if squeeze else result + return cast(np.ndarray, np.asarray(result[0] if squeeze else result)) def subtract_frame_transforms( diff --git a/src/unilab/envs/locomotion/g1/motion_tracking/motion_loader.py b/src/unilab/envs/locomotion/g1/motion_tracking/motion_loader.py index 1090885d4..a981a4eb3 100644 --- a/src/unilab/envs/locomotion/g1/motion_tracking/motion_loader.py +++ b/src/unilab/envs/locomotion/g1/motion_tracking/motion_loader.py @@ -208,10 +208,10 @@ def _sample_adaptive(self, env_ids: np.ndarray) -> np.ndarray: pmax = sampling_probs[pmax_idx] self.sampling_entropy = H_norm - self.sampling_top1_prob = pmax - self.sampling_top1_bin = pmax_idx / self.bin_count + self.sampling_top1_prob = float(pmax) + self.sampling_top1_bin = float(pmax_idx) / self.bin_count - return frames + return np.asarray(frames, dtype=np.int32) def update_failure_stats( self, terminated: np.ndarray, current_frames: np.ndarray | None = None diff --git a/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py b/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py index cf87c1a00..391063fd7 100644 --- a/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py +++ b/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Literal import mujoco import numpy as np @@ -100,7 +101,7 @@ class G1MotionTrackingCfg(G1BaseCfg): "right_elbow_link", "right_wrist_yaw_link", ) - sampling_mode: str = "adaptive" # "start", "uniform", "adaptive" + sampling_mode: Literal["start", "uniform", "adaptive"] = "adaptive" max_episode_seconds: float = 10.0 reward_config: RewardConfig = field(default_factory=RewardConfig) pose_randomization: PoseRandomization = field(default_factory=PoseRandomization) @@ -435,7 +436,9 @@ def _reward_motion_global_root_pos(self, info: dict) -> np.ndarray: anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] error = np.sum(np.square(anchor_pos_w - robot_anchor_pos_w), axis=-1) - return np.exp(-error / self._cfg.reward_config.std_root_pos**2) + return np.asarray( + np.exp(-error / self._cfg.reward_config.std_root_pos**2), dtype=get_global_dtype() + ) def _reward_motion_global_root_ori(self, info: dict) -> np.ndarray: motion_data = info["motion_data"] @@ -448,28 +451,40 @@ def _reward_motion_global_root_ori(self, info: dict) -> np.ndarray: def _reward_motion_body_pos(self, info: dict) -> np.ndarray: robot_body_pos_w = info["robot_body_pos_w"] error = np.sum(np.square(self.body_pos_relative_w - robot_body_pos_w), axis=-1) - return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_pos**2) + return np.asarray( + np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_pos**2), + dtype=get_global_dtype(), + ) def _reward_motion_body_ori(self, info: dict) -> np.ndarray: robot_body_quat_w = info["robot_body_quat_w"] error = quat_error_magnitude(self.body_quat_relative_w, robot_body_quat_w) ** 2 - return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_ori**2) + return np.asarray( + np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_ori**2), + dtype=get_global_dtype(), + ) def _reward_motion_body_lin_vel(self, info: dict) -> np.ndarray: motion_data = info["motion_data"] robot_body_lin_vel_w = info["robot_body_lin_vel_w"] error = np.sum(np.square(motion_data.body_lin_vel_w - robot_body_lin_vel_w), axis=-1) - return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_lin_vel**2) + return np.asarray( + np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_lin_vel**2), + dtype=get_global_dtype(), + ) def _reward_motion_body_ang_vel(self, info: dict) -> np.ndarray: motion_data = info["motion_data"] robot_body_ang_vel_w = info["robot_body_ang_vel_w"] error = np.sum(np.square(motion_data.body_ang_vel_w - robot_body_ang_vel_w), axis=-1) - return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_ang_vel**2) + return np.asarray( + np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_ang_vel**2), + dtype=get_global_dtype(), + ) def _reward_action_rate_l2(self, info: dict) -> np.ndarray: action_diff = info["current_actions"] - info["last_actions"] - return np.sum(np.square(action_diff), axis=1) + return np.asarray(np.sum(np.square(action_diff), axis=1), dtype=get_global_dtype()) def _reward_joint_limit(self, info: dict) -> np.ndarray: dof_pos = info["dof_pos"] @@ -483,7 +498,7 @@ def _reward_joint_limit(self, info: dict) -> np.ndarray: lower_violation = np.maximum(0, lower - dof_pos) upper_violation = np.maximum(0, dof_pos - upper) violation = lower_violation + upper_violation - return np.sum(np.square(violation), axis=1) + return np.asarray(np.sum(np.square(violation), axis=1), dtype=get_global_dtype()) def reset(self, env_indices: np.ndarray): """Reset specified environments.""" diff --git a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py index c355b3cf6..9c29e1ad3 100644 --- a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py +++ b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py @@ -263,9 +263,7 @@ def _get_obs(self, info: dict) -> dict[str, np.ndarray]: # ── Reset ──────────────────────────────────────────────────────── - def reset( - self, env_indices: np.ndarray - ) -> Tuple[dict[str, np.ndarray], dict]: + def reset(self, env_indices: np.ndarray) -> Tuple[dict[str, np.ndarray], dict]: num_reset = len(env_indices) dr = self._cfg.domain_rand diff --git a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py index 6591b2d7f..d316e82b0 100644 --- a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py +++ b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Tuple import numpy as np from etils import epath diff --git a/src/unilab/ipc/replay_buffer.py b/src/unilab/ipc/replay_buffer.py index ec577edac..39aee9338 100644 --- a/src/unilab/ipc/replay_buffer.py +++ b/src/unilab/ipc/replay_buffer.py @@ -157,6 +157,7 @@ def add( self.dones[idx : idx + n] = dones self.truncated[idx : idx + n] = truncated if self._privileged_dim > 0 and privileged is not None: + assert next_privileged is not None self.privileged_obs[idx : idx + n] = privileged self.next_privileged_obs[idx : idx + n] = next_privileged else: @@ -174,6 +175,7 @@ def add( self.truncated[idx:] = truncated[:split] self.truncated[: n - split] = truncated[split:] if self._privileged_dim > 0 and privileged is not None: + assert next_privileged is not None self.privileged_obs[idx:] = privileged[:split] self.privileged_obs[: n - split] = privileged[split:] self.next_privileged_obs[idx:] = next_privileged[:split] @@ -188,6 +190,7 @@ def add( truncated.unsqueeze(1), ] if self._privileged_dim > 0 and privileged is not None: + assert next_privileged is not None parts.extend([privileged, next_privileged]) row = torch.cat(parts, dim=1) diff --git a/src/unilab/utils/device_utils.py b/src/unilab/utils/device_utils.py index fcc63d392..4eff093ea 100644 --- a/src/unilab/utils/device_utils.py +++ b/src/unilab/utils/device_utils.py @@ -22,7 +22,9 @@ def get_env_dims( """ from unilab.utils.obs_utils import get_obs_dims as get_obs_dims_from_spec - env = registry.make(env_name, num_envs=1, sim_backend=sim_backend, env_cfg_override=env_cfg_override) + env = registry.make( + env_name, num_envs=1, sim_backend=sim_backend, env_cfg_override=env_cfg_override + ) obs_dim, privileged_dim = get_obs_dims_from_spec(env.obs_groups_spec) action_shape = env.action_space.shape assert action_shape is not None diff --git a/src/unilab/utils/reward_utils.py b/src/unilab/utils/reward_utils.py index 746d4bc2d..0f358d077 100644 --- a/src/unilab/utils/reward_utils.py +++ b/src/unilab/utils/reward_utils.py @@ -16,7 +16,9 @@ def extract_reward_config(cfg: DictConfig) -> dict: ValueError: If reward config is missing """ if not hasattr(cfg, "reward") or not cfg.reward: - raise ValueError("Missing 'reward' config in Hydra. Reward config must be explicitly provided.") + raise ValueError( + "Missing 'reward' config in Hydra. Reward config must be explicitly provided." + ) reward_dict = OmegaConf.to_container(cfg.reward, resolve=True) return {"reward_config": reward_dict} diff --git a/src/unilab/utils/xml_utils.py b/src/unilab/utils/xml_utils.py index b72507cef..5994bcbc7 100644 --- a/src/unilab/utils/xml_utils.py +++ b/src/unilab/utils/xml_utils.py @@ -3,15 +3,17 @@ import os import tempfile import xml.etree.ElementTree as ET +from typing import Any, cast import mujoco def _get_named_bodies(model_file: str) -> tuple[list[int], list[str]]: - _m = mujoco.MjModel.from_xml_path(model_file) + mj = cast(Any, mujoco) + _m = mj.MjModel.from_xml_path(model_file) ids, names = [], [] for i in range(1, _m.nbody): # skip body 0 (world body) - name = mujoco.mj_id2name(_m, mujoco.mjtObj.mjOBJ_BODY, i) + name = mj.mj_id2name(_m, mj.mjtObj.mjOBJ_BODY, i) if name: ids.append(i) names.append(name) diff --git a/tests/algos/test_mlx_ppo.py b/tests/algos/test_mlx_ppo.py index 7982b1416..bcad57ddf 100644 --- a/tests/algos/test_mlx_ppo.py +++ b/tests/algos/test_mlx_ppo.py @@ -327,7 +327,12 @@ def test_mlx_ppo_one_iteration_real_env(default_go2_reward_config): cfg = PPOStructuredConfig() algo_cfg = cfg.algorithm - env = registry.make(env_name, num_envs=num_envs, sim_backend="mujoco", env_cfg_override={"reward_config": default_go2_reward_config}) + env = registry.make( + env_name, + num_envs=num_envs, + sim_backend="mujoco", + env_cfg_override={"reward_config": default_go2_reward_config}, + ) obs_dim = sum(env.obs_groups_spec.values()) action_dim = env.action_space.shape[0] diff --git a/tests/algos/test_rsl_rl_runner.py b/tests/algos/test_rsl_rl_runner.py index a8d7aef3c..bc87f31d8 100644 --- a/tests/algos/test_rsl_rl_runner.py +++ b/tests/algos/test_rsl_rl_runner.py @@ -113,12 +113,16 @@ def get_privileged_observations(self): "G1JoystickFlatTerrain", ], ) -def test_rsl_rl_ppo_one_iteration(env_name: str, default_go2_reward_config, default_g1_reward_config): +def test_rsl_rl_ppo_one_iteration( + env_name: str, default_go2_reward_config, default_g1_reward_config +): """RSL-RL PPO can complete 1 training iteration on a real env.""" from rsl_rl.runners import OnPolicyRunner reward_cfg = default_go2_reward_config if "Go2" in env_name else default_g1_reward_config - env = registry.make(env_name, num_envs=256, sim_backend="mujoco", env_cfg_override={"reward_config": reward_cfg}) + env = registry.make( + env_name, num_envs=256, sim_backend="mujoco", env_cfg_override={"reward_config": reward_cfg} + ) wrapped = _RslRlVecEnvWrapper(env, device="cpu") cfg = PPOConfig() diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 5fe1ecf6f..bfa6bd9e9 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -492,15 +492,14 @@ def test_default_reward_extraction_yields_empty(): def test_appo_reward_extraction_pattern(): - """Verify APPO with go1 task extracts reward (even though only go1 has a custom reward).""" + """Verify APPO go1 task extracts its default reward override.""" cfg = _compose("appo", overrides=["task=go1_joystick"]) - # APPO go1 uses default reward (no override), so env_cfg_override should be empty env_cfg_override: dict = {} if hasattr(cfg, "reward") and cfg.reward: reward_dict = OmegaConf.to_container(cfg.reward, resolve=True) env_cfg_override["reward_config"] = reward_dict - # Currently go1_joystick in APPO uses default.yaml (empty) - assert "reward_config" not in env_cfg_override + assert "reward_config" in env_cfg_override + assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(1.0) def test_appo_custom_reward_extraction(): diff --git a/tests/conftest.py b/tests/conftest.py index a72bfa616..f36b2ccc2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -197,7 +197,37 @@ def default_g1_reward_config(): "base_height_target": 0.754, "min_base_height": 0.55, "max_tilt_deg": 25.0, - "pose_weights": [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0], + "pose_weights": [ + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + ], } @@ -224,6 +254,35 @@ def default_g1_sac_reward_config(): "feet_phase_swing_height": 0.09, "feet_phase_tracking_sigma": 0.008, "close_feet_threshold": 0.15, - "pose_weights": [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0], + "pose_weights": [ + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + ], } - diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 94c74d9ed..0235c7c35 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -126,7 +126,9 @@ def test_env_reset_and_step( elif "G1" in env_name: env_cfg_override = {"reward_config": default_g1_reward_config} - env = registry.make(env_name, num_envs=2, sim_backend="mujoco", env_cfg_override=env_cfg_override) + env = registry.make( + env_name, num_envs=2, sim_backend="mujoco", env_cfg_override=env_cfg_override + ) try: # 1. Spaces obs_space = env.observation_space diff --git a/tests/integration/test_reward_injection_integration.py b/tests/integration/test_reward_injection_integration.py index 2d5573b83..adec345d2 100644 --- a/tests/integration/test_reward_injection_integration.py +++ b/tests/integration/test_reward_injection_integration.py @@ -90,7 +90,7 @@ def test_backward_compatibility_no_reward_config(): # Should fail without reward_config with pytest.raises(ValueError, match="reward_config must be provided"): - env = registry.make( + registry.make( "Go1JoystickFlatTerrain", num_envs=2, sim_backend="mujoco", diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 0b49695b6..c4ab240d6 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -62,6 +62,24 @@ def _offpolicy_cfg(overrides=None): return compose("config", overrides=overrides or []) +def _ppo_cfg(overrides=None): + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): + return compose("config", overrides=overrides or []) + + +def _train_rsl_rl(monkeypatch: pytest.MonkeyPatch): + import types + + runners_mod = types.ModuleType("rsl_rl.runners") + runners_mod.OnPolicyRunner = object + rsl_pkg = types.ModuleType("rsl_rl") + rsl_pkg.runners = runners_mod + monkeypatch.setitem(sys.modules, "rsl_rl", rsl_pkg) + monkeypatch.setitem(sys.modules, "rsl_rl.runners", runners_mod) + return _load_script("train_rsl_rl") + + def test_offpolicy_hydra_default_algo(): cfg = _offpolicy_cfg() assert cfg.algo.algo == "sac" @@ -94,6 +112,51 @@ def test_offpolicy_hydra_algo_td3(): assert cfg.algo.algo == "td3" +def test_offpolicy_task_go1_exposes_motrix_legacy(): + cfg = _offpolicy_cfg(["task=go1_joystick"]) + + assert cfg.motrix_legacy.enabled is True + assert cfg.motrix_legacy.applies_to.algo == "sac" + assert cfg.motrix_legacy.algo_overrides.num_envs == 4096 + assert cfg.motrix_legacy.algo_overrides.max_iterations == 2000 + assert cfg.motrix_legacy.env_cfg_override.legacy_motrix_profile.enabled is True + + +def test_offpolicy_build_task_motrix_env_cfg_override_applies_go1_legacy(): + cfg = _offpolicy_cfg(["task=go1_joystick", "training.sim_backend=motrix"]) + + env_cfg_override = _offpolicy().build_task_motrix_offpolicy_env_cfg_override("sac", cfg) + + assert cfg.algo.num_envs == 4096 + assert cfg.algo.max_iterations == 2000 + assert env_cfg_override["legacy_motrix_profile"]["enabled"] is True + assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(1.0) + + +def test_offpolicy_build_task_motrix_env_cfg_override_skips_td3(): + cfg = _offpolicy_cfg(["algo=td3", "task=go1_joystick", "training.sim_backend=motrix"]) + + env_cfg_override = _offpolicy().build_task_motrix_offpolicy_env_cfg_override("td3", cfg) + + assert cfg.algo.num_envs == 2048 + assert cfg.algo.max_iterations == 3000 + assert "legacy_motrix_profile" not in env_cfg_override + + +def test_offpolicy_build_task_motrix_env_cfg_override_respects_cli_algo_override( + monkeypatch: pytest.MonkeyPatch, +): + cfg = _offpolicy_cfg( + ["task=go1_joystick", "training.sim_backend=motrix", "algo.max_iterations=1"] + ) + monkeypatch.setattr(sys, "argv", ["train_offpolicy.py", "algo.max_iterations=1"]) + + _offpolicy().build_task_motrix_offpolicy_env_cfg_override("sac", cfg) + + assert cfg.algo.num_envs == 4096 + assert cfg.algo.max_iterations == 1 + + def test_offpolicy_resolve_sac_use_symmetry_keeps_mujoco_setting(): cfg = _offpolicy_cfg(["task=g1_sac", "training.sim_backend=mujoco"]) @@ -106,6 +169,74 @@ def test_offpolicy_resolve_sac_use_symmetry_disables_motrix(): assert _offpolicy().resolve_sac_use_symmetry(cfg) is False +def test_ppo_task_go1_exposes_motrix_legacy(): + cfg = _ppo_cfg(["task=go1_joystick"]) + + assert cfg.motrix_legacy.enabled is True + assert cfg.motrix_legacy.algo_overrides.max_iterations == 151 + assert cfg.motrix_legacy.algo_overrides.policy.init_noise_std == pytest.approx(0.5) + assert cfg.motrix_legacy.env_cfg_override.legacy_motrix_profile.enabled is True + + +def test_ppo_task_g1_exposes_motrix_legacy(): + cfg = _ppo_cfg(["task=g1_joystick"]) + + assert cfg.motrix_legacy.enabled is True + assert cfg.motrix_legacy.algo_overrides.obs_groups.actor == ["policy"] + assert ( + cfg.motrix_legacy.env_cfg_override.backend_overrides.control_action_scale + == pytest.approx(0.5) + ) + + +def test_build_task_motrix_ppo_env_cfg_override_applies_go1_legacy( + monkeypatch: pytest.MonkeyPatch, +): + mod = _train_rsl_rl(monkeypatch) + cfg = _ppo_cfg(["task=go1_joystick", "training.sim_backend=motrix"]) + + env_cfg_override = mod.build_task_motrix_ppo_env_cfg_override(cfg) + + assert cfg.algo.max_iterations == 151 + assert cfg.algo.empirical_normalization is True + assert cfg.algo.policy.init_noise_std == pytest.approx(0.5) + assert cfg.algo.algorithm.learning_rate == pytest.approx(3.0e-4) + assert cfg.algo.algorithm.entropy_coef == pytest.approx(1.0e-3) + assert env_cfg_override["legacy_motrix_profile"]["enabled"] is True + assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(1.0) + + +def test_build_task_motrix_ppo_env_cfg_override_applies_g1_legacy( + monkeypatch: pytest.MonkeyPatch, +): + mod = _train_rsl_rl(monkeypatch) + cfg = _ppo_cfg(["task=g1_joystick", "training.sim_backend=motrix"]) + + env_cfg_override = mod.build_task_motrix_ppo_env_cfg_override(cfg) + + assert cfg.algo.max_iterations == 151 + assert cfg.algo.empirical_normalization is True + assert cfg.algo.policy.init_noise_std == pytest.approx(0.5) + assert cfg.algo.algorithm.learning_rate == pytest.approx(3.0e-4) + assert cfg.algo.algorithm.entropy_coef == pytest.approx(5.0e-3) + assert cfg.algo.obs_groups.actor == ["policy"] + assert env_cfg_override["backend_overrides"]["enabled"] is True + assert env_cfg_override["reward_config"]["scales"]["upper_body_pose"] == pytest.approx(-0.05) + + +def test_build_task_motrix_ppo_env_cfg_override_respects_cli_algo_override( + monkeypatch: pytest.MonkeyPatch, +): + mod = _train_rsl_rl(monkeypatch) + cfg = _ppo_cfg(["task=g1_joystick", "training.sim_backend=motrix", "algo.max_iterations=1"]) + monkeypatch.setattr(sys, "argv", ["train_rsl_rl.py", "algo.max_iterations=1"]) + + mod.build_task_motrix_ppo_env_cfg_override(cfg) + + assert cfg.algo.empirical_normalization is True + assert cfg.algo.max_iterations == 1 + + # --------------------------------------------------------------------------- # train_offpolicy.py — default_device() # --------------------------------------------------------------------------- @@ -219,12 +350,11 @@ def test_offpolicy_extract_reset_obs_handles_two_tuple(): assert result is obs -def test_offpolicy_extract_reset_obs_handles_three_tuple(): +def test_offpolicy_extract_reset_obs_rejects_three_tuple(): obs = {"obs": "value"} - result = _offpolicy().extract_reset_obs(("ignored", obs, {"info": 1})) - - assert result is obs + with pytest.raises(ValueError, match="Unexpected env.reset return format"): + _offpolicy().extract_reset_obs(("ignored", obs, {"info": 1})) def test_offpolicy_resolve_play_obs_dim_ignores_privileged(): From 0083cf2532ef40cb88ba04eb1a425ab86689eecc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=A2=E6=9D=B0?= Date: Thu, 19 Mar 2026 11:44:29 +0800 Subject: [PATCH 4/5] fix: align CI mypy path with src layout --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 79888e73e..6dd79b2e7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,7 +42,7 @@ jobs: - name: Install dependencies run: uv sync --extra dev - name: mypy - run: uv run mypy unilab + run: uv run mypy src/unilab - name: pyright run: uv run pyright From 86661d693639ee42f3073a8134000ffae9edd2e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=A2=E6=9D=B0?= Date: Thu, 19 Mar 2026 11:58:04 +0800 Subject: [PATCH 5/5] fix: make ndarray return types explicit for CI mypy --- src/unilab/envs/locomotion/g1/joystick.py | 2 +- src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 836809cf3..9d55a6f84 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -77,7 +77,7 @@ def sample_reset_base_qvel(rng, num_samples: int, limit: float) -> np.ndarray: def build_upper_body_pose_weights(pose_weights: list[float]) -> np.ndarray: weights = np.asarray(pose_weights, dtype=get_global_dtype()).copy() weights[:12] = 0.0 - return weights + return np.asarray(weights, dtype=get_global_dtype()) def compute_weighted_pose_penalty(diff: np.ndarray, weights: np.ndarray) -> np.ndarray: diff --git a/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py b/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py index a776a188b..a46c52322 100644 --- a/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py +++ b/src/unilab/envs/locomotion/g1/motion_tracking/math_utils.py @@ -50,9 +50,9 @@ def quat_conjugate(q: np.ndarray) -> np.ndarray: Conjugate quaternion(s) (N, 4) or (4,) """ if q.ndim == 1: - return np.array([q[0], -q[1], -q[2], -q[3]]) + return np.asarray([q[0], -q[1], -q[2], -q[3]], dtype=q.dtype) else: - result = q.copy() + result = np.asarray(q.copy(), dtype=q.dtype) result[:, 1:] *= -1 return result