From a5e5d5029e3e05dcb15e797113b3830bd3a48855 Mon Sep 17 00:00:00 2001
From: 8x54zj-m <8x54zj-m@motrixlab.local>
Date: Tue, 30 Jun 2026 16:42:48 +0800
Subject: [PATCH] fix: dt in state.reward, feet_air_time first_contact timing

---
 .../motrix_envs/locomotion/go1/dreamwaq.py    | 45 +++++++++++++++----
 motrix_rl/src/motrix_rl/tasks/go1.py          |  5 ++-
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py
index a3a4447..999e28c 100644
--- a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py
+++ b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py
@@ -547,21 +547,53 @@ class DreamWaQTask(Go1WalkTask):
             # stand_still 在上游被注释掉
         }
 
-    # _reward_feet_air_time 使用基类实现（walk_np.py），
-    # 之前 override 读了不存在的 key 导致 reward 永远为 0
+    def update_feet_air_time(self, info: dict):
+        """覆盖基类 — 在清零前保存 first_contact 供奖励函数使用。"""
+        feet_air_time = info["feet_air_time"]
+        contacts = info["contacts"]
+        # 保存着地瞬间的状态（脚刚触地，且之前腾空时间 > 0）
+        info["first_contact"] = (feet_air_time > 0.0) & contacts
+        info["air_time_at_contact"] = feet_air_time.copy()
+        # 基类逻辑：累加 dt，着地清零
+        feet_air_time = feet_air_time + self.cfg.ctrl_dt
+        feet_air_time = feet_air_time * (~contacts)
+        info["feet_air_time"] = feet_air_time
+        return feet_air_time
+
+    def _reward_feet_air_time(self, commands, info):
+        """足部腾空时间奖励——与上游 legged_robot.py 完全对齐。
+
+        使用 update_feet_air_time 保存的 first_contact / air_time_at_contact。
+        """
+        first_contact = info.get("first_contact")
+        air_time = info.get("air_time_at_contact")
+        if first_contact is None or air_time is None:
+            return np.zeros(self._num_envs, dtype=np.float32)
+        rew = np.sum((air_time - 0.5) * first_contact, axis=1)
+        rew *= np.linalg.norm(commands[:, :2], axis=1) > 0.1
+        return rew
 
     def update_reward(self, state):
-        """存储各项奖励到 TensorBoard + 累计 episode 总和。"""
+        """存储各项奖励到 TensorBoard + 更新 state.reward。
+
+        与上游 _prepare_reward_function 对齐：所有 scales × dt。
+        """
         reward_dict = self._get_reward(state.data, state.info)
-        # 乘系数 + dt（与上游 _prepare_reward_function 对齐）
         scales = self._cfg.reward_config.scales
         dt = self._cfg.ctrl_dt
         scaled_terms = {
             k: v * scales.get(k, 0.0) * dt
             for k, v in reward_dict.items()
         }
+        # state.reward 直接用带 dt 的值（不调 super，super 不带 dt）
+        state = state.replace(reward=np.zeros(self._num_envs, dtype=np.float32))
+        for v in scaled_terms.values():
+            state.reward += v
+        if self._cfg.reward_config.only_positive_rewards:
+            state.reward = np.maximum(state.reward, 0.0)
+
+        # TensorBoard 日志
         state.info["reward_terms"] = {k: float(np.mean(v)) for k, v in scaled_terms.items()}
-        # 累计 episode 总和
         for k, v in scaled_terms.items():
             ek = f"ep_{k}"
             state.info[ek] = state.info.get(ek, np.zeros(self._num_envs, dtype=np.float32)) + v
@@ -579,9 +611,6 @@ class DreamWaQTask(Go1WalkTask):
             ep_report["terrain_level"] = float(np.mean(tl))
             state.info["ep_report"] = ep_report
 
-        state = super().update_reward(state)
-        if self._cfg.reward_config.only_positive_rewards:
-            state = state.replace(reward=np.maximum(state.reward, 0.0))
         return state
 
     # ── 额外奖励函数 ──
diff --git a/motrix_rl/src/motrix_rl/tasks/go1.py b/motrix_rl/src/motrix_rl/tasks/go1.py
index 55360df..ae4cd2e 100644
--- a/motrix_rl/src/motrix_rl/tasks/go1.py
+++ b/motrix_rl/src/motrix_rl/tasks/go1.py
@@ -105,7 +105,7 @@ class rslrl:
     class Go1DreamWaQWalkRslrlPpo(RslrlCfg):
         """Go1 DreamWaQ walk — CENet VAE + 不对称特权观测。"""
 
-        num_envs: int = 1024  # 上游 4096，CPU/GPU 安全默认
+        num_envs: int = 2048
 
         def __post_init__(self):
             runner = self.runner
@@ -126,6 +126,7 @@ class rslrl:
             runner.algorithm.entropy_coef = 0.01      # 上游 Go1RoughCfgPPO
             runner.algorithm.desired_kl = 0.01        # 上游 0.01 (默认 0.008)
             runner.algorithm.clip_param = 0.2
+            runner.algorithm.schedule = "fixed"   # 固定 schedule，防止 noise_std 发散
             runner.algorithm.gamma = 0.99
             runner.algorithm.lam = 0.95
             runner.algorithm.max_grad_norm = 1.0
@@ -135,7 +136,7 @@ class rslrl:
             runner.actor.class_name = (
                 "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel")
             runner.actor.hidden_dims = [512, 256, 128]
-            runner.actor.init_noise_std = 1.0  # 上游原值
+            runner.actor.init_noise_std = 0.5  # 降低探索噪声，1024 envs 不需要太高
 
             # Critic：标准 MLPModel，输入 privileged_obs
             runner.critic.class_name = "MLPModel"