add Go2JoystickFlat task owner and retire G1_walk_flat_amp task

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/<task>/` |
 | APPO | `scripts/train_appo.py` | `logs/appo/<task>/` |
 | SAC | `scripts/train_offpolicy.py` | `logs/fast_sac/<task>/` |
+| FlashSAC | `scripts/train_offpolicy.py` | `logs/flash_sac/<task>/` |
 | TD3 | `scripts/train_offpolicy.py` | `logs/fast_td3/<task>/` |
 
 Training scripts automatically enter playback after training unless you set `training.no_play=true`.

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

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

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

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
 

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,
         }

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",

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

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,
-    },
 ]
 
 

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",
         ]
     )