diff --git a/README.md b/README.md index 147e2e225..5dcf38a86 100644 --- a/README.md +++ b/README.md @@ -90,6 +90,7 @@ uv run scripts/train_rsl_rl.py task=g1_motion_tracking/mujoco | PPO (MLX, macOS) | `scripts/train_mlx_ppo.py` | `logs/mlx_rl_train//` | | APPO | `scripts/train_appo.py` | `logs/appo//` | | SAC | `scripts/train_offpolicy.py` | `logs/fast_sac//` | +| FlashSAC | `scripts/train_offpolicy.py` | `logs/flash_sac//` | | TD3 | `scripts/train_offpolicy.py` | `logs/fast_td3//` | Training scripts automatically enter playback after training unless you set `training.no_play=true`. diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml index 76903fffc..621211530 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml @@ -7,6 +7,9 @@ algo: max_iterations: 5000 save_interval: 1000 env: + domain_rand: + randomize_kp: false + randomize_kd: false curriculum: enabled: true initial_scale: 0.5 diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat_amp/mujoco.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat_amp/mujoco.yaml deleted file mode 100644 index 88d54ab1b..000000000 --- a/conf/offpolicy/task/flashsac/g1_walk_flat_amp/mujoco.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# @package _global_ -training: - task_name: G1WalkFlat - sim_backend: mujoco - use_amp: true -algo: - num_envs: 1024 - replay_buffer_n: 9766 - learning_starts: 98 - updates_per_step: 2 - max_iterations: 48829 - save_interval: 5000 - gamma: 0.97 -env: - control_config: - action_scale: 0.5 - commands: - vel_limit: - - [-1.0, -0.5, -1.0] - - [1.0, 0.5, 1.0] - curriculum: - enabled: false -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.75 - feet_phase: 1.0 - lin_vel_z: 0.0 - ang_vel_xy: -0.15 - base_height: 0.0 - orientation: -2.0 - action_rate: 0.0 - pose: -0.1 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.15 - feet_phase_tracking_sigma: 0.01 - 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] diff --git a/conf/offpolicy/task/flashsac/go2_joystick_flat/motrix.yaml b/conf/offpolicy/task/flashsac/go2_joystick_flat/motrix.yaml new file mode 100644 index 000000000..5a5e34692 --- /dev/null +++ b/conf/offpolicy/task/flashsac/go2_joystick_flat/motrix.yaml @@ -0,0 +1,43 @@ +# @package _global_ +training: + task_name: Go2JoystickFlat + sim_backend: motrix +algo: + num_envs: 2048 + learning_starts: 49 + max_iterations: 3000 + save_interval: 1000 + updates_per_step: 8 + replay_buffer_n: 1024 + tau: 0.05 +env: + control_config: + action_scale: 1.0 + domain_rand: + randomize_kp: false + randomize_kd: false + randomize_base_mass: false + random_com: false + randomize_gravity: false + push_robots: false + noise_config: + level: 0.0 + scale_joint_angle: 0.01 + scale_joint_vel: 0.1 + scale_gyro: 0.0 + scale_gravity: 0.0 + scale_linvel: 0.0 +reward: + scales: + tracking_lin_vel: 3.0 + tracking_ang_vel: 0.4 + lin_vel_z: -5.0 + ang_vel_xy: -0.1 + base_height: -15.0 + action_rate: -0.005 + similar_to_default: -0.3 + contact: 0.24 + swing_feet_z: 4.0 + alive: 2.0 + tracking_sigma: 0.25 + base_height_target: 0.3 diff --git a/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml b/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml new file mode 100644 index 000000000..031d8e2b7 --- /dev/null +++ b/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml @@ -0,0 +1,43 @@ +# @package _global_ +training: + task_name: Go2JoystickFlat + sim_backend: mujoco +algo: + num_envs: 1024 + learning_starts: 50 + max_iterations: 4000 + save_interval: 1000 + updates_per_step: 2 + batch_size: 2048 + replay_buffer_n: 4096 + tau: 0.05 +env: + control_config: + action_scale: 0.4 + domain_rand: + randomize_kp: true + randomize_kd: true + randomize_base_mass: true + random_com: true + randomize_gravity: true + push_robots: true + noise_config: + level: 1.0 + scale_joint_angle: 0.01 + scale_joint_vel: 0.1 + scale_gyro: 0.0 + scale_gravity: 0.0 + scale_linvel: 0.0 +reward: + scales: + tracking_lin_vel: 1.0 + tracking_ang_vel: 0.2 + lin_vel_z: -5.0 + ang_vel_xy: -0.1 + base_height: -20.0 + action_rate: -0.02 + similar_to_default: -0.4 + contact: 1.5 + swing_feet_z: 4.0 + tracking_sigma: 0.4 + base_height_target: 0.3 \ No newline at end of file diff --git a/docs/users/zh_CN/02-simulation-backends.md b/docs/users/zh_CN/02-simulation-backends.md index 1b05c0cc0..63ac6f562 100644 --- a/docs/users/zh_CN/02-simulation-backends.md +++ b/docs/users/zh_CN/02-simulation-backends.md @@ -71,8 +71,8 @@ uv run scripts/generate_support_matrix.py --write | SAC (torch) | `g1_walk_rough` (G1 walk rough) | Tested | Registered | | SAC (torch) | `g1_sac_wbt` (g1 sac wbt) | Tested | - | | TD3 (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Registered | +| FlashSAC (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | Tested | | FlashSAC (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Tested | -| FlashSAC (torch) | `g1_walk_flat_amp` (g1 walk flat amp) | Tested | Registered | ### Source Index diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py index 9ef9c936d..29b523dc6 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/envs/locomotion/go2/joystick.py @@ -121,7 +121,7 @@ def _init_reward_functions(self): "base_height": rewards.base_height, "action_rate": rewards.action_rate, "similar_to_default": rewards.similar_to_default, - # "alive": rewards.alive, + "alive": rewards.alive, "swing_feet_z": self._reward_swing_feet_z, "contact": self._reward_contact, } diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index b53f3f737..5ee93e030 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -266,6 +266,45 @@ def test_offpolicy_g1_walk_flat_motrix_preserves_backend_env_overrides(): assert cfg.env.domain_rand.randomize_kd is False +def test_offpolicy_flashsac_go2_joystick_backend_overrides_diverge(): + mujoco_cfg = _compose( + "offpolicy", + overrides=["algo=flashsac", "task=flashsac/go2_joystick_flat/mujoco"], + ) + motrix_cfg = _compose( + "offpolicy", + overrides=["algo=flashsac", "task=flashsac/go2_joystick_flat/motrix"], + ) + + assert mujoco_cfg.training.task_name == "Go2JoystickFlat" + assert motrix_cfg.training.task_name == "Go2JoystickFlat" + assert mujoco_cfg.training.sim_backend == "mujoco" + assert motrix_cfg.training.sim_backend == "motrix" + + # MuJoCo enables the full DR stack; Motrix disables what its backend cannot + # honor (e.g. gravity, kp/kd, push) — see ADR-0002 backend capability boundary. + assert mujoco_cfg.env.domain_rand.randomize_kp is True + assert mujoco_cfg.env.domain_rand.randomize_kd is True + assert mujoco_cfg.env.domain_rand.randomize_base_mass is True + assert mujoco_cfg.env.domain_rand.random_com is True + assert mujoco_cfg.env.domain_rand.randomize_gravity is True + assert mujoco_cfg.env.domain_rand.push_robots is True + assert mujoco_cfg.env.noise_config.level == pytest.approx(1.0) + + assert motrix_cfg.env.domain_rand.randomize_kp is False + assert motrix_cfg.env.domain_rand.randomize_kd is False + assert motrix_cfg.env.domain_rand.randomize_base_mass is False + assert motrix_cfg.env.domain_rand.random_com is False + assert motrix_cfg.env.domain_rand.randomize_gravity is False + assert motrix_cfg.env.domain_rand.push_robots is False + assert motrix_cfg.env.noise_config.level == pytest.approx(0.0) + + # The alive reward scale must reach both backends — its registration was + # re-enabled in Go2WalkTask._init_reward_functions. + assert mujoco_cfg.reward.scales.alive == pytest.approx(2.0) + assert motrix_cfg.reward.scales.alive == pytest.approx(2.0) + + def test_cli_override_beats_task_defaults(): cfg = _compose( "ppo", diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index d6566514d..2cb966a6e 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -162,51 +162,45 @@ def test_offpolicy_flashsac_g1_task_overrides(): assert cfg.algo.algo_params.normalize_reward is True -def test_offpolicy_g1_rough_terrain_task_overrides(): +def test_offpolicy_flashsac_go2_task_overrides(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra - from unilab.envs.locomotion.g1.joystick import G1WalkRoughCfg - GlobalHydra.instance().clear() with initialize_config_dir(config_dir=str(CONF_DIR / "offpolicy"), version_base="1.3"): cfg = compose( "config", - overrides=["algo=sac", "task=sac/g1_walk_rough/mujoco"], + overrides=["algo=flashsac", "task=flashsac/go2_joystick_flat/mujoco"], ) - assert cfg.algo.algo == "sac" - assert cfg.training.task_name == "G1WalkRough" + assert cfg.algo.algo == "flashsac" + assert cfg.training.task_name == "Go2JoystickFlat" assert cfg.training.sim_backend == "mujoco" - assert G1WalkRoughCfg().model_file.endswith("scene_rough.xml") + assert cfg.algo.num_envs == 2048 + assert cfg.algo.max_iterations == 3000 + assert cfg.algo.tau == pytest.approx(0.05) + assert cfg.algo.replay_buffer_n == 1024 + assert cfg.algo.updates_per_step == 8 + assert cfg.reward.scales.alive == pytest.approx(2.0) + assert cfg.reward.scales.swing_feet_z == pytest.approx(4.0) + assert cfg.env.control_config.action_scale == pytest.approx(1.0) -def test_offpolicy_flashsac_g1_walk_flat_amp_task_overrides(): +def test_offpolicy_g1_rough_terrain_task_overrides(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra + from unilab.envs.locomotion.g1.joystick import G1WalkRoughCfg + GlobalHydra.instance().clear() with initialize_config_dir(config_dir=str(CONF_DIR / "offpolicy"), version_base="1.3"): cfg = compose( "config", - overrides=["algo=flashsac", "task=flashsac/g1_walk_flat_amp/mujoco"], + overrides=["algo=sac", "task=sac/g1_walk_rough/mujoco"], ) - assert cfg.algo.algo == "flashsac" - assert cfg.training.task_name == "G1WalkFlat" + assert cfg.algo.algo == "sac" + assert cfg.training.task_name == "G1WalkRough" assert cfg.training.sim_backend == "mujoco" - assert cfg.training.use_amp is True - assert cfg.algo.num_envs == 1024 - assert cfg.algo.updates_per_step == 2 - assert cfg.algo.num_envs == 1024 - assert cfg.algo.replay_buffer_n == 9766 - assert cfg.algo.learning_starts == 98 - assert cfg.algo.updates_per_step == 2 - assert cfg.algo.gamma == pytest.approx(0.97) - assert cfg.env.control_config.action_scale == pytest.approx(0.5) - assert cfg.env.commands.vel_limit[0] == [-1.0, -0.5, -1.0] - assert "obs_profile" not in cfg.env - assert cfg.env.curriculum.enabled is False - assert cfg.reward.scales.tracking_ang_vel == pytest.approx(0.75) - assert cfg.reward.scales.base_height == pytest.approx(0.0) + assert G1WalkRoughCfg().model_file.endswith("scene_rough.xml") def test_g1_task_owner_yamls_preserve_legacy_and_walk_observation_profiles(): @@ -226,14 +220,6 @@ def uses_walk_profile(config_group: str, overrides: list[str]) -> bool: assert uses_walk_profile("ppo", ["task=g1_walk_flat/mujoco"]) is False assert uses_walk_profile("appo", ["task=g1_walk_flat/mujoco"]) is False - assert ( - uses_walk_profile( - "offpolicy", - ["algo=flashsac", "task=flashsac/g1_walk_flat_amp/mujoco"], - ) - is False - ) - assert uses_walk_profile("offpolicy", ["algo=sac", "task=sac/g1_walk_flat/mujoco"]) is True assert uses_walk_profile("offpolicy", ["algo=sac", "task=sac/g1_walk_flat/motrix"]) is True assert uses_walk_profile("offpolicy", ["algo=sac", "task=sac/g1_walk_rough/mujoco"]) is True diff --git a/tests/envs/locomotion/g1/test_issue175_regression.py b/tests/envs/locomotion/g1/test_issue175_regression.py index fb8263f2b..b2171d6e2 100644 --- a/tests/envs/locomotion/g1/test_issue175_regression.py +++ b/tests/envs/locomotion/g1/test_issue175_regression.py @@ -109,16 +109,6 @@ "action_scale": 1.0, "curriculum_enabled": True, }, - { - "id": "flashsac_amp_mujoco", - "config_group": "offpolicy", - "overrides": ["algo=flashsac", "task=flashsac/g1_walk_flat_amp/mujoco"], - "task_name": "G1WalkFlat", - "backend": "mujoco", - "profile": "legacy", - "action_scale": 0.5, - "curriculum_enabled": False, - }, ] diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index c01b5feb6..02d519052 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -1269,13 +1269,13 @@ def test_offpolicy_td3_hydra_default_algo_log_name(): def test_offpolicy_flashsac_hydra_algo_log_name(): - cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_walk_flat_amp/mujoco"]) + cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_walk_flat/mujoco"]) assert cfg.algo.algo_log_name == "flash_sac" assert cfg.algo.load_run == "-1" -def test_offpolicy_flashsac_g1_walk_flat_amp_task_composes() -> None: - cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_walk_flat_amp/mujoco"]) +def test_offpolicy_flashsac_g1_walk_flat_task_composes() -> None: + cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_walk_flat/mujoco"]) assert cfg.training.task_name == "G1WalkFlat" assert cfg.training.sim_backend == "mujoco" @@ -1291,7 +1291,7 @@ def test_offpolicy_flashsac_rejects_multi_gpu(): cfg = _offpolicy_cfg( [ "algo=flashsac", - "task=flashsac/g1_walk_flat_amp/mujoco", + "task=flashsac/g1_walk_flat/mujoco", "training.num_gpus=2", ] )