Feat/privileged obs refactor

README.md

@@ -148,10 +148,10 @@ uv run python scripts/train_rsl_rl.py task=go1_joystick training.play_only=true
 
 ```bash
 # PPO (RSL-RL) - 使用 Hydra 配置
-uv run python scripts/train_rsl_rl.py task=go2_joystick
+uv run python scripts/train_rsl_rl.py task=go1_joystick
 
 # PPO (MLX - Apple Silicon)
-uv run python scripts/train_mlx_ppo.py task=go2_joystick
+uv run python scripts/train_mlx_ppo.py task=go1_joystick
 
 # APPO（异步 PPO，CPU/GPU 并行）
 uv run python scripts/train_appo.py task=go1_joystick

conf/appo/reward/g1_appo_mujoco.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+reward:
+  scales:
+    tracking_lin_vel: 2.0
+    tracking_ang_vel: 0.2
+    feet_phase: 1.0
+    lin_vel_z: -1.0
+    ang_vel_xy: -0.25
+    base_height: -500.0
+    orientation: -5.0
+    action_rate: -0.01
+    pose: -0.1
+  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
+  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]

conf/appo/reward/go2_appo_mujoco.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+reward:
+  scales:
+    tracking_lin_vel: 1.0
+    tracking_ang_vel: 0.2
+    lin_vel_z: -5.0
+    ang_vel_xy: -0.02
+    base_height: -100.0
+    action_rate: -0.005
+    similar_to_default: -0.1
+    alive: 0.0
+    foot_lift_reward: 0.2
+    foot_drag_penalty: -0.0
+  tracking_sigma: 0.25
+  base_height_target: 0.3
+  target_foot_height: 0.08
+  foot_clearance_sigma: 0.02

conf/appo/task/g1_joystick.yaml

@@ -1,4 +1,7 @@
 # @package _global_
+defaults:
+  - override /reward: g1_appo_mujoco
+
 training:
   task_name: G1JoystickFlatTerrain
 algo:

conf/appo/task/go1_joystick.yaml

@@ -1,3 +1,6 @@
 # @package _global_
+defaults:
+  - override /reward: go1_appo_mujoco
+
 training:
   task_name: Go1JoystickFlatTerrain

conf/appo/task/go2_joystick.yaml

@@ -1,3 +1,6 @@
 # @package _global_
+defaults:
+  - override /reward: go2_appo_mujoco
+
 training:
   task_name: Go2JoystickFlatTerrain

conf/offpolicy/reward/g1_sac_motrix.yaml

@@ -17,3 +17,6 @@ reward:
   gait_frequency: 1.5
   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]
+

conf/offpolicy/reward/g1_sac_mujoco.yaml

@@ -17,3 +17,6 @@ reward:
   gait_frequency: 1.5
   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]
+

conf/offpolicy/reward/go2_sac_mujoco.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+reward:
+  scales:
+    tracking_lin_vel: 1.0
+    tracking_ang_vel: 0.2
+    lin_vel_z: -5.0
+    ang_vel_xy: -0.02
+    base_height: -100.0
+    action_rate: -0.005
+    similar_to_default: -0.1
+    alive: 0.0
+    foot_lift_reward: 0.2
+    foot_drag_penalty: -0.0
+  tracking_sigma: 0.25
+  base_height_target: 0.3
+  target_foot_height: 0.08
+  foot_clearance_sigma: 0.02

conf/offpolicy/task/go2_joystick.yaml

@@ -1,6 +1,6 @@
 # @package _global_
 defaults:
-  - override /reward: default
+  - override /reward: go2_sac_mujoco
 
 training:
   task_name: Go2JoystickFlatTerrain

conf/ppo/reward/g1_ppo_mujoco.yaml

@@ -17,3 +17,5 @@ reward:
   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]
+

scripts/train_appo.py

@@ -55,12 +55,9 @@ 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
 
-    # Build env_cfg_override from reward config
-    env_cfg_override: dict | None = None
-    if hasattr(cfg, "reward") and cfg.reward:
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override = {"reward_config": reward_dict}
+    env_cfg_override = extract_reward_config(cfg)
 
     device = cfg.training.device or (
         "cuda"
@@ -77,7 +74,8 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
         sim_backend=cfg.training.sim_backend,
         env_cfg_override=env_cfg_override,
     )
-    obs_dim = sum(env.obs_groups_spec.values())
+    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]
 
     rl_cfg_dict = dict(rl_cfg)
@@ -164,8 +162,8 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
     env_indices = np.arange(cfg.training.play_env_num, dtype=np.int32)
     from unilab.utils.obs_utils import flatten_obs_dict
 
-    _, obs_out, _ = env.reset(env_indices)
-    obs_np = np.asarray(flatten_obs_dict(obs_out), dtype=np.float32)
+    obs_out, _ = env.reset(env_indices)
+    obs_np = np.asarray(obs_out["obs"], dtype=np.float32)
 
     state_list = []
     num_steps = 150
@@ -178,7 +176,7 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
             actions_torch = actor(td)
             actions_np = actions_torch.cpu().numpy().astype(np.float32)
             state = env.step(actions_np)
-            obs_np = np.asarray(flatten_obs_dict(state.obs), dtype=np.float32)
+            obs_np = np.asarray(state.obs["obs"], dtype=np.float32)
             state_list.append(np.asarray(env._backend.get_physics_state(), dtype=np.float32).copy())
 
     print("Rendering frames...")
@@ -195,11 +193,9 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
 def main(cfg: DictConfig) -> None:
     ensure_registries()
 
-    # Build env_cfg_override from reward config
-    env_cfg_override = {}
-    if hasattr(cfg, "reward") and cfg.reward:
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override["reward_config"] = reward_dict
+    from unilab.utils.reward_utils import extract_reward_config
+
+    env_cfg_override = extract_reward_config(cfg)
 
     # Convert algo config to plain dict for APPORunner / RSL-RL internals
     rl_cfg = OmegaConf.to_container(cfg.algo, resolve=True)
@@ -235,7 +231,6 @@ def main(cfg: DictConfig) -> None:
             collector_device=collector_device,
             num_envs=cfg.algo.num_envs,
             steps_per_env=cfg.algo.steps_per_env,
-            sim_backend=cfg.training.sim_backend,
             **runner_kwargs,
         )
 

scripts/train_mlx_ppo.py

@@ -157,11 +157,9 @@ def play_mlx_ppo(
     cfg: DictConfig, dtype, use_fp16: bool, resolved_sim_backend: str, task_log_root: Path
 ):
     """Play mode for MLX PPO."""
-    # Build env_cfg_override from reward config
-    env_cfg_override: dict | None = None
-    if hasattr(cfg, "reward") and cfg.reward:
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override = {"reward_config": reward_dict}
+    from unilab.utils.reward_utils import extract_reward_config
+
+    env_cfg_override = extract_reward_config(cfg)
 
     play_model_dtype = mx.float32 if use_fp16 else dtype
     play_env_num = cfg.training.play_env_num
@@ -209,7 +207,7 @@ def play_mlx_ppo(
     if env.state is None:
         env.init_state()
     play_reset_indices = np.arange(env.num_envs, dtype=np.int32)
-    _, obs_dict_play, _ = env.reset(play_reset_indices)
+    obs_dict_play, _ = env.reset(play_reset_indices)
     obs = mx.array(flatten_obs_dict(obs_dict_play))
 
     if resolved_sim_backend == "motrix":
@@ -380,11 +378,9 @@ def log(msg: str) -> None:
         except ImportError:
             log("[Warning] wandb not installed, skipping W&B logging")
 
-    # Build env_cfg_override from reward config
-    env_cfg_override: dict | None = None
-    if hasattr(cfg, "reward") and cfg.reward:
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override = {"reward_config": reward_dict}
+    from unilab.utils.reward_utils import extract_reward_config
+
+    env_cfg_override = extract_reward_config(cfg)
 
     preset = TASK_STEP_TUNING.get(task_name, {"threads": "32", "chunk": "16"})
     os.environ["UNILAB_MUJOCO_STEP_THREADS"] = preset["threads"]
@@ -398,7 +394,7 @@ def log(msg: str) -> None:
     if env.state is None:
         env.init_state()
     reset_indices = np.arange(env.num_envs, dtype=np.int32)
-    _, obs_dict, _ = env.reset(reset_indices)
+    obs_dict, _ = env.reset(reset_indices)
     obs = mx.array(flatten_obs_dict(obs_dict))
 
     obs_dim = sum(env.obs_groups_spec.values())

scripts/train_offpolicy.py

@@ -80,12 +80,9 @@ def resolve_checkpoint_path(
 
 def build_runner(algo_name: str, cfg: DictConfig):
     """Build algorithm runner from unified Hydra config."""
-    env_cfg_override = None
-    if hasattr(cfg, "reward") and cfg.reward:
-        from omegaconf import OmegaConf
+    from unilab.utils.reward_utils import extract_reward_config
 
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override = {"reward_config": reward_dict}
+    env_cfg_override = extract_reward_config(cfg)
 
     if algo_name == "sac":
         from unilab.algos.torch.fast_sac.learner import FastSACLearner
@@ -97,14 +94,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
+                cfg.training.task_name, num_envs=1, sim_backend=cfg.training.sim_backend,
+                env_cfg_override=env_cfg_override
             )
             assert env.action_space.shape
-            obs_dim = sum(env.obs_groups_spec.values())
+            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()
 
@@ -124,6 +125,7 @@ def build_runner(algo_name: str, cfg: DictConfig):
                 "use_layer_norm": cfg.algo.use_layer_norm,
                 "max_grad_norm": cfg.algo.algo_params.max_grad_norm,
                 "use_amp": cfg.training.use_amp,
+                "privileged_dim": privileged_dim,
                 # symmetry not supported in multi-GPU mode (mujoco_model not picklable)
                 "use_symmetry": False,
             }
@@ -229,11 +231,9 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> None:
     from unilab.utils import render_many
     from unilab.utils.algo_utils import build_actor
 
-    # Build env_cfg_override from reward config
-    env_cfg_override: dict | None = None
-    if hasattr(cfg, "reward") and cfg.reward:
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override = {"reward_config": reward_dict}
+    from unilab.utils.reward_utils import extract_reward_config
+
+    env_cfg_override = extract_reward_config(cfg)
 
     device = default_device(torch, cfg.training.device)
     print(f"Using device for play: {device}")

scripts/train_rsl_rl.py

@@ -70,7 +70,7 @@ def __init__(self, env, device="cuda"):
         self.reset()
 
     def _obs_to_tensordict(self, obs: dict[str, np.ndarray]) -> TensorDict:
-        actor = to_torch(obs["actor"], self.device)
+        actor = to_torch(obs["obs"], self.device)
         td = {"actor": actor}
         if "privileged" in obs:
             td["privileged"] = to_torch(obs["privileged"], self.device)
@@ -112,7 +112,7 @@ def reset(self):
         if self.env.state is None:
             self.env.init_state()
         env_indices = np.arange(self.num_envs, dtype=np.int32)
-        _, obs_out, _ = self.env.reset(env_indices)
+        obs_out, _ = self.env.reset(env_indices)
 
         self.episode_returns[:] = 0
         self.episode_lengths[:] = 0
@@ -131,11 +131,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: dict = {}
-    if hasattr(cfg, "reward") and cfg.reward:
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override["reward_config"] = reward_dict
+    env_cfg_override = extract_reward_config(cfg)
 
     env = registry.make(
         cfg.training.task_name,
@@ -265,12 +263,9 @@ def main(cfg: DictConfig) -> None:
     from omegaconf import OmegaConf
 
     from unilab.base import registry
+    from unilab.utils.reward_utils import extract_reward_config
 
-    # Build env_cfg_override from reward config
-    env_cfg_override = {}
-    if hasattr(cfg, "reward") and cfg.reward:
-        reward_dict = OmegaConf.to_container(cfg.reward, resolve=True)
-        env_cfg_override["reward_config"] = reward_dict
+    env_cfg_override = extract_reward_config(cfg)
 
     if torch.cuda.is_available():
         device = "cuda"