fix: dt in state.reward, feet_air_time first_contact timing

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
 
     # ── 额外奖励函数 ──