Backup DreamWaQ rslrl stability fixes

motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py

@@ -121,18 +121,20 @@ class DreamWaQCfg(Go1WalkNpEnvCfg):
             "tracking_lin_vel": 1.5,
             "tracking_ang_vel": 1.0,
             "lin_vel_z": -2.0,
-            "ang_vel_xy": -0.05,
-            "orientation": -0.2,
+            "ang_vel_xy": -0.10,   # 抑制晃动
+            "orientation": -0.5,   # 强制平稳姿态
             "dof_acc": -2.5e-7,
-            "base_height": -1.0,
+            "base_height": -10.0,  # 强惩罚身高偏差，避免蹲伏或踮脚
             "feet_air_time": 0.1,
             "action_rate": -0.01,
             "joint_power": -2e-5,
-            "smoothness": -0.01,
+            "smoothness": -0.02,
             "power_distribution": -10e-6,
-            # 注意：stand_still 在上游被注释掉
+            "stand_still": -0.5,
+            "dof_pos_limits": -5.0,  # 关节限位软约束（参考 M20 修改版）
+            "collision": -1.0,       # 惩罚身体碰撞
         })
-        self.reward_config.only_positive_rewards = True
+        self.reward_config.only_positive_rewards = False  # 让坏行为负反馈直达策略
         self.reward_config.tracking_sigma = 0.25
 
 
@@ -181,12 +183,25 @@ class DreamWaQTask(Go1WalkTask):
     # ── 动作裁剪 + 力矩计算 ──
 
     def apply_action(self, actions, state):
-        """裁剪动作防止奖励计算溢出（上游 clip_actions=100）。"""
-        actions = np.clip(actions, -100.0, 100.0)
-        return super().apply_action(actions, state)
+        """裁剪动作 + 随机延迟（模拟真实部署延迟，提高 sim-to-real 泛化）。"""
+        actions = np.clip(actions, -4.0, 4.0)
+        # 随机动作延迟：0-3 个控制步（0-60ms）
+        if not hasattr(self, '_action_buffer'):
+            self._action_buffer = np.zeros((self._num_envs, 3, self._num_action), dtype=np.float32)
+            self._latency_steps = np.random.randint(0, 4, size=self._num_envs)
+        # 滑动缓冲区，新动作插入末尾
+        self._action_buffer = np.concatenate([
+            self._action_buffer[:, 1:, :], actions[:, np.newaxis, :]
+        ], axis=1)
+        # 根据每个 env 的延迟取对应历史动作
+        delayed_actions = np.zeros_like(actions)
+        for i in range(self._num_envs):
+            lat = self._latency_steps[i]
+            delayed_actions[i] = self._action_buffer[i, 2 - lat]  # buffer[-1] 是最新
+        return super().apply_action(delayed_actions, state)
 
     def _compute_torques(self, actions, data):
-        """PD 控制器 + 域随机化。力矩裁剪防溢出。"""
+        """PD 控制器 + 域随机化 + 关节限位。力矩裁剪防溢出。"""
         state = getattr(self, '_state', None)
         if state is not None:
             motor_strength = state.info.get("motor_strength",
@@ -201,8 +216,13 @@ class DreamWaQTask(Go1WalkTask):
             kd_factor = np.ones(1, dtype=np.float32)
 
         actions_scaled = actions * self.cfg.control_config.action_scale * motor_strength[:, np.newaxis]
+        # 目标角度限制在关节范围内（与 MuJoCo 硬限位对齐）
+        target = actions_scaled + self.default_angles
+        lo = self._model.joint_limits[0][np.newaxis, :]  # (1, 12)
+        hi = self._model.joint_limits[1][np.newaxis, :]
+        target = np.clip(target, lo, hi)
         torques = (self.kps * kp_factor[:, np.newaxis]) * (
-            actions_scaled + self.default_angles - self.get_dof_pos(data)
+            target - self.get_dof_pos(data)
         ) - (self.kds * kd_factor[:, np.newaxis]) * self.get_dof_vel(data)
         return np.clip(torques, -80.0, 80.0)
 
@@ -326,14 +346,31 @@ class DreamWaQTask(Go1WalkTask):
 
     def update_observation(self, state):
         data = state.data
-        # 清理物理崩溃残留的 NaN/Inf（仅在检测到时才修改，避免无谓开销）
+        # 清理物理崩溃残留的 NaN/Inf
         dv = data.dof_vel
         if np.any(~np.isfinite(dv)):
-            data.set_dof_vel(np.nan_to_num(np.array(dv), nan=0.0, posinf=0.0, neginf=0.0))
+            dv_clean = np.array(dv)
+            dv_clean[:, 6:] = np.nan_to_num(dv_clean[:, 6:], nan=0.0, posinf=0.0, neginf=0.0)
+            data.set_dof_vel(dv_clean)
         dp = data.dof_pos
         if np.any(~np.isfinite(dp)):
-            data.set_dof_pos(np.nan_to_num(np.array(dp), nan=0.0, posinf=0.0, neginf=0.0), self._model)
+            dp_clean = np.array(dp)
+            # 四元数 NaN → 整行替换为单位四元数 [0,0,0,1]
+            quat_nan_row = np.any(~np.isfinite(dp_clean[:, 3:7]), axis=1)
+            if np.any(quat_nan_row):
+                dp_clean[quat_nan_row, 3:7] = [0.0, 0.0, 0.0, 1.0]
+            # 关节位置 NaN → 0
+            dp_clean[:, 7:] = np.nan_to_num(dp_clean[:, 7:], nan=0.0, posinf=0.0, neginf=0.0)
+            data.set_dof_pos(dp_clean, self._model)
+        else:
+            quat_norm = np.linalg.norm(dp[:, 3:7], axis=1)
+            bad_quat = quat_norm < 1e-6
+            if np.any(bad_quat):
+                dp_clean = np.array(dp)
+                dp_clean[bad_quat, 3:7] = [0.0, 0.0, 0.0, 1.0]
+                data.set_dof_pos(dp_clean, self._model)
         obs = self._get_obs(data, state.info)
+        obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)
 
         # 更新历史缓冲区：CENet 用不含当前帧的历史预测下一帧（论文 LVAE = MSE(õ_{t+1}, o_{t+1})）
         full_hist = state.info.get("obs_history_full",
@@ -360,8 +397,10 @@ class DreamWaQTask(Go1WalkTask):
         if do_rand.any():
             n_rand = int(do_rand.sum())
             state.info["motor_strength"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
-            state.info["kp_factor"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
-            state.info["kd_factor"][do_rand] = np.random.uniform(0.9, 1.1, size=n_rand).astype(np.float32)
+            state.info["kp_factor"][do_rand] = np.random.uniform(0.85, 1.15, size=n_rand).astype(np.float32)
+            state.info["kd_factor"][do_rand] = np.random.uniform(0.85, 1.15, size=n_rand).astype(np.float32)
+            # 动作延迟随机化（每个 rand 周期重新分配延迟步数）
+            self._latency_steps[do_rand] = np.random.randint(0, 4, size=n_rand)
             state.info["friction_coeff"][do_rand] = np.random.uniform(0.2, 1.25, size=n_rand).astype(np.float32)
             state.info["added_mass"][do_rand] = np.random.uniform(-1.0, 2.0, size=n_rand).astype(np.float32)
             state.info["com_displacement"][do_rand] = np.random.uniform(-0.05, 0.05, size=(n_rand, 3)).astype(np.float32)
@@ -369,6 +408,8 @@ class DreamWaQTask(Go1WalkTask):
 
         # 特权观测
         state.info["privileged_obs"] = self._get_privileged_obs(data, obs)
+        state.info["privileged_obs"] = np.nan_to_num(
+            state.info["privileged_obs"], nan=0.0, posinf=0.0, neginf=0.0)
         state.info["base_vel"] = self.get_local_linvel(data)
 
         # 足部接触
@@ -429,7 +470,9 @@ class DreamWaQTask(Go1WalkTask):
                     cy = half_y - self._border - row * self._cell_size - self._cell_size / 2
                     all_origins[row, col] = [cx, cy]
             self._terrain_origins = all_origins
-            self._max_init_level = 0  # 上游原值
+            # 跟踪课程进度：动态 max_init_level，避免将已学会的机器人拉回平地
+            self._max_init_level = 0
+            self._level_history = []
 
         if num_reset > 0 and self._init_done and state is not None and hasattr(state, 'info'):
             old_info = state.info
@@ -439,16 +482,117 @@ class DreamWaQTask(Go1WalkTask):
             base_pose = self._body.get_pose(state.data)
             base_pos = base_pose[done, :2]
             distance = np.linalg.norm(base_pos - old_origins, axis=1)
-            move_up = distance > 3.0
-            cmd_speed = np.linalg.norm(old_commands[:, :2], axis=1)
-            required_dist = cmd_speed * (self.cfg.max_episode_steps * self.cfg.ctrl_dt) * 0.5
-            move_down = (distance < required_dist) & ~move_up
+            # ── 多条件 Curriculum 升级 ──
+            # 1) 速度跟踪: exp(-L2²/σ), σ=0.25, >0.5 表示等效恒定 L2 误差约 0.42m/s
+            ep_tracking = old_info.get("ep_tracking_sum", np.zeros(self._num_envs, dtype=np.float32))[done]
+            ep_steps = old_info.get("ep_steps", np.zeros(self._num_envs, dtype=np.int32))[done]
+            mask = ep_steps > 0
+            avg_tracking = np.zeros(num_reset, dtype=np.float32)
+            avg_tracking[mask] = ep_tracking[mask] / ep_steps[mask].astype(np.float32)
+            # 调试：打印 avg_tracking 分布（单次）
+            if not hasattr(self, '_dbg_tracking_printed'):
+                self._dbg_tracking_printed = True
+                valid = avg_tracking[mask]
+                if len(valid) > 0:
+                    print(f"[Curric-DBG] avg_tracking: mean={np.mean(valid):.4f} "
+                          f"min={np.min(valid):.4f} p10={np.percentile(valid,10):.4f} "
+                          f"p50={np.percentile(valid,50):.4f} p90={np.percentile(valid,90):.4f} "
+                          f"max={np.max(valid):.4f} >0.5={100*np.mean(valid>0.5):.0f}%")
+            # 2) 存活率: 必须自然结束(truncated), 而非摔倒终止(terminated)
+            term = state.terminated[done] if hasattr(state, 'terminated') else np.zeros(num_reset, dtype=bool)
+            not_fallen = ~term  # True=活到 episode 结束
+            # 3) 姿态: avg_orient 保留 dt，-0.0005 对应等效倾角约 13 deg
+            ep_orient = old_info.get("ep_orientation", np.zeros(self._num_envs, dtype=np.float32))[done]
+            avg_orient = np.zeros(num_reset, dtype=np.float32)
+            avg_orient[mask] = ep_orient[mask] / ep_steps[mask].astype(np.float32)
+            # 4) 身高: avg_base_height 保留 dt，-0.001 对应 RMS 高度误差约 7cm
+            ep_bh = old_info.get("ep_base_height", np.zeros(self._num_envs, dtype=np.float32))[done]
+            avg_bh = np.zeros(num_reset, dtype=np.float32)
+            avg_bh[mask] = ep_bh[mask] / ep_steps[mask].astype(np.float32)
+            # orient: -0.5*dt*sin²(theta); bh: -10*dt*z_err²
+            move_up_raw = (
+                (avg_tracking > 0.5)
+                & not_fallen
+                & (avg_orient > -0.0005)
+                & (avg_bh > -0.001)
+            )
+            # ── 渐进升级：连续通过 3 次 + 冷却期 + 每次只升 1 级 ──
+            if not hasattr(self, '_consecutive_pass_count'):
+                self._consecutive_pass_count = np.zeros(self._num_envs, dtype=np.int32)
+                self._pending_upgrade = np.zeros(self._num_envs, dtype=bool)
+                self._upgrade_cooldown = np.zeros(self._num_envs, dtype=np.int32)
+
+            # 每 env: 连续通过 +1, 失败重置（仅更新被 reset 的 env）
+            self._consecutive_pass_count[done_idx] = np.where(
+                move_up_raw,
+                self._consecutive_pass_count[done_idx] + 1, 0)
+            # 冷却期以该 env 自己完成的 episode 为单位，而不是全局 reset 批次数。
+            self._upgrade_cooldown[done_idx] = np.maximum(0, self._upgrade_cooldown[done_idx] - 1)
+            streak_ready = self._consecutive_pass_count[done_idx] >= 3
+            cooldown_blocked = streak_ready & (self._upgrade_cooldown[done_idx] > 0)
+            # 连续 3 次通过 且 不在冷却期 → 升级（取被 reset env 的当前状态）
+            can_upgrade = streak_ready & (self._upgrade_cooldown[done_idx] == 0)
+            # 升级后设置 5 episode 冷却期，并从新等级重新累计连续通过次数。
+            self._upgrade_cooldown[done_idx] = np.where(
+                can_upgrade, 5, self._upgrade_cooldown[done_idx])
+            self._consecutive_pass_count[done_idx] = np.where(
+                can_upgrade, 0, self._consecutive_pass_count[done_idx])
+
+            move_up = can_upgrade
+            move_down_raw = (
+                (~not_fallen)
+                | (avg_tracking < 0.35)
+                | (avg_orient <= -0.001)
+                | (avg_bh <= -0.002)
+            )
+            move_down = move_down_raw & ~move_up
             new_levels = np.where(move_up, old_levels + 1, old_levels)
             new_levels = np.where(move_down, new_levels - 1, new_levels)
-            at_max = new_levels >= self._num_rows
-            if at_max.any():
-                new_levels[at_max] = np.random.randint(0, self._num_rows, size=int(at_max.sum()))
             new_levels = np.clip(new_levels, 0, self._num_rows - 1)
+            # 跟踪 level 分布，动态更新 max_init_level（P90 - 2）
+            self._level_history.extend(new_levels.tolist())
+            if len(self._level_history) > 500:
+                self._level_history = self._level_history[-500:]
+            if len(self._level_history) >= 10:
+                sorted_lv = sorted(self._level_history)
+                self._max_init_level = max(0, sorted_lv[int(len(sorted_lv) * 0.9)] - 2)
+            # ── 调试日志: curriculum 各条件通过率 ──
+            if not hasattr(self, '_curric_log_counter'):
+                self._curric_log_counter = 0
+                self._curric_track_ok = 0; self._curric_survive_ok = 0
+                self._curric_orient_ok = 0; self._curric_bh_ok = 0
+                self._curric_raw_all_ok = 0; self._curric_streak_ok = 0
+                self._curric_upgrade_ok = 0; self._curric_cooldown_block = 0
+                self._curric_down_ok = 0; self._curric_down_track = 0
+                self._curric_down_stability = 0; self._curric_total = 0
+            self._curric_log_counter += 1
+            self._curric_track_ok += int((avg_tracking > 0.5).sum())
+            self._curric_survive_ok += int(not_fallen.sum())
+            self._curric_orient_ok += int((avg_orient > -0.0005).sum())
+            self._curric_bh_ok += int((avg_bh > -0.001).sum())
+            self._curric_raw_all_ok += int(move_up_raw.sum())
+            self._curric_streak_ok += int(streak_ready.sum())
+            self._curric_upgrade_ok += int(move_up.sum())
+            self._curric_cooldown_block += int(cooldown_blocked.sum())
+            self._curric_down_ok += int(move_down.sum())
+            self._curric_down_track += int((avg_tracking < 0.35).sum())
+            self._curric_down_stability += int((~not_fallen | (avg_orient <= -0.001) | (avg_bh <= -0.002)).sum())
+            self._curric_total += num_reset
+            if self._curric_log_counter % 50 == 0:
+                def pct(n): return 100*n/max(self._curric_total,1)
+                print(f"[Curric] n={self._curric_total:5d} "
+                      f"track={pct(self._curric_track_ok):.0f}% "
+                      f"survive={pct(self._curric_survive_ok):.0f}% "
+                      f"orient={pct(self._curric_orient_ok):.0f}% "
+                      f"bh={pct(self._curric_bh_ok):.0f}% "
+                      f"rawALL={pct(self._curric_raw_all_ok):.0f}% "
+                      f"streak3={pct(self._curric_streak_ok):.0f}% "
+                      f"up={pct(self._curric_upgrade_ok):.0f}% "
+                      f"cool={pct(self._curric_cooldown_block):.0f}% "
+                      f"down={pct(self._curric_down_ok):.0f}% "
+                      f"downT={pct(self._curric_down_track):.0f}% "
+                      f"downS={pct(self._curric_down_stability):.0f}% "
+                      f"max_init={self._max_init_level}")
         else:
             new_levels = np.random.randint(0, self._max_init_level + 1, size=num_reset, dtype=np.int32)
         self._init_done = True
@@ -545,12 +689,32 @@ class DreamWaQTask(Go1WalkTask):
 
         return obs, info
 
+    # ── 新增奖励（参考 M20 修改版）──
+
+    def _reward_dof_pos_limits(self, data):
+        """关节软限位惩罚——接近关节极限时扣分。"""
+        lo = self._model.joint_limits[0]
+        hi = self._model.joint_limits[1]
+        margin = 0.1  # rad 软边界
+        pos = self.get_dof_pos(data)
+        lower = np.clip(lo + margin - pos, 0, None)
+        upper = np.clip(pos - (hi - margin), 0, None)
+        return np.sum(lower + upper, axis=1)
+
+    def _reward_collision(self, data):
+        """身体碰撞惩罚——膝/肩触地扣分。"""
+        cquerys = self._model.get_contact_query(data)
+        penal = cquerys.is_colliding(self.termination_check)
+        return np.any(penal.reshape(self._num_envs, -1), axis=1).astype(np.float32)
+
     def update_terminated(self, state):
         """基类接触终止 + hfield 边界终止（防止走出地形导致 NaN）。"""
         state = super().update_terminated(state)
+        if self._model.num_hfields == 0:
+            return state  # flat 场景无 hfield
+        hf = self._model.get_hfield(0)
         pose = self._body.get_pose(state.data)
         base_xy = pose[:, :2]
-        hf = self._model.get_hfield(0)
         b = hf.bound
         out_of_bounds = (
             (base_xy[:, 0] < b[0]) | (base_xy[:, 0] > b[3]) |
@@ -577,7 +741,9 @@ class DreamWaQTask(Go1WalkTask):
             "joint_power": self._reward_joint_power(data),
             "smoothness": self._reward_smoothness(info),
             "power_distribution": self._reward_power_distribution(data),
-            # stand_still 在上游被注释掉
+            "stand_still": self._reward_stand_still(data, commands),
+            "dof_pos_limits": self._reward_dof_pos_limits(data),
+            "collision": self._reward_collision(data),
         }
 
     def update_feet_air_time(self, info: dict):
@@ -606,6 +772,13 @@ class DreamWaQTask(Go1WalkTask):
         rew *= np.linalg.norm(commands[:, :2], axis=1) > 0.1
         return rew
 
+    def _reward_stand_still(self, data, commands):
+        """惩罚有命令但速度接近零的'卡住'行为。"""
+        cmd_norm = np.linalg.norm(commands[:, :2], axis=1)
+        vel_norm = np.linalg.norm(self.get_local_linvel(data)[:, :2], axis=1)
+        stifled = (cmd_norm > 0.2) & (vel_norm < 0.1)  # 有命令但基本不动
+        return stifled.astype(np.float32)
+
     def update_reward(self, state):
         """存储各项奖励到 TensorBoard + 更新 state.reward。
 

motrix_envs/src/motrix_envs/locomotion/go1/xmls/go1_motor_actuator.xml

@@ -49,19 +49,10 @@
           <geom size="0.046 0.02" quat="1 1 0 0" type="cylinder"/>
         </default>
         <default class="thigh1">
-          <geom size="0.015" fromto="-0.02 0 0 -0.02 0 -0.16"/>
-        </default>
-        <default class="thigh2">
-          <geom size="0.015" fromto="0 0 0 -0.02 0 -0.1"/>
-        </default>
-        <default class="thigh3">
-          <geom size="0.015" fromto="-0.02 0 -0.16 0 0 -0.2"/>
+          <geom type="capsule" size="0.02" fromto="0 0 0 0 0 -0.213"/>
         </default>
         <default class="calf1">
-          <geom size="0.01" fromto="0 0 0 0.02 0 -0.13"/>
-        </default>
-        <default class="calf2">
-          <geom size="0.01" fromto="0.02 0 -0.13 0 0 -0.2"/>
+          <geom type="capsule" size="0.015" fromto="0 0 0 0 0 -0.213"/>
         </default>
         <default class="foot">
           <geom type="sphere" size="0.023" pos="0 0 -0.213" solimp="0.9 .99 0.001" priority="10" condim="3"/>
@@ -97,20 +88,20 @@
         <inertial pos="-0.0049166 0.00762615 -8.865e-05" quat="0.507341 0.514169 0.495027 0.482891" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
         <joint class="abduction" name="FR_hip_joint"/>
         <geom class="visual" mesh="hip" quat="1 0 0 0"/>
-        <!-- <geom name="fr_hip" class="hip_right1"/> -->
+        <geom name="fr_hip" class="hip_right1"/> 
         <body name="FR_thigh" pos="0 -0.08 0">
           <inertial pos="-0.00304722 0.019315 -0.0305004" quat="0.65243 -0.0272313 0.0775126 0.753383" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
           <joint class="hip" name="FR_thigh_joint"/>
           <geom class="visual" mesh="thigh_mirror"/>
-          <!-- <geom name="fr_thigh1" class="thigh1"/> -->
-          <!-- <geom name="fr_thigh2" class="thigh2"/> -->
-          <!-- <geom name="fr_thigh3" class="thigh3"/> -->
+          <geom name="fr_thigh1" class="thigh1"/> 
+          <!-- removed: <geom name="fr_thigh2" class="thigh2"/> (simplified to single capsule) --> 
+          <!-- removed: <geom name="fr_thigh3" class="thigh3"/> (simplified to single capsule) --> 
           <body name="FR_calf" pos="0 0 -0.213">
             <inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
             <joint class="knee" name="FR_calf_joint"/>
             <geom class="visual" mesh="calf"/>
-            <!-- <geom name="fr_calf1" class="calf1"/> -->
-            <!-- <geom name="fr_calf2" class="calf2"/> -->
+            <geom name="fr_calf1" class="calf1"/> 
+            <!-- removed: <geom name="fr_calf2" class="calf2"/> (simplified to single capsule) --> 
             <geom name="FR_foot" class="foot"/>
             <site name="FR" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
           </body>
@@ -120,20 +111,20 @@
         <inertial pos="-0.0049166 -0.00762615 -8.865e-05" quat="0.482891 0.495027 0.514169 0.507341" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
         <joint class="abduction" name="FL_hip_joint"/>
         <geom class="visual" mesh="hip"/>
-        <!-- <geom name="fl_hip" class="hip_left1"/> -->
+        <geom name="fl_hip" class="hip_left1"/> 
         <body name="FL_thigh" pos="0 0.08 0">
           <inertial pos="-0.00304722 -0.019315 -0.0305004" quat="0.753383 0.0775126 -0.0272313 0.65243" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
           <joint class="hip" name="FL_thigh_joint"/>
           <geom class="visual" mesh="thigh"/>
-          <!-- <geom name="fl_thigh1" class="thigh1"/> -->
-          <!-- <geom name="fl_thigh2" class="thigh2"/> -->
-          <!-- <geom name="fl_thigh3" class="thigh3"/> -->
+          <geom name="fl_thigh1" class="thigh1"/> 
+          <!-- removed: <geom name="fl_thigh2" class="thigh2"/> (simplified to single capsule) --> 
+          <!-- removed: <geom name="fl_thigh3" class="thigh3"/> (simplified to single capsule) --> 
           <body name="FL_calf" pos="0 0 -0.213">
             <inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
             <joint class="knee" name="FL_calf_joint"/>
             <geom class="visual" mesh="calf"/>
-            <!-- <geom name="fl_calf1" class="calf1"/> -->
-            <!-- <geom name="fl_calf2" class="calf2"/> -->
+            <geom name="fl_calf1" class="calf1"/> 
+            <!-- removed: <geom name="fl_calf2" class="calf2"/> (simplified to single capsule) --> 
             <geom name="FL_foot" class="foot"/>
             <site name="FL" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
           </body>
@@ -143,20 +134,20 @@
         <inertial pos="0.0049166 0.00762615 -8.865e-05" quat="0.495027 0.482891 0.507341 0.514169" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
         <joint class="abduction" name="RR_hip_joint"/>
         <geom class="visual" quat="0 0 0 -1" mesh="hip"/>
-        <!-- <geom name="rr_hip" class="hip_right1"/> -->
+        <geom name="rr_hip" class="hip_right1"/> 
         <body name="RR_thigh" pos="0 -0.08 0">
           <inertial pos="-0.00304722 0.019315 -0.0305004" quat="0.65243 -0.0272313 0.0775126 0.753383" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
           <joint class="hip" name="RR_thigh_joint"/>
           <geom class="visual" mesh="thigh_mirror"/>
-          <!-- <geom name="rr_thigh1" class="thigh1"/>
-          <geom name="rr_thigh2" class="thigh2"/>
-          <geom name="rr_thigh3" class="thigh3"/> -->
+          <geom name="rr_thigh1" class="thigh1"/>
+          <!-- removed: <geom name="rr_thigh2" class="thigh2"/> (simplified to single capsule) -->
+          <!-- removed: <geom name="rr_thigh3" class="thigh3"/> (simplified to single capsule) --> 
           <body name="RR_calf" pos="0 0 -0.213">
             <inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
             <joint class="knee" name="RR_calf_joint"/>
             <geom class="visual" mesh="calf"/>
-            <!-- <geom name="rr_calf1" class="calf1"/>
-            <geom name="rr_calf2" class="calf2"/> -->
+            <geom name="rr_calf1" class="calf1"/>
+            <!-- removed: <geom name="rr_calf2" class="calf2"/> (simplified to single capsule) --> 
             <geom name="RR_foot" class="foot"/>
             <site name="RR" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
           </body>
@@ -166,20 +157,20 @@
         <inertial pos="0.0049166 -0.00762615 -8.865e-05" quat="0.514169 0.507341 0.482891 0.495027" mass="0.68" diaginertia="0.000734064 0.000468438 0.000398719"/>
         <joint class="abduction" name="RL_hip_joint"/>
         <geom class="visual" quat="0 0 1 0" mesh="hip"/>
-        <!-- <geom name="rl_hip" class="hip_left1"/> -->
+        <geom name="rl_hip" class="hip_left1"/> 
         <body name="RL_thigh" pos="0 0.08 0">
           <inertial pos="-0.00304722 -0.019315 -0.0305004" quat="0.753383 0.0775126 -0.0272313 0.65243" mass="1.009" diaginertia="0.00478717 0.00460903 0.000709268"/>
           <joint class="hip" name="RL_thigh_joint"/>
           <geom class="visual" mesh="thigh"/>
-          <!-- <geom name="rl_thigh1" class="thigh1"/>
-          <geom name="rl_thigh2" class="thigh2"/>
-          <geom name="rl_thigh3" class="thigh3"/> -->
+          <geom name="rl_thigh1" class="thigh1"/>
+          <!-- removed: <geom name="rl_thigh2" class="thigh2"/> (simplified to single capsule) -->
+          <!-- removed: <geom name="rl_thigh3" class="thigh3"/> (simplified to single capsule) --> 
           <body name="RL_calf" pos="0 0 -0.213">
             <inertial pos="0.00429862 0.000976676 -0.146197" quat="0.691246 0.00357467 0.00511118 0.722592" mass="0.195862" diaginertia="0.00149767 0.00148468 3.58427e-05"/>
             <joint class="knee" name="RL_calf_joint"/>
             <geom class="visual" mesh="calf"/>
-            <!-- <geom name="rl_calf1" class="calf1"/>
-            <geom name="rl_calf2" class="calf2"/> -->
+            <geom name="rl_calf1" class="calf1"/>
+            <!-- removed: <geom name="rl_calf2" class="calf2"/> (simplified to single capsule) --> 
             <geom name="RL_foot" class="foot"/>
             <site name="RL" pos="0 0 -0.213" type="sphere" size="0.023" group="5"/>
           </body>

motrix_envs/src/motrix_envs/np/env.py

@@ -182,17 +182,19 @@ class NpEnv(ABEnv):
         pass
 
     def physics_step(self):
+        self._physics_crashed_this_step = False
         # motrixsim.SceneModel.step only supports single step, so we loop
         try:
             for _ in range(self._cfg.sim_substeps):
                 self._model.step(self._state.data)
-        except Exception as e:
+        except BaseException as e:
             # Rust panic / 物理崩溃 → 标记终止 + 清除腐蚀数据
             n = self._state.data.shape[0]
             self._state.terminated[:] = True
             self._state.reward[:] = 0.0
             # 强制重置物理状态，防止 NaN 传播到 observations/rewards
             self._state.data.reset(self._model)
+            self._physics_crashed_this_step = True
             if not hasattr(self, '_physics_crash_count'):
                 self._physics_crash_count = 0
             self._physics_crash_count += 1
@@ -213,6 +215,9 @@ class NpEnv(ABEnv):
         self._state = self.apply_action(actions, self._state)
         assert self._state is not None, "apply_action must return a valid NpEnvState"
         self.physics_step()
+        if getattr(self, "_physics_crashed_this_step", False):
+            self._reset_done_envs()
+            return self._state
         self._state = self.update_state(self._state)
         self._state.info["steps"] += 1
         self._update_truncate()

motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py

@@ -4,6 +4,43 @@ import torch
 import torch.nn as nn
 from torch.distributions import Normal
 
+
+class RunningStats(nn.Module):
+    """Running mean/variance tracker for actor observation normalization.
+
+    Matches SKRL RunningStandardScaler behavior. Without it, actor sees
+    unscaled inputs at inference → garbage actions → robot collapses.
+    """
+    def __init__(self, num_features: int, eps: float = 1e-8, clip: float = 5.0):
+        super().__init__()
+        self.eps = eps
+        self.clip = clip
+        self.register_buffer("count", torch.zeros(1, dtype=torch.int64))
+        self.register_buffer("mean", torch.zeros(num_features))
+        self.register_buffer("var", torch.ones(num_features))
+
+    @torch.no_grad()
+    def update(self, x):
+        if x.dim() != 2:
+            return
+        n = x.shape[0]
+        batch_mean = x.mean(dim=0)
+        batch_var = x.var(dim=0, unbiased=False)
+        batch_count = torch.tensor(n, dtype=torch.int64, device=x.device)
+        delta = batch_mean - self.mean
+        total = self.count + batch_count
+        self.mean.add_(delta * batch_count.float() / total.float())
+        m_a = self.var * self.count.float()
+        m_b = batch_var * batch_count.float()
+        m2 = m_a + m_b + delta.pow(2) * self.count.float() * batch_count.float() / total.float()
+        self.var.copy_(m2 / total.float())
+        self.count.add_(batch_count)
+
+    def forward(self, x):
+        return torch.clamp(
+            (x - self.mean) / (self.var.sqrt() + self.eps), -self.clip, self.clip)
+
+
 class ActorCritic_DWAQ(nn.Module):
     def __init__(self, num_actor_obs, num_critic_obs, num_actions, cenet_in_dim, cenet_out_dim, activation="elu", init_noise_std=1.0,):
         super().__init__()
@@ -52,6 +89,7 @@ class ActorCritic_DWAQ(nn.Module):
         )
 
         self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+        self.actor_normalizer = RunningStats(actor_input_dim, eps=1e-8, clip=5.0)
         self.distribution = None
         # disable args validation for speedup
         Normal.set_default_validate_args = False
@@ -89,7 +127,7 @@ class ActorCritic_DWAQ(nn.Module):
         # code = mean_latent + var*code_temp
         # print("latent : ",code[0])
         mean_vel = self.encode_mean_vel(distribution)
-        logvar_vel = self.encode_mean_vel(distribution)
+        logvar_vel = self.encode_logvar_vel(distribution)  # FIXED: was encode_mean_vel
         code_latent = self.reparameterise(mean_latent,logvar_latent)
         code_vel = self.reparameterise(mean_vel,logvar_vel)
         code = torch.cat((code_vel,code_latent),dim=-1)
@@ -117,17 +155,23 @@ class ActorCritic_DWAQ(nn.Module):
         self.distribution = Normal(mean, mean * 0.0 + self.std)
 
     def act(self, observations, obs_history, **kwargs):
-        code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
-        observations = torch.cat((code,observations),dim=-1)
+        code, _, decode, _, _, _, _ = self.cenet_forward(obs_history)
+        observations = torch.cat((code, observations), dim=-1)
+        if self.training:
+            self.actor_normalizer.update(observations)
+        if self.actor_normalizer.count > 10:
+            observations = self.actor_normalizer(observations)
         self.update_distribution(observations)
         return self.distribution.sample()
 
     def get_actions_log_prob(self, actions):
         return self.distribution.log_prob(actions).sum(dim=-1)
 
-    def act_inference(self, observations,obs_history):
-        code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
-        observations = torch.cat((code,observations),dim=-1)
+    def act_inference(self, observations, obs_history):
+        code, _, decode, _, _, _, _ = self.cenet_forward(obs_history)
+        observations = torch.cat((code, observations), dim=-1)
+        if self.actor_normalizer.count > 10:
+            observations = self.actor_normalizer(observations)
         actions_mean = self.actor(observations)
         return actions_mean
 

motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py

@@ -20,7 +20,7 @@ import torch
 class DwaqVecEnv:
     """Wraps a MotrixLab DreamWaQ NpEnv for the upstream DreamWaQ rsl_rl runner."""
 
-    def __init__(self, env, device, num_obs=45, num_privileged_obs=235,
+    def __init__(self, env, device, num_obs=45, num_privileged_obs=247,
                  num_obs_hist=5, num_actions=12, clip_actions=100.0):
         self._env = env
         self.device = device

motrix_rl/src/motrix_rl/dwaq_rsl/runner.py

@@ -158,12 +158,19 @@ class OnPolicyRunner:
                   f"rew={rew:.2f} eplen={elen:.0f}")
 
     def save(self, path, infos=None):
-        torch.save({
+        save_dict = {
             "model_state_dict": self.alg.actor_critic.state_dict(),
             "optimizer_state_dict": self.alg.optimizer.state_dict(),
             "iter": self.current_learning_iteration,
             "infos": infos,
-        }, path)
+        }
+        # Persist RunningStats normalizer state for inference
+        if hasattr(self.alg.actor_critic, 'actor_normalizer'):
+            n = self.alg.actor_critic.actor_normalizer
+            save_dict["normalizer_count"] = n.count
+            save_dict["normalizer_mean"] = n.mean
+            save_dict["normalizer_var"] = n.var
+        torch.save(save_dict, path)
 
     def load(self, path, load_optimizer=True):
         loaded_dict = torch.load(path, map_location=self.device)
@@ -171,6 +178,12 @@ class OnPolicyRunner:
         if load_optimizer:
             self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
         self.current_learning_iteration = loaded_dict["iter"]
+        # Restore RunningStats normalizer state
+        if "normalizer_count" in loaded_dict and hasattr(self.alg.actor_critic, 'actor_normalizer'):
+            n = self.alg.actor_critic.actor_normalizer
+            n.count.copy_(loaded_dict["normalizer_count"])
+            n.mean.copy_(loaded_dict["normalizer_mean"])
+            n.var.copy_(loaded_dict["normalizer_var"])
         return loaded_dict["infos"]
 
     def get_inference_policy(self, device=None):

motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py

@@ -14,6 +14,7 @@
 
 import torch
 import torch.nn as nn
+from torch.distributions import Normal
 from tensordict import TensorDict
 
 from rsl_rl.models.mlp_model import MLPModel
@@ -113,23 +114,43 @@ class CENetActorModel(MLPModel):
 
         # 在 nn.Module.__init__ 之后创建 VAE 子模块
         self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation)
+        self.action_clip = 4.0
+        self.std_clip = 0.6
         self._last_cenet_output = None
         # AdaBoot: 自适应 bootstrapping（论文 Section II-C）
         self._adaboot_cv_buffer = []  # 速度估计误差的 CV 历史
         self._adaboot_prob = 1.0      # 当前 bootstrap 概率（1.0 = 完全信任 GT）
 
     def _update_distribution(self, obs: torch.Tensor) -> None:
-        """覆盖父类 — 强制 std > 0 再创建 Normal 分布（防止 NaN）。"""
-        # 先 clamp std，再调父类创建分布
+        """覆盖父类 — 限制动作分布，避免采样动作进入 PD/关节限位饱和区。"""
+        mean = torch.clamp(self.mlp(obs), -self.action_clip, self.action_clip)
         if self.stochastic and not self.state_dependent_std:
             with torch.no_grad():
                 if self.noise_std_type == "scalar":
                     self.std.nan_to_num_(nan=0.5, posinf=1.0, neginf=1.0)
-                    self.std.clamp_(min=1e-6)
+                    self.std.clamp_(min=1e-6, max=self.std_clip)
                 elif self.noise_std_type == "log":
                     self.log_std.nan_to_num_(nan=0.0, posinf=5.0, neginf=-5.0)
-                    self.log_std.clamp_(min=-20.0, max=10.0)
-        super()._update_distribution(obs)
+                    self.log_std.clamp_(min=-20.0, max=torch.log(torch.tensor(self.std_clip)).item())
+            if self.noise_std_type == "scalar":
+                std = self.std.expand_as(mean)
+            elif self.noise_std_type == "log":
+                std = torch.exp(self.log_std).expand_as(mean)
+            else:
+                raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}.")
+        else:
+            std = torch.full_like(mean, self.std_clip)
+        self.distribution = Normal(mean, std)
+
+    def forward(
+        self,
+        obs: TensorDict,
+        masks: torch.Tensor | None = None,
+        hidden_state: HiddenState = None,
+        stochastic_output: bool = False,
+    ) -> torch.Tensor:
+        actions = super().forward(obs, masks, hidden_state, stochastic_output)
+        return torch.clamp(actions, -self.action_clip, self.action_clip)
 
     def _get_latent_dim(self) -> int:
         """Actor 实际输入：code(19) + policy(45) = 64。
@@ -164,10 +185,10 @@ class CENetActorModel(MLPModel):
             code = torch.cat([code_vel, code[:, 3:]], dim=-1)
 
         # 防止 VAE NaN 传播到下游
-        if torch.isnan(code).any():
-            code = torch.nan_to_num(code, nan=0.0)
-        if torch.isnan(policy_obs).any():
-            policy_obs = torch.nan_to_num(policy_obs, nan=0.0)
+        if not torch.isfinite(code).all():
+            code = torch.nan_to_num(code, nan=0.0, posinf=0.0, neginf=0.0)
+        if not torch.isfinite(policy_obs).all():
+            policy_obs = torch.nan_to_num(policy_obs, nan=0.0, posinf=0.0, neginf=0.0)
 
         latent = torch.cat([code, policy_obs], dim=-1)  # (N, 64)
         return latent

motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py

@@ -81,8 +81,17 @@ class DreamWaQPPO(PPO):
             returns_batch, old_actions_log_prob_batch, old_mu_batch,
             old_sigma_batch, hid_states_batch, masks_batch,
         ) in generator:
-            # NaN 检测 — 数据有 NaN 就跳过这个 batch
-            if torch.isnan(obs_batch["policy"]).any() or torch.isnan(obs_batch["obs_history"]).any():
+            # 多层有限值检测 — 数据异常就跳过这个 batch，避免污染网络参数。
+            if (not torch.isfinite(obs_batch["policy"]).all()
+                or not torch.isfinite(obs_batch["obs_history"]).all()
+                or ("privileged_obs" in obs_batch and not torch.isfinite(obs_batch["privileged_obs"]).all())
+                or not torch.isfinite(actions_batch).all()
+                or not torch.isfinite(target_values_batch).all()
+                or not torch.isfinite(advantages_batch).all()
+                or not torch.isfinite(returns_batch).all()
+                or not torch.isfinite(old_actions_log_prob_batch).all()
+                or not torch.isfinite(old_mu_batch).all()
+                or not torch.isfinite(old_sigma_batch).all()):
                 continue
             # ── 标准 PPO 前向（各调用一次，复用结果）──
             self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
@@ -106,6 +115,11 @@ class DreamWaQPPO(PPO):
                 actions_batch, actions_log_prob_batch,
                 old_actions_log_prob_batch, advantages_batch)
 
+            # NaN guard: 损失值异常 → 跳过此 batch，保护模型权重
+            if (not torch.isfinite(surrogate_loss)
+                or not torch.isfinite(value_loss)):
+                continue
+
             # ── KL 自适应 schedule ──
             if self.desired_kl is not None and self.schedule == 'adaptive':
                 with torch.inference_mode():
@@ -136,20 +150,41 @@ class DreamWaQPPO(PPO):
                 - self.entropy_coef * entropy_loss
                 + autoenc_loss
             )
+            if not torch.isfinite(loss):
+                continue
 
             # ── 梯度更新 ──
             self.optimizer.zero_grad()
             loss.backward()
+            actor_grads_ok = all(
+                p.grad is None or torch.isfinite(p.grad).all()
+                for p in self.actor.parameters()
+            )
+            critic_grads_ok = all(
+                p.grad is None or torch.isfinite(p.grad).all()
+                for p in self.critic.parameters()
+            )
+            if not (actor_grads_ok and critic_grads_ok):
+                self.optimizer.zero_grad(set_to_none=True)
+                continue
             if self.max_grad_norm is not None:
-                nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
-                nn.utils.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
+                try:
+                    nn.utils.clip_grad_norm_(
+                        self.actor.parameters(), self.max_grad_norm, error_if_nonfinite=True)
+                    nn.utils.clip_grad_norm_(
+                        self.critic.parameters(), self.max_grad_norm, error_if_nonfinite=True)
+                except RuntimeError:
+                    self.optimizer.zero_grad(set_to_none=True)
+                    continue
                 nn.utils.clip_grad_value_(self.actor.parameters(), 10.0)
             self.optimizer.step()
             # 每次更新后强制 std > 0，防止数值异常导致 NaN
             if hasattr(self.actor, 'std') and self.actor.stochastic:
                 with torch.no_grad():
-                    self.actor.std.nan_to_num_(nan=0.5, posinf=1.0, neginf=1.0)
-                    self.actor.std.clamp_(min=1e-6)
+                    if not torch.isfinite(self.actor.std).all():
+                        self.actor.std.data.fill_(0.5)
+                    std_clip = getattr(self.actor, "std_clip", 0.6)
+                    self.actor.std.data.clamp_(min=1e-6, max=std_clip)
 
             # ── 累计日志 ──
             mean_value_loss += value_loss.item()

motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py

@@ -4,7 +4,7 @@
 CENet (VAE):  history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
 Decoder:       code(19) → [64,128] → next_obs(45)
 Actor:         code(19) + obs(45) = 64 → [512,256,128] → action(12)
-Critic:        privileged_obs(235) → [512,256,128] → value(1)
+Critic:        privileged_obs(247) → [512,256,128] → value(1)
 
 VAE Loss: reconstruction_MSE + velocity_MSE + beta * KL
 """
@@ -189,7 +189,7 @@ class DreamWaQWrapper:
     @property
     def privileged_obs(self):
         return self._env._state.info.get("privileged_obs",
-                np.zeros((self._num_envs, 235), dtype=np.float32))
+                np.zeros((self._num_envs, 247), dtype=np.float32))
 
     @property
     def obs_history(self):
@@ -398,7 +398,7 @@ class DreamWaQTrainer:
             def __call__(self, inputs, role):
                 kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
                 x = inputs["states"]
-                # Critic: obs(45) + base_vel(3) + heights(187) = 235
+                # Critic: obs(45) + base_vel(3) + heights(187) = 247
                 # Layout: [code(19) | obs(45) | base_vel(3) | heights(187)]
                 x_c = jnp.concatenate([x[:, 19:64], x[:, 64:254]], axis=-1)
                 for d in value_cfg.hiddens:

motrix_rl/src/motrix_rl/tasks/go1.py

@@ -112,7 +112,7 @@ class rslrl:
 
             # Runner 设置（严格对齐上游 LeggedRobotCfgPPO + Go1RoughCfgPPO）
             runner.seed = 5                     # 上游 seed=5
-            runner.max_iterations = 3000  # 续训到 3000 轮
+            runner.max_iterations = 5000
             runner.num_steps_per_env = 24
             runner.experiment_name = "go1_dreamwaq_walk"
             runner.save_interval = 50

scripts/diag_terrain_collision.py

@@ -0,0 +1,458 @@
+#!/usr/bin/env python3
+"""地形碰撞诊断工具 — 系统性检测 hfield PNG 地形的物理交互质量。
+
+检测项：
+  1. 脚部穿透（foot z < terrain z）
+  2. 接触力异常（为 0 但应该接触 / 过大）
+  3. NaN/Inf 传播
+  4. 物理爆炸（关节速度飙升）
+  5. 机器人卡住（速度接近 0 但有命令）
+  6. 躯干沉入地形（base z < terrain z）
+
+用法:
+    uv run scripts/diag_terrain_collision.py                     # 全量检测
+    uv run scripts/diag_terrain_collision.py --level 5           # 仅 level 5
+    uv run scripts/diag_terrain_collision.py --type stairs       # 仅楼梯地形
+    uv run scripts/diag_terrain_collision.py --render            # 渲染可视化
+"""
+import argparse, os, sys, time
+import numpy as np
+
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+os.environ.setdefault("DREAMWAQ_TERRAIN", "pyramid")
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import motrix_envs.locomotion.go1.dreamwaq  # noqa: F401
+from motrix_envs import registry as env_registry
+
+# ═══════════════════════════════════════════════════════════════════════
+# 配置
+# ═══════════════════════════════════════════════════════════════════════
+
+NUM_ROWS = 10
+NUM_COLS = 20
+CELL_M = 8.0
+BORDER_M = 5.0
+SETTLE_STEPS = 100     # 自由落体稳定步数（200Hz）
+TEST_STEPS = 400       # 测试步数
+STEPS_PER_CELL = 5     # 每个 cell 测试的 spawn 位置数
+
+# 地形类型名（与 gen_dreamwaq_terrain.py PROPORTIONS 对应）
+TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"]
+CUM = [0.1, 0.2, 0.55, 0.9, 1.0]
+
+
+def _cell_origin(row, col):
+    half_x = BORDER_M + NUM_COLS * CELL_M / 2.0
+    half_y = BORDER_M + NUM_ROWS * CELL_M / 2.0
+    cx = -half_x + BORDER_M + col * CELL_M + CELL_M / 2
+    cy = half_y - BORDER_M - row * CELL_M - CELL_M / 2
+    return cx, cy
+
+
+def _terrain_type(col):
+    choice = col / NUM_COLS + 0.001
+    for i, c in enumerate(CUM):
+        if choice < c:
+            return i
+    return len(TYPE_NAMES) - 1
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# 诊断核心
+# ═══════════════════════════════════════════════════════════════════════
+
+class CollisionDiag:
+    def __init__(self, env, render=False):
+        self.env = env
+        self.render = render
+        self.renderer = None
+        if render:
+            from motrix_envs.np.renderer import NpRenderer
+            self.renderer = NpRenderer(env)
+
+        # 统计数据
+        self.stats = {
+            "penetrations": [],       # (row, col, ttype, foot_z - terrain_z)
+            "contact_anomalies": [],  # (row, col, ttype, max_contact_force)
+            "nan_events": [],         # (row, col, ttype, step)
+            "explosions": [],         # (row, col, ttype, max_dof_vel)
+            "stuck_events": [],       # (row, col, ttype, speed)
+            "base_sink": [],          # (row, col, ttype, base_z - terrain_z)
+            "ok_count": 0,
+            "total_tested": 0,
+        }
+
+    def _get_terrain_z(self, x, y):
+        """查询 (x, y) 处的地形高度。"""
+        hm = self.env._hm_cache
+        if hm is None:
+            return 0.0
+        nr, nc = hm.shape
+        xmin = self.env._h_xmin
+        ymax = self.env._h_ymax
+        w = self.env._h_xmax - xmin
+        h = ymax - self.env._h_ymin
+        z_base = self.env._h_zbase
+        col = int(np.clip((x - xmin) / max(w, 1e-6) * (nc - 1), 0, nc - 1))
+        row = int(np.clip((ymax - y) / max(h, 1e-6) * (nr - 1), 0, nr - 1))
+        return float(hm[row, col]) + z_base
+
+    def _get_foot_positions(self, state):
+        """获取所有足部在世界坐标系中的位置 (N, 4, 3)。"""
+        data = state.data
+        n = data.shape[0]
+        feet = np.zeros((n, 4, 3), dtype=np.float32)
+        foot_names = ["FR_foot", "FL_foot", "RR_foot", "RL_foot"]
+        for i, name in enumerate(foot_names):
+            gidx = self.env._model.get_geom_index(name)
+            if gidx is not None:
+                pose = self.env._model.geoms[gidx].get_pose(data)
+                feet[:, i, :] = pose[:, :3]
+        return feet
+
+    def test_cell(self, row, col, level):
+        """在单个 cell 上测试多个 spawn 位置。"""
+        cx, cy = _cell_origin(row, col)
+        ttype = _terrain_type(col)
+        tname = TYPE_NAMES[ttype]
+
+        # 在 cell 内取几个 spawn 点（中央 + 四角附近）
+        offsets = [
+            (0, 0),
+            (CELL_M * 0.3, 0), (-CELL_M * 0.3, 0),
+            (0, CELL_M * 0.3), (0, -CELL_M * 0.3),
+        ][:STEPS_PER_CELL]
+
+        for ox, oy in offsets:
+            sx, sy = cx + ox, cy + oy
+            terrain_z = float(self.env._sample_terrain_height(
+                np.array([[sx, sy]], dtype=np.float32), radius=0.35)[0])
+
+            self.stats["total_tested"] += 1
+            self._test_single(sx, sy, terrain_z, row, col, level, ttype, tname)
+
+    def _test_single(self, sx, sy, terrain_z, row, col, level, ttype, tname):
+        """单个 spawn 点的完整测试。"""
+        env = self.env
+        data = env._state.data
+        n = 1
+
+        # Spawn
+        init_pos = env._init_dof_pos.copy().reshape(1, -1)
+        init_pos[0, 0] = sx
+        init_pos[0, 1] = sy
+        spawn_z = terrain_z + 0.45  # 标准 clearance
+        init_pos[0, 2] = spawn_z
+        yaw = np.random.uniform(-np.pi, np.pi)
+        init_pos[0, 3:7] = [0, 0, np.sin(yaw / 2), np.cos(yaw / 2)]
+
+        data.reset(env._model)
+        data.set_dof_pos(init_pos, env._model)
+        env._model.forward_kinematic(data)
+
+        # 给一个随机动作防止完全静止（小幅度 PD 控制）
+        action = 0.1 * (np.random.randn(n, 12).astype(np.float32))
+
+        has_nan = False
+        max_dof_vel = 0.0
+        max_contact_force = 0.0
+        min_foot_z = spawn_z
+        min_base_z = spawn_z
+        total_speed = 0.0
+
+        for step in range(SETTLE_STEPS + TEST_STEPS):
+            # Apply PD
+            actions_scaled = action * env.cfg.control_config.action_scale
+            target = actions_scaled + env.default_angles.reshape(1, -1)
+            lo = env._model.joint_limits[0].reshape(1, -1)
+            hi = env._model.joint_limits[1].reshape(1, -1)
+            target = np.clip(target, lo, hi)
+            torques = env.kps * (target - env.get_dof_pos(data)) - env.kds * env.get_dof_vel(data)
+            data.actuator_ctrls = np.clip(torques, -80.0, 80.0)
+
+            env._model.step(data)
+
+            # ── 检测 ──
+
+            # 1. NaN 检测
+            dv = np.array(data.dof_vel)
+            dp = np.array(data.dof_pos)
+            if np.any(~np.isfinite(dv)) or np.any(~np.isfinite(dp)):
+                self.stats["nan_events"].append((row, col, ttype, step))
+                has_nan = True
+                break
+
+            # 2. 速度爆炸
+            cur_max_vel = float(np.max(np.abs(dv[:, 6:])))
+            max_dof_vel = max(max_dof_vel, cur_max_vel)
+            if cur_max_vel > 100.0:
+                self.stats["explosions"].append((row, col, ttype, cur_max_vel))
+                break
+
+            # 3. 接触力
+            # Read contact forces via MotrixSim sensor API
+            cf_vals = []
+            for foot in ["FR", "FL", "RR", "RL"]:
+                try:
+                    v = env._model.get_sensor_value(f"{foot}_foot_contact", data)
+                    cf_vals.append(v[0])
+                except Exception:
+                    cf_vals.append(np.zeros(3, dtype=np.float32))
+            cf_raw = np.concatenate(cf_vals)
+            cur_cf = float(np.max(np.abs(cf_raw)))
+            max_contact_force = max(max_contact_force, cur_cf)
+
+            # 4. 足部穿透（settle 之后测）
+            if step >= SETTLE_STEPS:
+                feet = self._get_foot_positions(env._state)
+                for fi in range(4):
+                    fz = float(feet[0, fi, 2])
+                    tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1]))
+                    penetration = tz - fz  # 正 = 穿透
+                    min_foot_z = min(min_foot_z, fz - tz)  # 相对地形的高度
+                    if penetration > 0.02:  # > 2cm 穿透
+                        self.stats["penetrations"].append(
+                            (row, col, ttype, float(penetration)))
+
+            # 5. 躯干沉入
+            base_pose = env._body.get_pose(data)
+            bz = float(base_pose[0, 2])
+            tz = self._get_terrain_z(float(base_pose[0, 0]), float(base_pose[0, 1]))
+            min_base_z = min(min_base_z, bz - tz)
+            if bz < tz + 0.10 and step >= SETTLE_STEPS:  # 躯干离地面 < 10cm
+                self.stats["base_sink"].append((row, col, ttype, float(bz - tz)))
+
+            # 6. 速度（卡住检测）
+            if step >= SETTLE_STEPS:
+                linvel = env.get_local_linvel(data)[0]
+                total_speed += float(np.linalg.norm(linvel[:2]))
+
+            if self.renderer and step % 20 == 0:
+                self.renderer.render()
+                time.sleep(0.005)
+
+        # ── 统计 ──
+        if not has_nan and max_dof_vel < 100.0:
+            avg_speed = total_speed / max(TEST_STEPS, 1)
+            if avg_speed < 0.01 and step == SETTLE_STEPS + TEST_STEPS - 1:
+                self.stats["stuck_events"].append((row, col, ttype, avg_speed))
+
+        if not has_nan and max_dof_vel < 50.0 and min_base_z > 0.05:
+            self.stats["ok_count"] += 1
+
+        # 记录接触力异常（完全无接触力但足部应该着地）
+        if max_contact_force < 1e-3 and step == SETTLE_STEPS + TEST_STEPS - 1:
+            feet = self._get_foot_positions(env._state)
+            for fi in range(4):
+                fz = float(feet[0, fi, 2])
+                tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1]))
+                if abs(fz - tz) < 0.05:  # 足部接近地面
+                    self.stats["contact_anomalies"].append(
+                        (row, col, ttype, 0.0))  # 0 = 无接触力
+                    break
+
+    def run(self, target_level=None, target_type=None):
+        """遍历所有地形 cell 测试。"""
+        total = 0
+        start_t = time.time()
+
+        for row in range(NUM_ROWS):
+            if target_level is not None and row != target_level:
+                continue
+            for col in range(NUM_COLS):
+                ttype = _terrain_type(col)
+                if target_type is not None and ttype != target_type:
+                    continue
+                total += 1
+
+        print(f"[Diag] 测试 {NUM_ROWS}×{NUM_COLS} 网格, "
+              f"每 cell {STEPS_PER_CELL} 个 spawn 点 = {total * STEPS_PER_CELL} 次")
+        print(f"[Diag] Settle={SETTLE_STEPS}步 Test={TEST_STEPS}步")
+
+        count = 0
+        for row in range(NUM_ROWS):
+            if target_level is not None and row != target_level:
+                continue
+            for col in range(NUM_COLS):
+                ttype = _terrain_type(col)
+                if target_type is not None and ttype != target_type:
+                    continue
+                self.test_cell(row, col, row)
+                count += 1
+                if count % 10 == 0:
+                    elapsed = time.time() - start_t
+                    print(f"  [{count}/{total}] {elapsed:.1f}s  "
+                          f"OK={self.stats['ok_count']}/{self.stats['total_tested']}  "
+                          f"穿透={len(self.stats['penetrations'])}  "
+                          f"NaN={len(self.stats['nan_events'])}  "
+                          f"爆炸={len(self.stats['explosions'])}")
+
+        elapsed = time.time() - start_t
+        self._report(elapsed)
+
+    def _report(self, elapsed):
+        s = self.stats
+        t = s["total_tested"]
+        print("\n" + "=" * 70)
+        print(f"地形碰撞诊断报告 ({elapsed:.1f}s, {t} 次测试)")
+        print("=" * 70)
+
+        # 总体统计
+        ok_rate = s["ok_count"] / max(t, 1) * 100
+        print(f"\n  ✅ 通过: {s['ok_count']}/{t} ({ok_rate:.1f}%)")
+
+        # 穿透
+        if s["penetrations"]:
+            pens = s["penetrations"]
+            pvals = [p[3] for p in pens]
+            print(f"\n  🔴 足部穿透: {len(pens)} 次")
+            print(f"     最大穿透: {max(pvals):.3f}m  平均: {np.mean(pvals):.3f}m")
+
+            # 按地形类型分组
+            by_type = {}
+            for p in pens:
+                tn = TYPE_NAMES[p[2]]
+                by_type.setdefault(tn, []).append(p[3])
+            for tn, vals in sorted(by_type.items()):
+                print(f"     {tn}: {len(vals)}次  max={max(vals):.3f}m  avg={np.mean(vals):.3f}m")
+
+            # 按 level 分组
+            by_level = {}
+            for p in pens:
+                by_level.setdefault(p[0], []).append(p[3])
+            print(f"     按 level: ", end="")
+            for lv in sorted(by_level.keys()):
+                vals = by_level[lv]
+                print(f"L{lv}({len(vals)}x max={max(vals):.3f}) ", end="")
+            print()
+        else:
+            print(f"\n  🟢 足部穿透: 0 次 ✓")
+
+        # NaN
+        if s["nan_events"]:
+            print(f"\n  🔴 NaN/Inf: {len(s['nan_events'])} 次")
+            for n in s["nan_events"][:5]:
+                print(f"     L{n[0]} C{n[1]:02d} {TYPE_NAMES[n[2]]} @step {n[3]}")
+        else:
+            print(f"  🟢 NaN/Inf: 0 次 ✓")
+
+        # 物理爆炸
+        if s["explosions"]:
+            print(f"\n  🟡 速度爆炸 (>{100}rad/s): {len(s['explosions'])} 次")
+            for e in s["explosions"][:5]:
+                print(f"     L{e[0]} C{e[1]:02d} {TYPE_NAMES[e[2]]} max_vel={e[3]:.1f}")
+        else:
+            print(f"  🟢 速度爆炸: 0 次 ✓")
+
+        # 接触力异常
+        if s["contact_anomalies"]:
+            print(f"\n  🟡 接触力异常 (无接触但足部近地): {len(s['contact_anomalies'])} 次")
+        else:
+            print(f"  🟢 接触力异常: 0 次 ✓")
+
+        # 卡住
+        if s["stuck_events"]:
+            print(f"\n  🟡 卡住 (速度<0.01m/s): {len(s['stuck_events'])} 次")
+            by_type = {}
+            for e in s["stuck_events"]:
+                tn = TYPE_NAMES[e[2]]
+                by_type.setdefault(tn, []).append(1)
+            for tn, vals in sorted(by_type.items()):
+                print(f"     {tn}: {len(vals)}次")
+        else:
+            print(f"  🟢 卡住: 0 次 ✓")
+
+        # 躯干沉入
+        if s["base_sink"]:
+            sinks = s["base_sink"]
+            print(f"\n  🟡 躯干贴地 (<10cm): {len(sinks)} 次")
+            vals = [s[3] for s in sinks]
+            print(f"     最差: {min(vals):.3f}m  平均: {np.mean(vals):.3f}m")
+        else:
+            print(f"  🟢 躯干贴地: 0 次 ✓")
+
+        # 按地形类型汇总
+        print(f"\n── 按地形类型汇总 ──")
+        for ti, tn in enumerate(TYPE_NAMES):
+            pens_t = sum(1 for p in s["penetrations"] if p[2] == ti)
+            nans_t = sum(1 for n in s["nan_events"] if n[2] == ti)
+            exps_t = sum(1 for e in s["explosions"] if e[2] == ti)
+            icon = "🔴" if (pens_t + nans_t + exps_t) > 0 else "🟢"
+            print(f"  {icon} {tn}: 穿透{pens_t} NaN{nans_t} 爆炸{exps_t}")
+
+        # 结论
+        print(f"\n── 全局结论 ──")
+        if ok_rate > 95 and len(s["nan_events"]) == 0 and len(s["penetrations"]) < 5:
+            print(f"  ✅ 仿真环境正常 — hfield 地形碰撞基本可靠")
+        elif ok_rate > 80:
+            print(f"  ⚠ 仿真环境存在一些问题 — 部分地形类型/level 有问题")
+        else:
+            print(f"  🔴 仿真环境问题严重 — 需要排查 MotrixSim hfield 实现")
+
+        if len(s["penetrations"]) > 10:
+            print(f"  💡 穿透较多 → hfield 碰撞面可能偏软或 contact margin 过大")
+            print(f"     可尝试增大 geom margin 或检查 MotrixSim 碰撞参数")
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Main
+# ═══════════════════════════════════════════════════════════════════════
+
+def main():
+    p = argparse.ArgumentParser(description="DreamWaQ 地形碰撞诊断")
+    p.add_argument("--level", type=int, default=None, help="仅测试指定 level (0-9)")
+    p.add_argument("--type", dest="ttype", default=None,
+                   choices=["slope", "rough", "stairs_down", "stairs_up", "obstacles"],
+                   help="仅测试指定地形类型")
+    p.add_argument("--render", action="store_true", help="渲染可视化")
+    p.add_argument("--settle", type=int, default=SETTLE_STEPS)
+    p.add_argument("--test", type=int, default=TEST_STEPS)
+    args = p.parse_args()
+
+    type_map = {"slope": 0, "rough": 1, "stairs_down": 2, "stairs_up": 3, "obstacles": 4}
+    target_type = type_map.get(args.ttype)
+
+    print("[Diag] 创建 DreamWaQ 环境...")
+    env = env_registry.make("go1-dreamwaq-walk", num_envs=1)
+
+    # 如果指定 level，强制执行
+    if args.level is not None:
+        env._force_level = args.level
+        print(f"[Diag] 强制 level={args.level}")
+
+    env.init_state()
+
+    # 确保 terrain 初始化（触发 hfield cache）
+    if env._model.num_hfields > 0:
+        hf = env._model.get_hfield(0)
+        env._hm_cache = hf.height_matrix
+        env._h_xmin, env._h_ymin = hf.bound[0], hf.bound[1]
+        env._h_xmax, env._h_ymax = hf.bound[3], hf.bound[4]
+        env._h_nr, env._h_nc = hf.height_matrix.shape
+        cfg_base = getattr(env.cfg, "hfield_z_base", None)
+        env._h_zbase = float(cfg_base) if cfg_base is not None else float(-hf.height_matrix[0, 0])
+        print(f"[Diag] HField: {hf.height_matrix.shape}  bound={hf.bound}")
+        print(f"[Diag] z_base={env._h_zbase:.3f}")
+    else:
+        print("[Diag] WARNING: 无 hfield（flat 场景），跳过诊断")
+        return
+
+    diag = CollisionDiag(env, render=args.render)
+    diag.SETTLE_STEPS = args.settle
+    diag.TEST_STEPS = args.test
+
+    try:
+        diag.run(target_level=args.level, target_type=target_type)
+    except KeyboardInterrupt:
+        print("\n[Diag] 中断")
+    finally:
+        if diag.renderer:
+            diag.renderer.close()
+        print("[Diag] 完成")
+
+
+if __name__ == "__main__":
+    main()

scripts/dreamwaq_sim2sim_mujoco.py

@@ -31,7 +31,8 @@ HISTORY_LEN = 5
 ACTION_SCALE = 0.25
 KP = 28.0
 KD = 0.7
-CLIP_ACTIONS = 23.7
+CLIP_ACTIONS = 4.0
+CLIP_TORQUES = 80.0
 CLIP_OBS = 100.0
 MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0
 
@@ -122,7 +123,7 @@ def main():
 
     # Select XML scene
     terrain_map = {
-        "flat": "scene_motor_actuator.xml",
+        "flat": "scene_dreamwaq_flat.xml",
         "rough": "scene_rough_terrain.xml",
         "stairs": "scene_stairs_terrain.xml",
         "dreamwaq": "scene_dreamwaq_terrain.xml",
@@ -243,10 +244,22 @@ def main():
             action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
             last_action = action.copy()
 
-        # PD control
+        # PD control（对齐训练：目标限位 + 力矩裁剪）
         target = DEFAULT_ANGLES + action * ACTION_SCALE
+        # 关节目标限位（与训练一致）
+        jnt_lo = model.jnt_range[:, 0].copy() if hasattr(model, 'jnt_range') else None
+        jnt_hi = model.jnt_range[:, 1].copy() if hasattr(model, 'jnt_range') else None
+        # MuJoCo model.actuator_trnid 可能不直接暴露，改用 model.jnt_range
+        try:
+            trnid = model.actuator_trnid[:, 0]  # transmission joint indices
+            lo = model.jnt_range[trnid, 0]
+            hi = model.jnt_range[trnid, 1]
+        except Exception:
+            lo = np.full(12, -12.0)
+            hi = np.full(12, 12.0)
+        target = np.clip(target, lo, hi)
         torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
-        data.ctrl[:] = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
+        data.ctrl[:] = np.clip(torques, -CLIP_TORQUES, CLIP_TORQUES)
         mujoco.mj_step(model, data)
         view.sync()
 

scripts/export_dreamwaq_onnx_new.py

@@ -28,7 +28,7 @@ class DwaqInfer(torch.nn.Module):
         """obs: (N,45), obs_history: (N,225) → action: (N,12)"""
         code = self.vae.deterministic_code(obs_history)  # (N,19)
         latent = torch.cat([code, obs], dim=-1)  # (N,64)
-        return self.actor_mlp(latent)
+        return torch.clamp(self.actor_mlp(latent), -4.0, 4.0)
 
 
 def main():

scripts/export_dreamwaq_rsl_onnx.py

@@ -22,13 +22,20 @@ PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
 
 class DwaqInfer(tnn.Module):
-    """Deterministic inference: CENet mean code + actor."""
+    """Deterministic inference: CENet mean code + RunningStats normalize + actor."""
     def __init__(self, ac: ActorCritic_DWAQ):
         super().__init__()
         self.encoder = ac.encoder
         self.encode_mean_vel = ac.encode_mean_vel
         self.encode_mean_latent = ac.encode_mean_latent
         self.actor = ac.actor
+        # Bake RunningStats into the exported model
+        if hasattr(ac, 'actor_normalizer') and ac.actor_normalizer.count > 10:
+            self.register_buffer("norm_mean", ac.actor_normalizer.mean.clone())
+            self.register_buffer("norm_std", ac.actor_normalizer.var.sqrt().clone())
+            self.has_norm = True
+        else:
+            self.has_norm = False
 
     def forward(self, obs, obs_history):
         h = self.encoder(obs_history.reshape(obs_history.shape[0], -1))  # (B,225)->(B,64)
@@ -36,6 +43,8 @@ class DwaqInfer(tnn.Module):
         latent = self.encode_mean_latent(h)  # (B,16) mean latent
         code = torch.cat([vel, latent], dim=-1)   # (B,19) = [vel, latent]
         x = torch.cat([code, obs], dim=-1)        # (B,64) = [code, obs]
+        if self.has_norm:
+            x = torch.clamp((x - self.norm_mean) / (self.norm_std + 1e-8), -5.0, 5.0)
         return self.actor(x)                      # (B,12)
 
 

scripts/play_amp_mujoco.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Deploy Go1 AMP model in MuJoCo."""
+import mujoco, numpy as np, os, torch, time, sys
+from mujoco import viewer
+
+# Use AMP's rsl_rl (NOT MotrixLab's)
+sys.path.insert(0, '/home/8x54zj-m/amp_go2/rsl_rl')
+sys.path.insert(0, '/home/8x54zj-m/amp_go2/legged_gym')
+# Force reload of rsl_rl to pick up AMP version
+for mod in list(sys.modules.keys()):
+    if 'rsl_rl' in mod:
+        del sys.modules[mod]
+
+XML_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+                       'motrix_envs/src/motrix_envs/locomotion/go1/xmls')
+CKPT = '/home/8x54zj-m/amp_go2/legged_gym/logs/go1_amp/Jul05_01-15-41_flat/model_6000.pt'
+
+import argparse
+p = argparse.ArgumentParser()
+p.add_argument("--terrain", default="flat", choices=["flat", "flat_stairs"])
+args = p.parse_args()
+
+scene = "scene_flat_stairs.xml" if args.terrain == "flat_stairs" else "go1_motor_actuator.xml"
+os.chdir(XML_DIR)
+m = mujoco.MjModel.from_xml_string(open(scene).read())
+d = mujoco.MjData(m)
+
+# Spawn on level 1 platform if using flat_stairs
+spawn_x, spawn_y, spawn_z = 0.0, 0.0, 0.45
+if args.terrain == "flat_stairs":
+    col = np.random.randint(0, 4)
+    spawn_x = -12.0 + col * 8.0
+    spawn_y = -4.0  # level 1 center
+    # Sample terrain height at spawn
+    floor_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "floor")
+    if floor_id >= 0 and m.geom_type[floor_id] == mujoco.mjtGeom.mjGEOM_HFIELD:
+        hf = m.geom_dataid[floor_id]
+        nr, nc = int(m.hfield_nrow[hf]), int(m.hfield_ncol[hf])
+        sx, sy, zt, sb = m.hfield_size[hf]
+        hd = m.hfield_data[m.hfield_adr[hf]:m.hfield_adr[hf]+nr*nc].reshape(nr, nc)
+        gp = m.geom_pos[floor_id]
+        col_idx = int(np.clip(((spawn_x-gp[0])/sx*0.5+0.5)*(nc-1), 0, nc-1))
+        row_idx = int(np.clip(((spawn_y-gp[1])/sy*0.5+0.5)*(nr-1), 0, nr-1))
+        spawn_z = float(gp[2] + sb + hd[row_idx, col_idx] * zt) + 0.45
+    print(f"[Terrain] flat_stairs level=1 spawn=({spawn_x:.1f},{spawn_y:.1f},{spawn_z:.2f})")
+
+ckpt = torch.load(CKPT, map_location='cpu', weights_only=False)
+normalizer = ckpt['amp_normalizer']
+
+actor = torch.nn.Sequential(
+    torch.nn.Linear(45, 512), torch.nn.ELU(),
+    torch.nn.Linear(512, 256), torch.nn.ELU(),
+    torch.nn.Linear(256, 128), torch.nn.ELU(),
+    torch.nn.Linear(128, 12)
+)
+sd = {k.replace('actor.',''): v for k,v in ckpt['model_state_dict'].items() if k.startswith('actor.')}
+actor.load_state_dict(sd)
+actor.eval()
+
+DEFAULT = np.array([-0.1,0.8,-1.5, 0.1,0.8,-1.5, -0.1,1.0,-1.5, 0.1,1.0,-1.5], dtype=np.float32)
+LIN_VEL_SCALE, ANG_VEL_SCALE = 2.0, 0.25
+DOF_POS_SCALE, DOF_VEL_SCALE = 1.0, 0.05
+CMD_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32)
+KP, KD, ACT_SCALE = 20.0, 0.5, 0.25
+
+def get_gravity(q):
+    qw,qx,qy,qz = q[3],q[0],q[1],q[2]
+    return np.array([2*(-qz*qx+qw*qy), -2*(qz*qy+qw*qx), 1-2*(qw*qw+qz*qz)])
+
+def compute_obs(d, last_action, cmd):
+    obs = np.zeros(45, dtype=np.float32)
+    obs[0:3] = d.qvel[3:6] * ANG_VEL_SCALE
+    obs[3:6] = get_gravity(d.qpos[3:7])
+    obs[6:9] = cmd * CMD_SCALE
+    obs[9:21] = (d.qpos[7:19] - DEFAULT) * DOF_POS_SCALE
+    obs[21:33] = d.qvel[6:18] * DOF_VEL_SCALE
+    obs[33:45] = last_action
+    return obs
+
+d.qpos[7:19] = DEFAULT; d.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; d.qvel[:] = 0
+mujoco.mj_forward(m, d)
+
+view = viewer.launch_passive(m, d)
+cmd = np.array([0.8, 0.0, 0.0], dtype=np.float32)
+last_action = np.zeros(12, dtype=np.float32)
+
+print("AMP Go1 | W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
+
+while view.is_running():
+    if view.is_running():
+        obs = compute_obs(d, last_action, cmd)
+        obs_t = torch.from_numpy(obs).unsqueeze(0).float()
+        obs_norm = obs_t.clone()
+        obs_norm[:, :43] = normalizer.normalize(obs_t[:, :43])
+        with torch.no_grad():
+            act = actor(obs_norm)[0].numpy()
+        act = np.clip(act, -18, 18)
+        last_action = act.copy()
+        target = DEFAULT + act * ACT_SCALE
+        d.ctrl[:] = KP * (target - d.qpos[7:19]) - KD * d.qvel[6:18]
+    mujoco.mj_step(m, d)
+    view.sync()
+    time.sleep(0.001)

scripts/test_abrupt_stop.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""无头 MuJoCo 测试：前进 → 突然停止 → 观察姿态变化。"""
+import numpy as np
+import mujoco
+import onnxruntime as ort
+import os, sys, time
+
+_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+XML_DIR = os.path.join(_PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
+ONNX = os.path.join(_PROJECT, "runs/go1-dreamwaq-walk/rslrl/26-07-02_20-03-23-_36054_PPO/policy.onnx")
+
+NUM_OBS = 45
+NUM_ACTIONS = 12
+HISTORY_LEN = 5
+ACTION_SCALE = 0.25
+KP = 28.0
+KD = 0.7
+DECIMATION = 4
+DEFAULT_ANGLES = np.array([0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8], dtype=np.float32)
+
+
+def compute_obs(model, data, commands, last_action):
+    obs = np.zeros(NUM_OBS, dtype=np.float32)
+    g = None  # no gyro sensor in headless
+    obs[0:3] = (g if g is not None else data.qvel[3:6]) * 0.25
+    grav_world = model.opt.gravity.copy()
+    grav_world = grav_world / np.linalg.norm(grav_world)
+    R = data.xmat[1].reshape(3, 3)
+    obs[3:6] = (R.T @ grav_world).astype(np.float32)
+    obs[6:9] = commands * np.array([2.0, 2.0, 0.25], dtype=np.float32)
+    obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
+    obs[21:33] = data.qvel[6:18] * 0.05
+    obs[33:45] = last_action
+    return obs
+
+
+def get_base_pose(data):
+    """返回 base 的 z 高度和 roll/pitch（度）。"""
+    quat = data.xquat[1]  # base body quaternion
+    w, x, y, z = quat[0], quat[1], quat[2], quat[3]
+    # roll, pitch from quaternion
+    sinr_cosp = 2 * (w * x + y * z)
+    cosr_cosp = 1 - 2 * (x * x + y * y)
+    roll = np.arctan2(sinr_cosp, cosr_cosp)
+    sinp = 2 * (w * y - z * x)
+    pitch = np.arcsin(np.clip(sinp, -1, 1))
+    return float(data.xpos[1, 2]), np.degrees(roll), np.degrees(pitch)
+
+
+def main():
+    session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider'])
+    xml_file = os.path.join(XML_DIR, "scene_dreamwaq_flat.xml")
+    model = mujoco.MjModel.from_xml_path(xml_file)
+    data = mujoco.MjData(model)
+
+    # 初始化
+    data.qpos[7:19] = DEFAULT_ANGLES
+    data.qpos[2] = 0.35
+    mujoco.mj_forward(model, data)
+
+    history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32)
+    last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
+
+    print(f"{'Step':>5s}  {'Cmd_vx':>7s}  {'base_z':>8s}  {'roll':>8s}  {'pitch':>8s}  {'speed_xy':>8s}")
+    print("-" * 60)
+
+    for step in range(3000):
+        # 命令：前 1500 步前进，之后突然停止
+        if step < 1500:
+            vx, vy, wz = 0.5, 0.0, 0.0  # 前进
+        else:
+            vx, vy, wz = 0.0, 0.0, 0.0  # 突然停止！
+
+        for _ in range(DECIMATION):
+            mujoco.mj_step(model, data)
+
+        if step % DECIMATION == 0:
+            cmd = np.array([vx, vy, wz], dtype=np.float32)
+            obs = compute_obs(model, data, cmd, last_action)
+            history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1)
+            outputs = session.run(None, {
+                'obs': obs.reshape(1, -1).astype(np.float32),
+                'obs_history': history.reshape(1, -1).astype(np.float32),
+            })
+            action = outputs[0][0]
+            last_action = action
+
+        target = DEFAULT_ANGLES + action * ACTION_SCALE
+        data.ctrl[:] = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
+
+        # 每 100 步打印
+        if step % 100 == 0:
+            bz, roll, pitch = get_base_pose(data)
+            speed_xy = np.linalg.norm(data.qvel[0:2])
+            print(f"{step:5d}  {vx:7.1f}  {bz:8.3f}  {roll:+8.1f}  {pitch:+8.1f}  {speed_xy:8.3f}")
+
+    # 最终状态
+    bz, roll, pitch = get_base_pose(data)
+    print(f"\n最终: base_z={bz:.3f} roll={roll:.1f}° pitch={pitch:.1f}°")
+    if abs(roll) > 30 or abs(pitch) > 30:
+        print("⚠ 机器人倾覆！突然停止导致翻跟头")
+    else:
+        print("✅ 机器人在突然停止后保持稳定")
+
+
+if __name__ == "__main__":
+    main()