feat: DreamWaQ full replication — env, terrain, CENet, PPO

docs/dreamwaq_replication.md

@@ -0,0 +1,57 @@
+# DreamWaQ 复刻进度
+
+## 总览
+
+| Phase | 状态 | 完成时间 |
+|-------|------|----------|
+| Phase 1: 环境差距补齐 | ✅ | 2026-06-30 |
+| Phase 2: 地形生成 10×20 | ✅ | 2026-06-30 |
+| Phase 3: CENet 网络集成 | ✅ | 2026-06-30 |
+| Phase 4: 训练 Pipeline | ✅ | 2026-06-30 |
+| Phase 5: 验证 | ⏳ | - |
+
+## Phase 3+4: CENet + 训练
+
+**新增文件**：
+- `motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py` — CENetVAE + CENetActorModel
+- `motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py` — DreamWaQPPO（PPO + VAE loss）
+
+**修改文件**：
+- `wrap_vec_env.py` — 多键 TensorDict（policy + obs_history + privileged_obs）
+- `cfg.py` — 添加 vae_beta/cenet_in_dim/cenet_out_dim
+- `go1.py` — go1-dreamwaq-walk 训练配置
+
+**冒烟测试**：1 iteration 通过，VAE 损失 ≈ 2.85
+
+### 数据流
+
+```
+DreamWaQTask.update_observation()
+  → state.info["obs_history"]       (N, 5, 45)
+  → state.info["privileged_obs"]    (N, 247)
+
+RslrlNpEnvWrap._build_obs_dict()
+  → TensorDict({
+      "policy":         (N, 45),
+      "obs_history":    (N, 225),
+      "privileged_obs": (N, 247),
+    })
+
+CENetActorModel.get_latent()
+  → VAE(obs_history) → code(19)
+  → cat(code, policy) → latent(64)
+  → MLP[512,256,128] → action(12)
+
+DreamWaQPPO.update()
+  → + VAE loss（速度估计 + 重建 + KL）
+```
+
+## 训练命令
+
+```bash
+# 平坦地形
+uv run scripts/train.py --env go1-dreamwaq-walk --rllib rslrl
+
+# 金字塔地形
+DREAMWAQ_TERRAIN=pyramid uv run scripts/train.py --env go1-dreamwaq-walk --rllib rslrl
+```

docs/dreamwaq_usage.md

@@ -0,0 +1,78 @@
+# DreamWaQ 使用手册
+
+## 环境切换
+
+通过环境变量 `DREAMWAQ_TERRAIN` 选择地形场景：
+
+| 值 | 场景 | 描述 |
+|----|------|------|
+| `flat` (默认) | 无限平面 | 学习基础行走 |
+| `pyramid` | 10×20 混合地形 | hfield + mesh 楼梯 |
+| `flat_stairs` | 2 级 flat+stairs | 简易楼梯测试 |
+| `stairs` | 纯楼梯 | stair box 场景 |
+
+```bash
+# 金字塔地形训练
+DREAMWAQ_TERRAIN=pyramid uv run scripts/train.py --env go1-dreamwaq-walk --rllib rslrl
+
+# 平坦地形训练
+uv run scripts/train.py --env go1-dreamwaq-walk --rllib rslrl
+```
+
+## 地形生成
+
+```bash
+# 完整 10×20 地形（生成 PNG + OBJ + XML）
+uv run python3 scripts/gen_dreamwaq_terrain.py
+
+# 仅前 N 个难度级别（测试用）
+uv run python3 scripts/gen_dreamwaq_terrain.py --max-level 3
+
+# 仅平坦（不生成楼梯）
+uv run python3 scripts/gen_dreamwaq_terrain.py --flat-only
+```
+
+## 可视化
+
+```bash
+# 查看金字塔地形（随机动作）
+uv run scripts/view_dreamwaq.py
+
+# 平坦地形 + 单机器人 + 固定难度
+uv run scripts/view_dreamwaq.py --flat --num-envs 1 --level 5
+
+# 持续前进（不站立）
+uv run scripts/view_dreamwaq.py --no-stand --vx 0.8
+```
+
+## 训练
+
+```bash
+# 启动训练
+uv run scripts/train.py --env go1-dreamwaq-walk --rllib rslrl
+
+# 指定环境数量
+uv run scripts/train.py --env go1-dreamwaq-walk --rllib rslrl --num-envs 4096
+```
+
+训练结果保存在 `runs/go1-dreamwaq-walk/rslrl/` 下，TensorBoard 日志自动记录：
+- `value_loss`, `surrogate_loss` — 标准 PPO 损失
+- `autoenc_loss` — CENet VAE 自编码器损失
+
+```bash
+# 查看训练曲线
+uv run tensorboard --logdir runs/go1-dreamwaq-walk
+```
+
+## Play（评估）
+
+```bash
+# 自动发现最新 checkpoint
+uv run scripts/play_dreamwaq_rsl.py
+
+# 指定 checkpoint + 命令
+uv run scripts/play_dreamwaq_rsl.py --checkpoint runs/go1-dreamwaq-walk/rslrl/.../model_1000.pt --vx 0.5
+
+# 固定地形级别查看
+uv run scripts/play_dreamwaq_rsl.py --terrain --level 5 --num-envs 1
+```

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

@@ -0,0 +1,625 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+"""DreamWaQ environment for MotrixLab — 对齐 Manaro-Alpha/DreamWaQ。
+
+架构:
+  CENet (VAE):  history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
+  Actor:        code(19) + obs(45) = 64 → [512,256,128] → action(12)
+  Critic:       privileged_obs(247) → [512,256,128] → value(1)
+
+观测 (45-dim, 上游顺序):
+  ang_vel(3) + gravity(3) + commands(3) + joint_pos(12) + joint_vel(12) + actions(12)
+
+特权观测 (247-dim):
+  obs(45) + base_vel(3) + 足部接触力(12) + heights(187)
+  注意：上游为 286（含全部 17 body × 3-axis 接触力），MotrixSim 仅用四足接触力
+
+注册: "go1-dreamwaq-walk"
+"""
+import gymnasium as gym
+import motrixsim as mtx
+import numpy as np
+from dataclasses import dataclass, field
+import os
+
+from motrix_envs import registry
+from motrix_envs.locomotion.go1.cfg import Go1WalkNpEnvCfg
+from motrix_envs.locomotion.go1.walk_np import Go1WalkTask
+from motrix_envs.math import quaternion
+
+
+_SCENE_PRINTED = False
+
+
+def _scene_file():
+    """选择地形场景。DREAMWAQ_TERRAIN=flat|pyramid|flat_stairs。默认 FLAT。"""
+    global _SCENE_PRINTED
+    terrain = os.environ.get("DREAMWAQ_TERRAIN", "flat").lower()
+    if terrain in ("flat_stairs", "flatstairs"):
+        fname = "scene_flat_stairs.xml"
+    elif terrain in ("pyramid", "terrain", "1"):
+        fname = "scene_dreamwaq_terrain.xml"
+    elif terrain in ("stairs", "stairs_terrain"):
+        fname = "scene_stairs_terrain.xml"
+    else:
+        fname = "scene_dreamwaq_flat.xml"
+    if not _SCENE_PRINTED:
+        print(f"[DreamWaQ] scene = {fname}  [set DREAMWAQ_TERRAIN to switch]")
+        _SCENE_PRINTED = True
+    return os.path.join(os.path.dirname(__file__), "xmls", fname)
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Config
+# ═══════════════════════════════════════════════════════════════════════
+
+@registry.envcfg("go1-dreamwaq-walk")
+@dataclass
+class DreamWaQCfg(Go1WalkNpEnvCfg):
+    """DreamWaQ 配置 — 10×20 地形网格，5 种类型，10 个难度级别。"""
+
+    model_file: str = field(default_factory=_scene_file)
+
+    # CENet 参数
+    num_latent: int = 16
+    cenet_out_dim: int = 19         # code = vel(3) + latent(16)
+    num_history: int = 5
+
+    # 特权观测维度：obs(45) + base_vel(3) + 足部接触力(12) + heights(187) = 247
+    num_privileged_obs: int = 247
+
+    max_episode_steps: int = 1000  # 与上游一致: 20s / 0.02s ctrl_dt
+    sim_dt: float = 0.005
+    ctrl_dt: float = 0.02
+
+    # 高度测量网格（与上游一致: 17×11）
+    height_points_x: tuple = tuple(np.linspace(-0.8, 0.8, 17).tolist())
+    height_points_y: tuple = tuple(np.linspace(-0.5, 0.5, 11).tolist())
+
+    # 地形网格
+    terrain_rows: int = 10
+    terrain_cols: int = 20  # 上游 10×20 网格：5 类型 × 10 难度
+    cell_size: float = 8.0
+    border_size: float = 5.0
+
+    # 命令范围（上游: [-1,1] 对称）
+    @dataclass
+    class Commands:
+        vel_limit = [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]
+
+    commands: Commands = field(default_factory=Commands)
+
+    def __init__(self):
+        super().__init__()
+        self.model_file = _scene_file()
+        self.num_latent = 16
+        self.cenet_out_dim = 19
+        self.num_history = 5
+        self.num_privileged_obs = 247
+        self.sim_dt = 0.005
+        self.ctrl_dt = 0.02
+        if "flat_stairs" in self.model_file:
+            self.terrain_rows = 2
+        else:
+            self.terrain_rows = 10
+        self.terrain_cols = 20
+        self.cell_size = 8.0
+        self.border_size = 5.0
+        self.height_points_x = tuple(np.linspace(-0.8, 0.8, 17).tolist())
+        self.height_points_y = tuple(np.linspace(-0.5, 0.5, 11).tolist())
+        self.commands = DreamWaQCfg.Commands()
+        self.control_config.stiffness = 28.0
+        self.control_config.damping = 0.7
+        self._apply_dreamwaq()
+
+    def _apply_dreamwaq(self):
+        """DreamWaQ 奖励尺度 — 与上游 Go1RoughCfg 一致（上游会在 _prepare 中 × dt）。"""
+        r = self.reward_config.scales
+        r.clear()
+        r.update({
+            "tracking_lin_vel": 1.0,
+            "tracking_ang_vel": 0.5,
+            "lin_vel_z": -2.0,
+            "ang_vel_xy": -0.05,
+            "orientation": -0.2,
+            "dof_acc": -2.5e-7,
+            "base_height": -1.0,
+            "feet_air_time": 0.1,
+            "action_rate": -0.01,
+            "joint_power": -2e-5,
+            "smoothness": -0.01,
+            "power_distribution": -10e-6,
+            # 注意：stand_still 在上游被注释掉
+        })
+        self.reward_config.only_positive_rewards = True
+        self.reward_config.tracking_sigma = 0.25
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# DreamWaQ 环境
+# ═══════════════════════════════════════════════════════════════════════
+
+@registry.env("go1-dreamwaq-walk", sim_backend="np")
+class DreamWaQTask(Go1WalkTask):
+    """DreamWaQ 环境：不对称观测、历史缓冲区、高度测量。"""
+
+    _cfg: DreamWaQCfg
+
+    def __init__(self, cfg: DreamWaQCfg = None, num_envs=1):
+        if cfg is None:
+            cfg = DreamWaQCfg()
+        super().__init__(cfg, num_envs)
+
+        # 覆盖观测空间（DreamWaQ 45-dim，比基类少 linvel 3-dim）
+        self._observation_space = gym.spaces.Box(
+            low=-np.inf, high=np.inf, shape=(45,), dtype=np.float32)
+        self._num_observation = 45
+
+        self._hx = np.array(cfg.height_points_x, dtype=np.float32)
+        self._hy = np.array(cfg.height_points_y, dtype=np.float32)
+        self._num_rows = cfg.terrain_rows
+        self._num_cols = cfg.terrain_cols
+        self._cell_size = cfg.cell_size
+        self._border = cfg.border_size
+        self._init_done = False
+        self._all_terrain_types = np.tile(
+            np.arange(self._num_cols), num_envs // self._num_cols + 1)[:num_envs]
+        np.random.shuffle(self._all_terrain_types)
+        self._hf_cache = None
+        self._hm_cache = None
+
+    # ── 命令：全范围，无课程（匹配上游）──
+
+    def resample_commands(self, num_envs: int) -> np.ndarray:
+        lim = np.array(self.cfg.commands.vel_limit, dtype=np.float32)
+        cmds = np.random.uniform(lim[0], lim[1], size=(num_envs, 3)).astype(np.float32)
+        small = np.linalg.norm(cmds[:, :2], axis=1) < 0.2
+        cmds[small, :2] = 0.0
+        return cmds
+
+    # ── 动作裁剪 + 力矩计算 ──
+
+    def apply_action(self, actions, state):
+        """裁剪动作防止奖励计算溢出（上游 clip_actions=100）。"""
+        actions = np.clip(actions, -100.0, 100.0)
+        return super().apply_action(actions, state)
+
+    def _compute_torques(self, actions, data):
+        """PD 控制器 + 域随机化。力矩裁剪防溢出。"""
+        state = getattr(self, '_state', None)
+        if state is not None:
+            motor_strength = state.info.get("motor_strength",
+                np.ones(self._num_envs, dtype=np.float32))
+            kp_factor = state.info.get("kp_factor",
+                np.ones(self._num_envs, dtype=np.float32))
+            kd_factor = state.info.get("kd_factor",
+                np.ones(self._num_envs, dtype=np.float32))
+        else:
+            motor_strength = np.ones(1, dtype=np.float32)
+            kp_factor = np.ones(1, dtype=np.float32)
+            kd_factor = np.ones(1, dtype=np.float32)
+
+        actions_scaled = actions * self.cfg.control_config.action_scale * motor_strength[:, np.newaxis]
+        torques = (self.kps * kp_factor[:, np.newaxis]) * (
+            actions_scaled + self.default_angles - self.get_dof_pos(data)
+        ) - (self.kds * kd_factor[:, np.newaxis]) * self.get_dof_vel(data)
+        return np.clip(torques, -80.0, 80.0)
+
+    # ── 观测 ──
+
+    def _get_obs(self, data: mtx.SceneData, info: dict) -> np.ndarray:
+        """45-dim: ang_vel + gravity + commands + joint_pos + joint_vel + actions。"""
+        gyro = self.get_gyro(data)
+        pose = self._body.get_pose(data)
+        base_quat = pose[:, 3:7]
+        gravity = quaternion.rotate_inverse(base_quat, self.gravity_vec)
+        commands = info["commands"] * self.commands_scale
+        joint_pos = (self.get_dof_pos(data) - self.default_angles) * self.cfg.normalization.dof_pos
+        joint_vel = self.get_dof_vel(data) * self.cfg.normalization.dof_vel
+        noisy_gyro = gyro * self.cfg.normalization.ang_vel
+        actions = info["current_actions"]
+        obs = np.concatenate([
+            noisy_gyro, gravity, commands, joint_pos, joint_vel, actions
+        ], axis=-1)
+        # 观测噪声
+        nc = self.cfg.noise_config
+        noise_scale_vec = np.zeros(45, dtype=np.float32)
+        noise_scale_vec[0:3] = nc.scale_gyro
+        noise_scale_vec[3:6] = nc.scale_gravity
+        noise_scale_vec[6:9] = 0.0
+        noise_scale_vec[9:21] = nc.scale_joint_angle
+        noise_scale_vec[21:33] = nc.scale_joint_vel
+        noise_scale_vec[33:45] = 0.0
+        noise = (2.0 * np.random.rand(*obs.shape).astype(np.float32) - 1.0) * nc.level
+        obs = obs + noise * noise_scale_vec[np.newaxis, :]
+        return obs
+
+    # ── 高度测量 ──
+
+    def _get_heights(self, data: mtx.SceneData) -> np.ndarray:
+        n = data.shape[0]
+        nx, ny = len(self._hx), len(self._hy)
+        pose = self._body.get_pose(data)
+        base_pos = pose[:, 0:3]
+        yaw = quaternion.get_yaw(pose[:, 3:7])
+        cos_yaw, sin_yaw = np.cos(yaw), np.sin(yaw)
+        try:
+            if self._hm_cache is None:
+                hf = self._model.get_hfield(0)
+                self._hf_cache = hf
+                hm = hf.height_matrix
+                self._hm_cache = hm
+                nr, nc = hm.shape
+                b = hf.bound
+                self._h_xmin, self._h_ymin = b[0], b[1]
+                self._h_xmax, self._h_ymax = b[3], b[4]
+                self._h_nr, self._h_nc = nr, nc
+                cfg_base = getattr(self.cfg, "hfield_z_base", None)
+                self._h_zbase = float(cfg_base) if cfg_base is not None else float(-hm[0, 0])
+            hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
+            xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
+            z_base = self._h_zbase
+            w, h = xmax - xmin, ymax - ymin
+        except Exception:
+            return np.zeros((n, nx * ny), dtype=np.float32)
+
+        heights = np.zeros((n, nx * ny), dtype=np.float32)
+        idx = 0
+        for iy in range(ny):
+            gy = self._hy[iy]
+            for ix in range(nx):
+                gx = self._hx[ix]
+                wx = base_pos[:, 0] + cos_yaw * gx - sin_yaw * gy
+                wy = base_pos[:, 1] + sin_yaw * gx + cos_yaw * gy
+                col_f = (wx - xmin) / max(w, 1e-6) * (nc - 1)
+                row_f = (ymax - wy) / max(h, 1e-6) * (nr - 1)
+                col = np.clip(np.nan_to_num(col_f, nan=0).astype(np.int32), 0, nc - 1)
+                row = np.clip(np.nan_to_num(row_f, nan=0).astype(np.int32), 0, nr - 1)
+                heights[:, idx] = hm[row, col] + z_base
+                idx += 1
+        return heights
+
+    def _sample_terrain_height(self, xy: np.ndarray, radius: float = 0.0) -> np.ndarray:
+        n = xy.shape[0]
+        try:
+            hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
+            xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
+            z_base = self._h_zbase
+            w, h = xmax - xmin, ymax - ymin
+        except Exception:
+            return np.zeros(n, dtype=np.float32)
+        col = np.clip(np.nan_to_num((xy[:, 0] - xmin) / max(w, 1e-6) * (nc - 1), nan=0).astype(np.int32), 0, nc - 1)
+        row = np.clip(np.nan_to_num((ymax - xy[:, 1]) / max(h, 1e-6) * (nr - 1), nan=0).astype(np.int32), 0, nr - 1)
+        if radius <= 0:
+            return (hm[row, col] + z_base).astype(np.float32)
+        dc = max(1, int(radius / max(w, 1e-6) * (nc - 1)))
+        dr = max(1, int(radius / max(h, 1e-6) * (nr - 1)))
+        out = np.empty(n, dtype=np.float32)
+        for i in range(n):
+            r, c = int(row[i]), int(col[i])
+            out[i] = hm[max(0, r - dr):r + dr + 1, max(0, c - dc):c + dc + 1].max()
+        return out + z_base
+
+    # ── 足部接触力 ──
+
+    def _read_contact_forces(self, data: mtx.SceneData) -> np.ndarray:
+        """读取四足接触力（body frame）→ (N, 12)。"""
+        pose = self._body.get_pose(data)
+        base_quat = pose[:, 3:7]
+        forces = []
+        for foot in ["FR", "FL", "RR", "RL"]:
+            v = self._model.get_sensor_value(foot + "_foot_contact", data)
+            v_body = quaternion.rotate_inverse(base_quat, v)
+            forces.append(v_body)
+        return np.concatenate(forces, axis=1)
+
+    def _get_privileged_obs(self, data: mtx.SceneData, obs: np.ndarray) -> np.ndarray:
+        base_vel = self.get_local_linvel(data) * self.cfg.normalization.lin_vel
+        heights = self._get_heights(data)
+        contact_forces = self._read_contact_forces(data)
+        return np.concatenate([obs, base_vel, contact_forces, heights], axis=-1)
+
+    # ── update_observation ──
+
+    def update_observation(self, state):
+        data = state.data
+        obs = self._get_obs(data, state.info)
+
+        # 更新历史缓冲区
+        old_history = state.info.get("obs_history",
+            np.zeros((self._num_envs, self._cfg.num_history, 45), dtype=np.float32))
+        new_history = np.concatenate([old_history[:, 1:, :], obs[:, np.newaxis, :]], axis=1)
+        state.info["obs_history"] = new_history
+
+        current_step = state.info.get("steps", np.zeros(self._num_envs, dtype=np.int32))
+
+        # 中期 command 重采样（每 10 秒）
+        resample_steps = int(10.0 / self.cfg.ctrl_dt)
+        do_resample = (current_step % resample_steps) == 0
+        if do_resample.any():
+            n_resample = int(do_resample.sum())
+            state.info["commands"][do_resample] = self.resample_commands(n_resample)
+
+        # 周期性域随机化（每 4s）
+        rand_interval = int(4.0 / self.cfg.ctrl_dt)
+        last_rand = state.info.get("last_rand_step", np.zeros(self._num_envs, dtype=np.int32))
+        do_rand = (current_step - last_rand) >= rand_interval
+        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["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)
+            state.info["last_rand_step"][do_rand] = current_step[do_rand]
+
+        # 特权观测
+        state.info["privileged_obs"] = self._get_privileged_obs(data, obs)
+        state.info["base_vel"] = self.get_local_linvel(data)
+
+        # 足部接触
+        cquerys = self._model.get_contact_query(data)
+        state.info["contacts"] = cquerys.is_colliding(self.foot_check).reshape(
+            (self._num_envs, self.foot_check_num))
+        state.info["feet_air_time"] = self.update_feet_air_time(state.info)
+
+        # 累计命令距离和跟踪
+        cmd_speed = np.linalg.norm(state.info["commands"][:, :2], axis=1)
+        state.info["ep_cmd_distance"] = state.info.get("ep_cmd_distance",
+            np.zeros(self._num_envs, dtype=np.float32)) + cmd_speed * self.cfg.ctrl_dt
+        state.info["ep_steps"] = state.info.get("ep_steps", np.zeros(self._num_envs, dtype=np.int32)) + 1
+        local_vel = self.get_local_linvel(data)[:, :2]
+        cmd_vel = state.info["commands"][:, :2]
+        vel_error = np.sum(np.square(cmd_vel - local_vel), axis=1)
+        tracking = np.exp(-vel_error / self.cfg.reward_config.tracking_sigma)
+        state.info["ep_tracking_sum"] = state.info.get("ep_tracking_sum",
+            np.zeros(self._num_envs, dtype=np.float32)) + tracking
+
+        return state.replace(obs=obs)
+
+    # ── 出生点 ──
+
+    def _make_origins(self, levels, indices):
+        n = len(levels)
+        half_x = self._border + self._num_cols * self._cell_size / 2.0
+        half_y = self._border + self._num_rows * self._cell_size / 2.0
+        origins = np.zeros((n, 2), dtype=np.float32)
+        for i in range(n):
+            row = levels[i]
+            col = self._all_terrain_types[int(indices[i]) % len(self._all_terrain_types)]
+            cx = -half_x + self._border + col * self._cell_size + self._cell_size / 2
+            cy = half_y - self._border - row * self._cell_size - self._cell_size / 2
+            origins[i, 0] = cx
+            origins[i, 1] = cy
+        return origins
+
+    # ── Reset ──
+
+    def reset(self, data) -> tuple[np.ndarray, dict]:
+        num_reset = data.shape[0]
+        state = getattr(self, '_state', None)
+        if state is not None and hasattr(state, 'done'):
+            done = state.done
+            done_idx = np.where(done)[0]
+        else:
+            done_idx = np.arange(num_reset)
+
+        # 游戏启发式地形课程（per-env，与上游一致）
+        if not hasattr(self, '_terrain_origins'):
+            all_levels = np.repeat(np.arange(self._num_rows), self._num_cols)
+            all_indices = np.tile(np.arange(self._num_cols), self._num_rows)
+            all_origins = self._make_origins(all_levels, all_indices)
+            self._terrain_origins = all_origins.reshape(self._num_rows, self._num_cols, 2)
+            self._max_init_level = 5
+
+        if num_reset > 0 and self._init_done and state is not None and hasattr(state, 'info'):
+            old_info = state.info
+            old_origins = old_info.get("env_origins", np.zeros((self._num_envs, 2), dtype=np.float32))[done]
+            old_levels = old_info.get("terrain_level", np.zeros(self._num_envs, dtype=np.int32))[done]
+            old_commands = old_info.get("commands", np.zeros((self._num_envs, 3), dtype=np.float32))[done]
+            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 > (self._cell_size / 2.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
+            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)
+        else:
+            new_levels = np.random.randint(0, self._max_init_level + 1, size=num_reset, dtype=np.int32)
+        self._init_done = True
+
+        _force = getattr(self, "_force_level", None)
+        if _force is not None:
+            new_levels = np.full(num_reset, int(_force), dtype=np.int32)
+
+        if "scene_stairs_terrain" in self.cfg.model_file:
+            height_list = np.array([-1.0, 0.5, 1.5], dtype=np.float32)
+            offset_h = [[2, 0, 2, 1, 1], [2, 2, 1, 0, 0], [1, 1, 2, 1, 2],
+                        [0, 1, 0, 2, 0], [0, 1, 1, 0, 2]]
+            offsets = []
+            for i in range(5):
+                for j in range(5):
+                    offsets.append([(i-2)*8.0, (j-2)*8.0, height_list[offset_h[j][i]]])
+            offset_arr = np.array(offsets, dtype=np.float32)
+            idx = np.random.choice(len(offsets), size=num_reset)
+            new_origins = offset_arr[idx, :2]
+            self._stairs_z_offsets = offset_arr[idx, 2]
+        else:
+            indices = done_idx
+            if hasattr(self, '_terrain_origins'):
+                new_origins = np.zeros((num_reset, 2), dtype=np.float32)
+                for i in range(num_reset):
+                    row = int(new_levels[i])
+                    col = self._all_terrain_types[int(indices[i]) % len(self._all_terrain_types)]
+                    new_origins[i] = self._terrain_origins[row, col]
+            else:
+                new_origins = self._make_origins(new_levels, indices)
+            self._stairs_z_offsets = np.zeros(num_reset, dtype=np.float32)
+
+        data.reset(self._model)
+
+        init_dof_pos = np.tile(self._init_dof_pos, (num_reset, 1))
+        init_dof_pos[:, 0] = new_origins[:, 0]
+        init_dof_pos[:, 1] = new_origins[:, 1]
+        if hasattr(self, '_stairs_z_offsets') and self._stairs_z_offsets.any():
+            init_dof_pos[:, 2] = self._stairs_z_offsets
+        else:
+            terrain_z = self._sample_terrain_height(new_origins, radius=0.35)
+            spawn_abs = getattr(self, "_spawn_absolute", None)
+            if spawn_abs is not None:
+                init_dof_pos[:, 2] = float(spawn_abs)
+            else:
+                init_dof_pos[:, 2] = terrain_z + getattr(self, "_spawn_clearance", 0.45)
+        yaw = np.random.uniform(-np.pi, np.pi, size=num_reset)
+        init_dof_pos[:, 3] = 0.0
+        init_dof_pos[:, 4] = 0.0
+        init_dof_pos[:, 5] = np.sin(yaw / 2)
+        init_dof_pos[:, 6] = np.cos(yaw / 2)
+
+        init_dof_vel = np.tile(self._init_dof_vel, (num_reset, 1))
+        data.set_dof_vel(init_dof_vel)
+        data.set_dof_pos(init_dof_pos, self._model)
+        self._model.forward_kinematic(data)
+
+        info = {
+            "current_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
+            "last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
+            "last_last_actions": np.zeros((num_reset, self._num_action), dtype=np.float32),
+            "commands": self.resample_commands(num_reset),
+            "last_dof_vel": np.zeros((num_reset, self._num_action), dtype=np.float32),
+            "feet_air_time": np.zeros((num_reset, self.foot_check_num), dtype=np.float32),
+            "contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
+            "last_contacts": np.zeros((num_reset, self.foot_check_num), dtype=np.bool),
+            "motor_strength": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
+            "kp_factor": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
+            "kd_factor": np.random.uniform(0.9, 1.1, size=num_reset).astype(np.float32),
+            "friction_coeff": np.random.uniform(0.2, 1.25, size=num_reset).astype(np.float32),
+            "added_mass": np.random.uniform(-1.0, 2.0, size=num_reset).astype(np.float32),
+            "com_displacement": np.random.uniform(-0.05, 0.05, size=(num_reset, 3)).astype(np.float32),
+            "last_rand_step": np.zeros(num_reset, dtype=np.int32),
+            "ep_cmd_distance": np.zeros(num_reset, dtype=np.float32),
+            "ep_steps": np.zeros(num_reset, dtype=np.int32),
+            "ep_tracking_sum": np.zeros(num_reset, dtype=np.float32),
+        }
+        obs = self._get_obs(data, info)
+
+        hist = np.zeros((num_reset, self._cfg.num_history, 45), dtype=np.float32)
+        hist[:, -1, :] = obs
+        info["obs_history"] = hist
+
+        info["privileged_obs"] = self._get_privileged_obs(data, obs)
+        info["base_vel"] = self.get_local_linvel(data)
+        info["terrain_level"] = new_levels
+        info["env_origins"] = new_origins
+
+        # 清除 ep_ 奖励累计器（跳过 ep_report 它是 dict）
+        if state is not None:
+            for k in list(state.info.keys()):
+                if k.startswith("ep_") and k != "ep_report" and k not in info:
+                    info[k] = np.zeros(num_reset, dtype=np.float32)
+
+        return obs, info
+
+    # ── 奖励（与上游对齐，update_reward 中 × dt）──
+
+    def _get_reward(self, data: mtx.SceneData, info: dict) -> dict[str, np.ndarray]:
+        """DreamWaQ 奖励项——与上游 Manaro-Alpha 对齐。"""
+        commands = info["commands"]
+        return {
+            "tracking_lin_vel": self._reward_tracking_lin_vel(data, commands),
+            "tracking_ang_vel": self._reward_tracking_ang_vel(data, commands),
+            "lin_vel_z": self._reward_lin_vel_z(data),
+            "ang_vel_xy": self._reward_ang_vel_xy(data),
+            "orientation": self._reward_orientation(data),
+            "dof_acc": self._reward_dof_acc(data, info),
+            "base_height": self._reward_base_height(data),
+            "feet_air_time": self._reward_feet_air_time(commands, info),
+            "action_rate": self._reward_action_rate(info),
+            "joint_power": self._reward_joint_power(data),
+            "smoothness": self._reward_smoothness(info),
+            "power_distribution": self._reward_power_distribution(data),
+            # stand_still 在上游被注释掉
+        }
+
+    def _reward_feet_air_time(self, commands, info):
+        """足部腾空时间奖励——与上游 legged_robot.py 公式一致。"""
+        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 总和。"""
+        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.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
+
+        # 对即将结束的 env 存 ep_report
+        steps = state.info.get("steps", np.zeros(self._num_envs, dtype=np.int32))
+        will_end = state.terminated | (steps >= self._cfg.max_episode_steps - 1)
+        if will_end.any():
+            ep_report = state.info.get("ep_report", {})
+            for k in scaled_terms:
+                ek = f"ep_{k}"
+                vals = state.info.get(ek, np.zeros(self._num_envs, dtype=np.float32))[will_end]
+                ep_report[f"rew_{k}"] = float(np.mean(vals))
+            tl = state.info.get("terrain_level", np.zeros(self._num_envs, dtype=np.int32))[will_end]
+            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
+
+    # ── 额外奖励函数 ──
+
+    def _reward_base_height(self, data):
+        """惩罚偏离目标基础高度——与上游 Go1RoughCfg.base_height_target=0.30 一致。"""
+        pose = self._body.get_pose(data)
+        base_z = pose[:, 2]
+        heights = self._get_heights(data)
+        ground_level = np.mean(heights, axis=1)
+        target = 0.30
+        return np.square(base_z - ground_level - target)
+
+    def _reward_joint_power(self, data):
+        torque = np.clip(data.actuator_ctrls, -100, 100)
+        vel = np.clip(self.get_dof_vel(data), -100, 100)
+        return np.sum(np.abs(torque * vel), axis=1)
+
+    def _reward_smoothness(self, info):
+        scale = self.cfg.control_config.action_scale
+        da = self.default_angles
+        a0 = info["current_actions"] * scale + da
+        a1 = info["last_actions"] * scale + da
+        a2 = info.get("last_last_actions", info["last_actions"]) * scale + da
+        diff = np.square(np.clip(a0 - 2.0 * a1 + a2, -1e4, 1e4))
+        diff = diff * (info["last_actions"] != 0)
+        diff = diff * (info.get("last_last_actions", info["last_actions"]) != 0)
+        return np.sum(diff, axis=1)
+
+    def _reward_power_distribution(self, data):
+        torque = np.clip(data.actuator_ctrls, -100, 100)
+        vel = np.clip(self.get_dof_vel(data), -100, 100)
+        power = torque * vel
+        return np.var(np.abs(power), axis=1)

motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/dreamwaq_terrain.png

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5fe0435e385f736b46a910b53d30135d7f6280c0d8daa2bd2f64b7df7962d998
+size 1439313

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

@@ -0,0 +1,25 @@
+<mujoco model="go1 dreamwaq flat scene">
+  <include file="go1_motor_actuator.xml" />
+  <include file="materials.xml" />
+  <statistic center="0 0 0.1" extent="0.8" meansize="0.04" />
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
+    <rgba haze="0.15 0.25 0.35 1" />
+    <global azimuth="120" elevation="-20" />
+    <map force="0.01" />
+    <scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
+    <quality shadowsize="8192" />
+  </visual>
+  <worldbody>
+    <light pos="0 0 1.5" dir="0 0 -1" directional="true" />
+    <geom name="floor" pos="0 0 0" size="0 0 0.01" type="plane"
+          material="motphys-ground" contype="1" conaffinity="0" priority="1"
+          friction="0.6" condim="3" />
+  </worldbody>
+  <sensor>
+    <contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
+    <contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
+    <contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
+    <contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
+  </sensor>
+</mujoco>

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

@@ -0,0 +1,34 @@
+<mujoco model="go1 dreamwaq terrain scene">
+  <include file="go1_motor_actuator.xml" />
+  <include file="materials.xml" />
+  <statistic center="0 0 0.2" extent="5" meansize="0.04" />
+
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
+    <rgba haze="0.15 0.25 0.35 1" />
+    <global azimuth="120" elevation="-20" />
+    <map force="0.01" />
+    <scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
+    <quality shadowsize="8192" />
+  </visual>
+
+  <asset>
+    <hfield name="dreamwaq_terrain"
+            file="assets/dreamwaq_terrain.png"
+            size="85.0 45.0 1.680 0.001" />
+  </asset>
+
+  <worldbody>
+    <light pos="0 0 4" dir="0 0 -1" directional="true" />
+    <geom name="floor" pos="0 0 0" type="hfield" hfield="dreamwaq_terrain"
+          material="motphys-ground" contype="1" conaffinity="0"
+          priority="1" friction="0.6" />
+  </worldbody>
+
+  <sensor>
+    <contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
+    <contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
+    <contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
+    <contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
+  </sensor>
+</mujoco>

motrix_rl/src/motrix_rl/rslrl/cfg.py

@@ -74,6 +74,10 @@ class RslRlPpoAlgorithmCfg:
     max_grad_norm: float = 1.0
     normalize_advantage_per_mini_batch: bool = False
     rnd_cfg: dict | None = None
+    # DreamWaQ CENet 参数
+    vae_beta: float = 1.0       # VAE KL 散度权重
+    cenet_in_dim: int = 225     # 观测历史维度 (num_history × obs_dim)
+    cenet_out_dim: int = 19     # CENet code 维度 (vel_est 3 + latent 16)
     symmetry_cfg: dict | None = None
 
 

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

@@ -0,0 +1,2 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+"""RSLRL 自定义模型模块。"""

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

@@ -0,0 +1,162 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+"""CENet Actor 模型——DreamWaQ 的 VAE 编码器 + MLP Actor。
+
+架构：
+  CENetVAE:  obs_history(225) → encoder[128,64] → latent(16) + vel_est(3) = code(19)
+             code(19) → decoder[64,128] → 重建 obs(45)
+  CENetActorModel (继承 MLPModel):
+              code(19) + obs(45) = 64 → MLP[512,256,128] → action(12)
+
+用法：
+  class_name = "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel"
+  obs_groups = {"actor": ["policy", "obs_history"], "critic": ["privileged_obs"]}
+"""
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+
+from rsl_rl.models.mlp_model import MLPModel
+from rsl_rl.modules import EmpiricalNormalization, HiddenState
+
+
+# ═══ CENet VAE ═══
+
+class CENetVAE(nn.Module):
+    """CENet VAE 编码器-解码器。
+
+    obs_history(225) → encoder → latent(16) + vel_est(3) = code(19)
+    code(19) → decoder → 重建 obs(45)
+    """
+
+    def __init__(self, cenet_in_dim: int = 225, cenet_out_dim: int = 19,
+                 activation: str = "elu"):
+        super().__init__()
+        act = _get_activation(activation)
+        self.cenet_in_dim = cenet_in_dim
+        self.cenet_out_dim = cenet_out_dim
+
+        # 编码器：225 → 128 → 64
+        self.encoder = nn.Sequential(
+            nn.Linear(cenet_in_dim, 128), act,
+            nn.Linear(128, 64), act,
+        )
+        # 潜变量头：64 → 16 (mean + logvar)
+        self.encode_mean_latent = nn.Linear(64, cenet_out_dim - 3)
+        self.encode_logvar_latent = nn.Linear(64, cenet_out_dim - 3)
+        # 速度估计头：64 → 3 (mean + logvar)
+        self.encode_mean_vel = nn.Linear(64, 3)
+        self.encode_logvar_vel = nn.Linear(64, 3)
+        # 解码器：19 → 64 → 128 → 45
+        self.decoder = nn.Sequential(
+            nn.Linear(cenet_out_dim, 64), act,
+            nn.Linear(64, 128), act,
+            nn.Linear(128, 45),
+        )
+
+    def reparameterise(self, mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        """重参数化技巧：从 N(mean, exp(logvar/2)) 采样。"""
+        std = torch.exp(logvar * 0.5)
+        eps = torch.randn_like(std)
+        return mean + std * eps
+
+    def forward(self, obs_history: torch.Tensor):
+        """前向传播。
+
+        Returns:
+            code:       (N, 19)  潜变量 [vel_sample(3) + latent_sample(16)]
+            code_vel:   (N, 3)   速度估计采样
+            decode:     (N, 45)  重建观测
+            mean_vel:   (N, 3)   速度估计均值
+            logvar_vel: (N, 3)   速度估计对数方差
+            mean_latent:   (N, 16)  潜变量均值
+            logvar_latent: (N, 16)  潜变量对数方差
+        """
+        h = self.encoder(obs_history)
+        mean_latent = self.encode_mean_latent(h)
+        logvar_latent = self.encode_logvar_latent(h)
+        mean_vel = self.encode_mean_vel(h)
+        logvar_vel = self.encode_logvar_vel(h)
+        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)
+        decode = self.decoder(code)
+        return code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent
+
+    def deterministic_code(self, obs_history: torch.Tensor) -> torch.Tensor:
+        """推理模式：使用均值而非采样，产生确定性 code(19)。"""
+        h = self.encoder(obs_history)
+        mean_latent = self.encode_mean_latent(h)
+        mean_vel = self.encode_mean_vel(h)
+        return torch.cat((mean_vel, mean_latent), dim=-1)
+
+
+# ═══ CENet Actor Model ═══
+
+class CENetActorModel(MLPModel):
+    """CENet Actor：继承 MLPModel，在 get_latent() 中注入 VAE code。
+
+    class_name = "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel"
+    """
+
+    def __init__(self, obs: TensorDict, obs_groups: dict[str, list[str]],
+                 obs_set: str, output_dim: int,
+                 cenet_in_dim: int = 225, cenet_out_dim: int = 19,
+                 activation: str = "elu", **kwargs):
+        # 必须在 super().__init__ 之前设置，因为 _get_latent_dim() 会被父类构造函数调用
+        self._history_dim = cenet_in_dim
+        self._code_dim = cenet_out_dim
+        # 禁用观测归一化（VAE 输出已是归一化后的 code，维度也不匹配）
+        kwargs["obs_normalization"] = False
+
+        super().__init__(obs, obs_groups, obs_set, output_dim, **kwargs)
+
+        # 在 nn.Module.__init__ 之后创建 VAE 子模块
+        self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation)
+        self._last_cenet_output = None
+
+    def _update_distribution(self, obs: torch.Tensor) -> None:
+        """覆盖父类 — 确保 std 始终为正，防止 NaN。"""
+        super()._update_distribution(obs)
+        # 如果 std 因数值问题变负，clamp 到最小值
+        if self.stochastic and not self.state_dependent_std:
+            with torch.no_grad():
+                if self.noise_std_type == "scalar":
+                    self.std.clamp_(min=1e-6)
+                elif self.noise_std_type == "log":
+                    self.log_std.clamp_(min=-20.0, max=10.0)
+
+    def _get_latent_dim(self) -> int:
+        """Actor 实际输入：code(19) + policy(45) = 64。"""
+        return self.obs_dim - self._history_dim + self._code_dim
+
+    def get_latent(self, obs: TensorDict, masks: torch.Tensor | None = None,
+                   hidden_state: HiddenState = None) -> torch.Tensor:
+        """提取观测 → VAE 编码 → 拼接 code + policy → 返回 latent(64)。"""
+        policy_obs = obs["policy"]           # (N, 45)
+        obs_history = obs["obs_history"]     # (N, 225)
+
+        out = self.vae(obs_history)
+        self._last_cenet_output = out
+        code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = out
+
+        latent = torch.cat([code, policy_obs], dim=-1)  # (N, 64)
+        return latent
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        """CENetActor 使用 obs_normalization=False，此方法为空。"""
+        pass
+
+
+# ═══ 工具函数 ═══
+
+def _get_activation(act_name: str) -> nn.Module:
+    """解析激活函数名称。"""
+    _map = {
+        "elu": nn.ELU, "selu": nn.SELU, "relu": nn.ReLU,
+        "lrelu": nn.LeakyReLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid,
+    }
+    if act_name in _map:
+        return _map[act_name]()
+
+    raise ValueError(f"未知激活函数: {act_name}。可选: {list(_map.keys())}")

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

@@ -0,0 +1,205 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+"""DreamWaQ PPO——标准 PPO + CENet VAE 自编码器损失。
+
+在 PPO.update() 的每个 mini-batch 中，额外计算：
+  - 速度估计损失：MSE(code_vel, 真实 base_vel)（来自 privileged_obs）
+  - 观测重建损失：MSE(decode, 当前 policy obs)
+  - KL 散度损失：beta * KL(N(mean, var) || N(0, 1))
+
+class_name = "motrix_rl.rslrl.torch.train.dreamwaq_ppo:DreamWaQPPO"
+"""
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+
+from rsl_rl.algorithms import PPO
+from rsl_rl.storage import RolloutStorage
+from rsl_rl.utils import resolve_callable, resolve_obs_groups
+
+
+class DreamWaQPPO(PPO):
+    """PPO + Beta-VAE 自编码器损失——DreamWaQ CENet 训练。"""
+
+    # VAE 损失权重（与上游 beta=1.0 一致）
+    vae_beta: float = 1.0
+
+    @staticmethod
+    def construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> "DreamWaQPPO":
+        """构造 DreamWaQ PPO 算法——创建 CENetActor + MLP Critic。
+
+        与父类 PPO.construct_algorithm 的区别：
+          - actor 使用 CENetActorModel（含 VAE）
+          - obs_groups 中 actor=["policy", "obs_history"], critic=["privileged_obs"]
+        """
+        # 提取 DreamWaQ 特有参数
+        vae_beta = cfg.pop("vae_beta", 1.0)
+
+        # 解析 actor / critic 类
+        actor_class = resolve_callable(cfg["actor"].pop("class_name"))
+        critic_class = resolve_callable(cfg["critic"].pop("class_name"))
+
+        # 解析观测分组
+        obs_groups = resolve_obs_groups(obs, cfg["obs_groups"], ["actor", "critic"])
+
+        # 创建 actor（CENetActorModel）
+        actor = actor_class(obs, obs_groups, "actor", env.num_actions, **cfg["actor"]).to(device)
+
+        # 创建 critic（标准 MLPModel，输入 privileged_obs）
+        critic = critic_class(obs, obs_groups, "critic", 1, **cfg["critic"]).to(device)
+
+        # 初始化 rollout 存储
+        storage = RolloutStorage(
+            "rl", env.num_envs, cfg["num_steps_per_env"], obs, [env.num_actions], device
+        )
+
+        # 提取算法参数（移除 DreamWaQPPO 特有 key，剩余传给父类 PPO.__init__）
+        algo_cfg = dict(cfg["algorithm"])
+        for dw_key in ("class_name", "vae_beta", "cenet_in_dim", "cenet_out_dim"):
+            algo_cfg.pop(dw_key, None)
+
+        # 创建 DreamWaQPPO 实例
+        alg = DreamWaQPPO(actor, critic, storage, device=device, **algo_cfg)
+        alg.vae_beta = vae_beta
+        return alg
+
+    def update(self) -> dict[str, float]:
+        """标准 PPO update + VAE 自编码器损失。"""
+        mean_value_loss = 0.0
+        mean_surrogate_loss = 0.0
+        mean_autoenc_loss = 0.0
+
+        if self.actor.is_recurrent or self.critic.is_recurrent:
+            generator = self.storage.recurrent_mini_batch_generator(
+                self.num_mini_batches, self.num_learning_epochs)
+        else:
+            generator = self.storage.mini_batch_generator(
+                self.num_mini_batches, self.num_learning_epochs)
+
+        for (
+            obs_batch, actions_batch, target_values_batch, advantages_batch,
+            returns_batch, old_actions_log_prob_batch, old_mu_batch,
+            old_sigma_batch, hid_states_batch, masks_batch,
+        ) in generator:
+            # ── 标准 PPO 前向 ──
+            self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
+            self.critic(obs_batch, masks=masks_batch)
+
+            # ── 动作对数概率 ──
+            actions_log_prob_batch = self.actor.get_output_log_prob(actions_batch)
+
+            # ── 价值损失 ──
+            value_batch = target_values_batch
+            if self.use_clipped_value_loss:
+                value_clipped = target_values_batch + (self.critic(obs_batch).detach() - target_values_batch).clamp(
+                    -self.clip_param, self.clip_param)
+                value_losses = torch.square(self.critic(obs_batch) - returns_batch)
+                value_losses_clipped = torch.square(value_clipped - returns_batch)
+                value_loss = torch.max(value_losses, value_losses_clipped).mean()
+            else:
+                value_loss = torch.square(returns_batch - self.critic(obs_batch)).mean()
+
+            # ── 代理损失 ──
+            surrogate_loss = self._compute_surrogate_loss(
+                actions_batch, actions_log_prob_batch,
+                old_actions_log_prob_batch, advantages_batch)
+
+            # ── 熵 ──
+            entropy_batch = self.actor.output_entropy
+            entropy_loss = entropy_batch.mean()
+
+            # ── VAE 自编码器损失 ──
+            autoenc_loss = self._compute_vae_loss(obs_batch)
+
+            # ── 总损失 ──
+            loss = (
+                surrogate_loss
+                + self.value_loss_coef * value_loss
+                - self.entropy_coef * entropy_loss
+                + autoenc_loss
+            )
+
+            # ── 梯度更新 ──
+            self.optimizer.zero_grad()
+            loss.backward()
+            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)
+                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.clamp_(min=1e-6)
+
+            # ── 累计日志 ──
+            mean_value_loss += value_loss.item()
+            mean_surrogate_loss += surrogate_loss.item()
+            mean_autoenc_loss += autoenc_loss.item()
+
+        num_updates = self.num_learning_epochs * self.num_mini_batches
+        mean_value_loss /= num_updates
+        mean_surrogate_loss /= num_updates
+        mean_autoenc_loss /= num_updates
+
+        self.storage.clear()
+
+        return {
+            "value_loss": mean_value_loss,
+            "surrogate_loss": mean_surrogate_loss,
+            "autoenc_loss": mean_autoenc_loss,
+        }
+
+    def _compute_surrogate_loss(
+        self, actions_batch, actions_log_prob_batch,
+        old_actions_log_prob_batch, advantages_batch
+    ) -> torch.Tensor:
+        """计算 PPO 代理损失（从父类 PPO.update() 中提取）。"""
+        ratio = torch.exp(actions_log_prob_batch - old_actions_log_prob_batch)
+        surrogate = -advantages_batch * ratio
+        surrogate_clipped = -advantages_batch * torch.clamp(
+            ratio, 1.0 - self.clip_param, 1.0 + self.clip_param)
+        return torch.max(surrogate, surrogate_clipped).mean()
+
+    def _compute_vae_loss(self, obs_batch: TensorDict) -> torch.Tensor:
+        """计算 CENet VAE 损失。
+
+        obs_batch 包含：
+          - "policy": 当前观测 (N, 45) = 重建目标
+          - "obs_history": 观测历史 (N, 225) = VAE 编码器输入
+          - "privileged_obs": 特权观测 (N, 247), 其中 [45:48] 是 base_vel
+
+        损失组成：
+          1. 速度估计损失：MSE(code_vel, base_vel_gt)
+          2. 观测重建损失：MSE(decoded_obs, policy_obs)
+          3. KL 散度：beta * KL(q(z|history) || N(0,1))
+        """
+        # 从 actor 获取最近一次 CENet 前向输出
+        cenet_out = getattr(self.actor, "_last_cenet_output", None)
+        if cenet_out is None:
+            return torch.tensor(0.0, device=obs_batch.device)
+
+        code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = cenet_out
+
+        # 速度估计目标：privileged_obs 中的 base_vel（索引 45:48）
+        vel_target = obs_batch["privileged_obs"][:, 45:48]
+        # 观测重建目标：当前 policy obs
+        obs_target = obs_batch["policy"]
+
+        mse = nn.functional.mse_loss
+        estimation_loss = mse(code_vel, vel_target)
+        reconstruction_loss = mse(decode, obs_target)
+        # KL 散度：-0.5 * sum(1 + logvar - mean^2 - exp(logvar))
+        # clamp logvar 防止 exp 溢出
+        logvar_latent = torch.clamp(logvar_latent, -20.0, 10.0)
+        kl_loss = -0.5 * torch.sum(
+            1 + logvar_latent - mean_latent.pow(2) - logvar_latent.exp(), dim=-1
+        ).mean()
+
+        autoenc_loss = (
+            estimation_loss + reconstruction_loss + self.vae_beta * kl_loss
+        )
+        # 防止 NaN 传播
+        if torch.isnan(autoenc_loss) or torch.isinf(autoenc_loss):
+            return torch.tensor(0.0, device=obs_batch.device)
+        return autoenc_loss

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

@@ -16,6 +16,7 @@
 """PPO Trainer for RSLRL integration."""
 
 import logging
+import os
 
 import torch
 from rsl_rl.runners import OnPolicyRunner
@@ -48,6 +49,7 @@ class Trainer:
         sim_backend: str = None,
         enable_render: bool = False,
         cfg_override: dict = None,
+        env_cfg_override: dict = None,
     ) -> None:
         """Initialize the RSLRL PPO trainer.
 
@@ -56,6 +58,7 @@ class Trainer:
             sim_backend: Simulation backend to use (e.g., "mujoco", "npcm")
             enable_render: Whether to enable rendering during training
             cfg_override: Optional configuration overrides
+            env_cfg_override: Optional env config overrides passed to make()
         """
         rlcfg = rl_registry.default_rl_cfg(env_name, "rslrl", backend="torch")
         if cfg_override is not None:
@@ -64,23 +67,29 @@ class Trainer:
         self._env_name = env_name
         self._sim_backend = sim_backend
         self._enable_render = enable_render
+        self._env_cfg_override = env_cfg_override
 
-    def train(self) -> None:
+    def train(self, checkpoint: str = None) -> None:
         """Start training the agent.
 
         Creates the environment, wraps it for RSLRL, and runs the training loop.
+
+        Args:
+            checkpoint: Optional path to a checkpoint (.pt) to resume from.
         """
         rlcfg = self._rlcfg
 
         # Create environment
-        env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
+        env = env_registry.make(self._env_name, sim_backend=self._sim_backend,
+                                num_envs=rlcfg.num_envs, env_cfg_override=self._env_cfg_override)
 
         # Set random seed
         if rlcfg.runner.seed is not None:
             torch.manual_seed(rlcfg.runner.seed)
 
-        # Determine device
+        # Determine device（可通过 MOTRIX_DEVICE=cpu 强制 CPU 训练）
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        device_str = os.environ.get("MOTRIX_DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu")
+        device = torch.device(device_str)
         logger.info(f"Using device: {device}")
 
         # Wrap environment for RSLRL
@@ -94,6 +103,11 @@ class Trainer:
             vec_env, rslrl_cfg, log_dir=get_log_dir(self._env_name, rllib="rslrl", agent_name="PPO"), device=device
         )
 
+        # Load checkpoint if specified
+        if checkpoint:
+            runner.load(checkpoint)
+            logger.info(f"Resumed from checkpoint: {checkpoint}")
+
         # Start training
         logger.info(f"Starting training for {self._env_name}")
         logger.info(f"Number of environments: {rlcfg.num_envs}")
@@ -117,14 +131,25 @@ class Trainer:
         rlcfg = self._rlcfg
 
         # Create environment with play_num_envs
-        env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
+        # Enable play_mode if the env config supports it (for multi-terrain random spawn)
+        play_override = {"play_mode": True}
+        if self._env_cfg_override:
+            play_override.update(self._env_cfg_override)
+        try:
+            env = env_registry.make(
+                self._env_name, sim_backend=self._sim_backend,
+                num_envs=rlcfg.play_num_envs, env_cfg_override=play_override,
+            )
+        except ValueError:
+            env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
 
         # Set random seed
         if rlcfg.runner.seed is not None:
             torch.manual_seed(rlcfg.runner.seed)
 
         # Determine device
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        device_str = os.environ.get("MOTRIX_DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu")
+        device = torch.device(device_str)
 
         # Wrap environment for RSLRL
         vec_env = RslrlNpEnvWrap(env, device)

motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py

@@ -83,6 +83,21 @@ class RslrlNpEnvWrap(VecEnv):
         """Return the unwrapped environment (self for this wrapper)."""
         return self
 
+    def _build_obs_dict(self, state) -> dict[str, torch.Tensor]:
+        """将 NpEnvState 中的观测字段组装为 TensorDict 字典。
+
+        支持 env 通过 state.info 传递 obs_history 和 privileged_obs。
+        """
+        obs_dict = {"policy": torch.from_numpy(state.obs).to(self._device)}
+        if "obs_history" in state.info:
+            hist = state.info["obs_history"]  # (N, num_history, obs_dim)
+            obs_dict["obs_history"] = torch.from_numpy(hist).reshape(
+                self._num_envs, -1).to(self._device)
+        if "privileged_obs" in state.info:
+            obs_dict["privileged_obs"] = torch.from_numpy(
+                state.info["privileged_obs"]).to(self._device)
+        return obs_dict
+
     def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:
         # Convert torch actions to numpy
         actions_np = actions.cpu().numpy()
@@ -98,20 +113,24 @@ class RslrlNpEnvWrap(VecEnv):
         self.episode_length_buf[dones_np] = 0
 
         # Convert to torch tensors
-        obs_tensor = torch.from_numpy(state.obs).to(self._device)
         rewards = torch.from_numpy(state.reward).to(self._device)
 
         # Merge terminated and truncated into dones
         dones = torch.from_numpy(state.done.astype(np.float32)).to(self._device)
 
-        # Create TensorDict for observations
+        # 构建多键 TensorDict（policy + obs_history + privileged_obs）
-        obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+        obs = TensorDict(self._build_obs_dict(state),
+                         batch_size=[self._num_envs], device=self._device)
 
         # Build extras dict (RSLRL calls it "extras" not "infos")
         extras = {}
         if "time_outs" in state.info:
             extras["time_outs"] = torch.from_numpy(state.info["time_outs"]).to(self._device)
 
+        # 将 episode 各项奖励传入 TensorBoard
+        if "ep_report" in state.info:
+            extras["episode"] = state.info["ep_report"]
+
         return obs, rewards, dones, extras
 
     def reset(self) -> tuple[TensorDict, dict]:
@@ -128,10 +147,9 @@ class RslrlNpEnvWrap(VecEnv):
         # Reset episode length buffer
         self.episode_length_buf.zero_()
 
-        obs_tensor = torch.from_numpy(state.obs).to(self._device)
+        # 构建多键 TensorDict
-
+        obs = TensorDict(self._build_obs_dict(state),
-        # Create TensorDict for observations
+                         batch_size=[self._num_envs], device=self._device)
-        obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
 
         # Build extras dict
         extras = {}
@@ -139,17 +157,17 @@ class RslrlNpEnvWrap(VecEnv):
         return obs, extras
 
     def get_observations(self) -> TensorDict:
-        """Get current observations without stepping the environment.
+        """获取当前观测（不步进环境）。
 
         Returns:
-            Current observations as TensorDict
+            当前观测的 TensorDict（含 policy, obs_history, privileged_obs）
         """
         if self._state is None:
             obs, _ = self.reset()
             return obs
 
-        obs_tensor = torch.from_numpy(self._state.obs).to(self._device)
+        obs = TensorDict(self._build_obs_dict(self._state),
-        obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+                         batch_size=[self._num_envs], device=self._device)
         return obs
 
     def render(self) -> None:

motrix_rl/src/motrix_rl/tasks/go1.py

@@ -66,6 +66,19 @@ class skrl:
     @dataclass
     class Go1WalkStairsPPO(Go1WalkRoughSkrlPpo): ...
 
+    @rlcfg("go1-stairs-terrain-walk-no-linevel")
+    @dataclass
+    class Go1WalkStairsNoLinvelSkrlPpo(Go1WalkRoughSkrlPpo):
+        """Go1 stairs terrain walk (no linear velocity obs) - SKRL PPO config.
+
+        Uses [512, 256, 128] network from rough terrain config.
+        Increased timesteps for 4-phase curriculum.
+        """
+
+        def __post_init__(self):
+            super().__post_init__()
+            self.runner.trainer.timesteps = 60000
+
 
 class rslrl:
     @rlcfg("go1-flat-terrain-walk")
@@ -93,6 +106,53 @@ class rslrl:
             algo.num_learning_epochs = 5
             algo.num_mini_batches = 3
 
+    @rlcfg("go1-dreamwaq-walk")
+    @dataclass
+    class Go1DreamWaQWalkRslrlPpo(RslrlCfg):
+        """Go1 DreamWaQ walk — CENet VAE + 不对称特权观测。"""
+
+        num_envs: int = 1024  # 上游 4096，CPU/GPU 安全默认
+
+        def __post_init__(self):
+            runner = self.runner
+
+            # Runner 设置（严格对齐上游 LeggedRobotCfgPPO + Go1RoughCfgPPO）
+            runner.seed = 5                     # 上游 seed=5
+            runner.max_iterations = 3000
+            runner.num_steps_per_env = 24
+            runner.experiment_name = "go1_dreamwaq_walk"
+            runner.save_interval = 50
+
+            # 算法：DreamWaQPPO（含 VAE loss）—— 严格对齐上游
+            runner.algorithm.class_name = (
+                "motrix_rl.rslrl.torch.train.dreamwaq_ppo:DreamWaQPPO")
+            runner.algorithm.learning_rate = 1e-3     # 上游 1.e-3
+            runner.algorithm.num_learning_epochs = 5
+            runner.algorithm.num_mini_batches = 4
+            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.gamma = 0.99
+            runner.algorithm.lam = 0.95
+            runner.algorithm.max_grad_norm = 1.0
+            runner.algorithm.vae_beta = 1.0
+
+            # Actor：CENetActorModel（code 替换 obs_history）
+            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
+
+            # Critic：标准 MLPModel，输入 privileged_obs
+            runner.critic.class_name = "MLPModel"
+            runner.critic.hidden_dims = [512, 256, 128]
+
+            # 观测分组：actor 用 policy+history，critic 用 privileged_obs
+            runner.obs_groups = {
+                "actor": ["policy", "obs_history"],
+                "critic": ["privileged_obs"],
+            }
+
     @rlcfg("go1-rough-terrain-walk")
     @dataclass
     class Go1WalkRoughRslrlPpo(Go1WalkFlatRslrlPpo):
@@ -123,3 +183,18 @@ class rslrl:
         def __post_init__(self):
             super().__post_init__()
             self.runner.experiment_name = "go1_stairs_terrain_walk"
+
+    @rlcfg("go1-stairs-terrain-walk-no-linevel")
+    @dataclass
+    class Go1WalkStairsNoLinvelRslrlPpo(Go1WalkRoughRslrlPpo):
+        """Go1 stairs terrain walk (no linear velocity obs) - RSLRL PPO config.
+
+        Uses [512, 256, 128] network from rough terrain config.
+        Increased iterations for 4-phase curriculum.
+        """
+
+        def __post_init__(self):
+            super().__post_init__()
+            self.runner.experiment_name = "go1_stairs_terrain_walk_no_linevel"
+            self.runner.max_iterations = 2000
+            self.runner.experiment_name = "go1_stairs_terrain_walk_no_linevel"

scripts/gen_dreamwaq_terrain.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+"""生成 DreamWaQ 10×20 纯 hfield 地形——OpenCV 绘制。
+
+5 种地形类型 × 10 难度，全部在单张 PNG 高度图中。
+楼梯用 1px riser 近垂直面（HS=0.05 时每像素 5cm）。
+
+用法：
+    uv run python3 scripts/gen_dreamwaq_terrain.py
+"""
+import cv2
+import numpy as np
+import os
+import argparse
+
+# ═══ 参数 ═══
+HS = 0.05             # 水平分辨率 [m/px]
+VS = 0.005            # 垂直分辨率 [m/unit]
+CELL_M = 8.0
+NUM_ROWS = 10
+NUM_COLS = 20
+BORDER_M = 5.0
+PROPORTIONS = [0.1, 0.1, 0.35, 0.35, 0.1]
+CUM = [sum(PROPORTIONS[:i + 1]) for i in range(len(PROPORTIONS))]
+PLATFORM_M = 3.0
+_SLOPE_SCALE = 0.4   # 上游原值（已验证 z_scale 上限远超 0.54）
+
+CELL_PX = int(CELL_M / HS)       # 160
+BORDER_PX = int(BORDER_M / HS)   # 100
+PLATFORM_PX = int(PLATFORM_M / HS)  # 60
+TOT_ROWS_PX = NUM_ROWS * CELL_PX + 2 * BORDER_PX  # 1800
+TOT_COLS_PX = NUM_COLS * CELL_PX + 2 * BORDER_PX  # 3400
+TOTAL_X = TOT_COLS_PX * HS
+TOTAL_Y = TOT_ROWS_PX * HS
+
+
+# ═══ 地形绘制 ═══
+
+def draw_slope(canvas, x0, y0, difficulty, noise=False):
+    """平滑/粗糙斜坡——与上游 pyramid_sloped_terrain 对齐。
+
+    上游逻辑：先建金字塔（中心高→边缘低），再用平台边缘高度 clip 整个 terrain，
+    形成与周围地形齐平的平台（而非硬清零到 0）。
+    """
+    if difficulty <= 0:
+        return
+    slope = difficulty * _SLOPE_SCALE
+    max_h = int(slope * (1.0 / VS) * (CELL_M / 2.0))
+    if max_h <= 0:
+        return
+    cx, cy = CELL_PX // 2, CELL_PX // 2
+    x = np.arange(0, CELL_PX)
+    y = np.arange(0, CELL_PX)
+    xx, yy = np.meshgrid(x, y, sparse=True)
+    xx = (cx - np.abs(cx - xx)) / cx
+    yy = (cy - np.abs(cy - yy)) / cy
+    hf = (max_h * xx.reshape(CELL_PX, 1) * yy.reshape(1, CELL_PX)).astype(np.int32)
+    p2 = PLATFORM_PX // 2
+    # 上游 clip: 取平台边缘高度作为上下界
+    edge_h = int(hf[cx - p2, cy - p2])
+    lo = min(edge_h, 0)
+    hi = max(edge_h, 0)
+    hf = np.clip(hf, lo, hi).astype(np.uint16)
+    if noise:
+        na = int(0.05 / VS)
+        n = np.random.randint(-na, na + 1, (CELL_PX, CELL_PX), dtype=np.int16)
+        # 噪声也只在平台外
+        n[cx - p2:cx + p2, cy - p2:cy + p2] = 0
+        hf = np.clip(hf.astype(np.int32) + n, 0, 65535).astype(np.uint16)
+    canvas[y0:y0 + CELL_PX, x0:x0 + CELL_PX] += hf
+
+
+def draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=False):
+    """金字塔楼梯——OpenCV 同心矩形（近垂直 riser）。
+
+    每级台阶 2px 宽（10cm tread），高度缩放保持 z_scale < 0.54。
+    """
+    if difficulty <= 0:
+        return
+    # 上游公式：step_height = 0.05 + 0.18 * difficulty [m]
+    step_h_m = 0.05 + 0.18 * difficulty
+    step_h = max(1, int(step_h_m / VS))
+    cx = x0 + CELL_PX // 2
+    cy = y0 + CELL_PX // 2
+    p2 = PLATFORM_PX // 2
+
+    # 上游踏面 31cm → 6px (HS=0.05), 最多约 8 级
+    tread_px = max(1, int(0.31 / HS))
+    n_steps = min(8, (CELL_PX // 2 - p2) // tread_px)
+
+    if concave:
+        base_h = step_h * n_steps
+        cv2.rectangle(canvas, (x0, y0), (x0 + CELL_PX, y0 + CELL_PX), int(base_h), -1)
+        for i in range(n_steps + 1):
+            half = p2 + (n_steps - i) * tread_px
+            h = int(base_h - step_h * i)
+            cv2.rectangle(canvas, (cx - half, cy - half), (cx + half, cy + half), h, -1)
+    else:
+        for i in range(n_steps + 1):
+            half = p2 + (n_steps - i) * tread_px
+            h = int(step_h * i)
+            cv2.rectangle(canvas, (cx - half, cy - half), (cx + half, cy + half), h, -1)
+
+
+def draw_obstacles(canvas, x0, y0, difficulty):
+    """离散障碍物（随机矩形块）。"""
+    if difficulty <= 0:
+        return
+    max_h = int((0.05 + 0.2 * difficulty) / VS)
+    if max_h <= 0:
+        return
+    p2 = PLATFORM_PX // 2
+    # 上游: min_size=1.0m, max_size=2.0m, 20 个矩形
+    min_sz = int(1.0 / HS); max_sz = int(2.0 / HS)
+    for _ in range(20):
+        w = np.random.randint(min_sz, max_sz + 1)
+        ln = np.random.randint(min_sz, max_sz + 1)
+        si = np.random.randint(0, CELL_PX - w)
+        sj = np.random.randint(0, CELL_PX - ln)
+        cv2.rectangle(canvas, (x0 + si, y0 + sj),
+                      (x0 + si + w, y0 + sj + ln),
+                      int(np.random.choice([max_h // 2, max_h])), -1)
+    cx, cy = x0 + CELL_PX // 2, y0 + CELL_PX // 2
+    cv2.rectangle(canvas, (cx - p2, cy - p2), (cx + p2, cy + p2), 0, -1)
+
+
+# ═══ 主流程 ═══
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--flat-only", action="store_true")
+    p.add_argument("--max-level", type=int, default=None)
+    args = p.parse_args()
+
+    max_row = NUM_ROWS if args.max_level is None else min(args.max_level + 1, NUM_ROWS)
+    print(f"DreamWaQ 纯 hfield ({max_row}×{NUM_COLS})  {TOT_COLS_PX}×{TOT_ROWS_PX}px")
+
+    canvas = np.zeros((TOT_ROWS_PX, TOT_COLS_PX), dtype=np.uint16)
+
+    for row in range(max_row):
+        difficulty = row / NUM_ROWS
+        for col in range(NUM_COLS):
+            if args.flat_only or difficulty == 0:
+                continue
+            x0 = BORDER_PX + col * CELL_PX
+            y0 = BORDER_PX + row * CELL_PX
+            choice = col / NUM_COLS + 0.001
+            if choice < CUM[0]:
+                draw_slope(canvas, x0, y0, difficulty)
+            elif choice < CUM[1]:
+                draw_slope(canvas, x0, y0, difficulty, noise=True)
+            elif choice < CUM[2]:
+                draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=True)
+            elif choice < CUM[3]:
+                draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=False)
+            else:
+                draw_obstacles(canvas, x0, y0, difficulty)
+
+    hf_m = canvas.astype(np.float32) * VS
+    z_min, z_max = float(hf_m.min()), float(hf_m.max())
+    z_range = max(z_max - z_min, 0.001)
+    print(f"  高度范围: [{z_min:.3f}, {z_max:.3f}]m  z_scale={z_range:.3f}")
+
+    if z_range > 0.54:
+        print(f"  ⚠ z_scale={z_range:.3f} > 0.54!")
+
+    out_d = os.path.join(os.path.dirname(__file__), "..",
+                         "motrix_envs", "src", "motrix_envs",
+                         "locomotion", "go1", "xmls", "assets")
+    os.makedirs(out_d, exist_ok=True)
+    png = ((hf_m - z_min) / z_range * 65535.0).astype(np.uint16)
+    cv2.imwrite(os.path.join(out_d, "dreamwaq_terrain.png"), png)
+
+    # XML
+    xml = f"""<mujoco model="go1 dreamwaq terrain scene">
+  <include file="go1_motor_actuator.xml" />
+  <include file="materials.xml" />
+  <statistic center="0 0 0.2" extent="5" meansize="0.04" />
+
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
+    <rgba haze="0.15 0.25 0.35 1" />
+    <global azimuth="120" elevation="-20" />
+    <map force="0.01" />
+    <scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
+    <quality shadowsize="8192" />
+  </visual>
+
+  <asset>
+    <hfield name="dreamwaq_terrain"
+            file="assets/dreamwaq_terrain.png"
+            size="{TOTAL_X / 2:.1f} {TOTAL_Y / 2:.1f} {z_range:.3f} {max(z_min, 0.001):.3f}" />
+  </asset>
+
+  <worldbody>
+    <light pos="0 0 4" dir="0 0 -1" directional="true" />
+    <geom name="floor" pos="0 0 0" type="hfield" hfield="dreamwaq_terrain"
+          material="motphys-ground" contype="1" conaffinity="0"
+          priority="1" friction="0.6" />
+  </worldbody>
+
+  <sensor>
+    <contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
+    <contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
+    <contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
+    <contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
+  </sensor>
+</mujoco>
+"""
+    xml_dir = os.path.join(os.path.dirname(__file__), "..",
+                           "motrix_envs", "src", "motrix_envs",
+                           "locomotion", "go1", "xmls")
+    with open(os.path.join(xml_dir, "scene_dreamwaq_terrain.xml"), "w") as f:
+        f.write(xml)
+
+    half_x = TOTAL_X / 2
+    half_y = TOTAL_Y / 2
+    print(f"  XML: size=\"{half_x:.1f} {half_y:.1f} {z_range:.3f} {max(z_min, 0.001):.3f}\"")
+    print(f"  楼梯: 1px tread (5cm), 1px riser → 近垂直面")
+
+
+if __name__ == "__main__":
+    np.random.seed(42)
+    main()

scripts/test_mesh_collision.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""最小 mesh 碰撞测试——验证 MotrixSim 的 OBJ mesh 是否支持碰撞。"""
+
+import os, numpy as np
+
+xml_dir = '/home/8x54zj-m/MotrixLab/motrix_envs/src/motrix_envs/locomotion/go1/xmls'
+assets_dir = os.path.join(xml_dir, 'assets', 'tmp_test')
+os.makedirs(assets_dir, exist_ok=True)
+obj_path = os.path.join(assets_dir, 'test_box.obj')
+
+# 封闭 box OBJ (1x1x0.1m)，带法线
+with open(obj_path, 'w') as f:
+    f.write("""# closed box
+v -0.5 -0.5 0.0
+v 0.5 -0.5 0.0
+v 0.5 0.5 0.0
+v -0.5 0.5 0.0
+v -0.5 -0.5 0.1
+v 0.5 -0.5 0.1
+v 0.5 0.5 0.1
+v -0.5 0.5 0.1
+f 1 3 2
+f 1 4 3
+f 5 6 7
+f 5 7 8
+f 1 5 6
+f 1 6 2
+f 2 6 7
+f 2 7 3
+f 3 7 8
+f 3 8 4
+f 4 8 5
+f 4 5 1
+""")
+
+# 生成测试 XML
+xml_path = os.path.join(xml_dir, 'scene_test_mesh.xml')
+with open(xml_path, 'w') as f:
+    f.write("""<mujoco model="test mesh">
+  <include file="go1_motor_actuator.xml" />
+  <include file="materials.xml" />
+  <statistic center="0 0 0.3" extent="1" />
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
+  </visual>
+  <asset>
+    <mesh name="test_box" file="assets/tmp_test/test_box.obj" />
+  </asset>
+  <worldbody>
+    <light pos="0 0 2" dir="0 0 -1" directional="true" />
+    <!-- 无 plane floor -- 纯 mesh 碰撞测试 -->
+    <geom name="box_mesh" type="mesh" mesh="test_box" pos="0 0 0.15"
+          contype="1" conaffinity="1" rgba="0.8 0.3 0.3 1" friction="0.8 0.3 0.3"/>
+  </worldbody>
+</mujoco>
+""")
+
+import motrixsim as mtx
+model = mtx.load_model(xml_path)
+print('加载成功')
+
+data = mtx.SceneData(model, batch=[1])
+data.reset(model)
+body = model.get_body(0)
+
+init_pos = model.compute_init_dof_pos().reshape(1, -1)
+init_pos[0, 0:2] = 0.0   # 在 box 正上方
+init_pos[0, 2] = 0.8      # 从 0.8m 自由落体 (box 顶面在 z=0.25)
+data.set_dof_pos(init_pos, model)
+model.forward_kinematic(data)
+
+print('自由落体到 box mesh (顶部 z=0.25):')
+for i in range(80):
+    model.step(data)
+    bz = body.get_pose(data)[0, 2]
+    if i < 15 or i % 15 == 0:
+        print(f'  步{i+1}: base_z={bz:.4f}')
+
+bz_final = body.get_pose(data)[0, 2]
+print(f'\n最终 base_z={bz_final:.4f}')
+if bz_final > 0.45:
+    print('✅ mesh 碰撞正常！机器人站在 box 上')
+elif bz_final < 0.10:
+    print('❌ mesh 碰撞不工作！机器人穿透 box 坠入深渊')
+else:
+    print(f'⚠ 不确定: base_z={bz_final:.4f}')

scripts/test_zscale_limit.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+"""测试 MotrixSim hfield 的 z_scale 上限——自己动手测，不信文档。"""
+import os, sys, numpy as np, time
+
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import motrix_envs.locomotion.go1.dreamwaq  # noqa
+from motrix_envs import registry as env_registry
+
+NUM_ENVS = 256
+TEST_STEPS = 100
+os.environ["DREAMWAQ_TERRAIN"] = "flat"  # 用 flat 场景，手动覆盖 hfield 参数
+
+
+def _build_custom_hfield(x_radius, y_radius, z_scale, z_base=0.001):
+    """构建自定义 hfield 描述字符串，用于覆盖 XML 中的 hfield 参数。"""
+    import tempfile, cv2
+    # 生成一个纯斜坡的 hfield PNG 用于测试
+    nx, ny = 200, 200
+    canvas = np.zeros((ny, nx), dtype=np.uint16)
+    # 从左上到右下的斜坡：高度从 0 到 z_scale
+    for i in range(ny):
+        for j in range(nx):
+            # 对角线斜坡，最高点在右下角
+            h = int((i + j) / (nx + ny) * 65535)
+            canvas[i, j] = h
+    # 中间 1/3 区域做平台（平坦）
+    cx, cy = nx // 2, ny // 2
+    p = nx // 6
+    canvas[cy - p:cy + p, cx - p:cx + p] = 0
+
+    out_d = os.path.join(os.path.dirname(__file__), "..",
+                         "motrix_envs", "src", "motrix_envs",
+                         "locomotion", "go1", "xmls", "assets")
+    os.makedirs(out_d, exist_ok=True)
+    png_path = os.path.join(out_d, "zscale_test.png")
+    cv2.imwrite(png_path, canvas)
+
+    return png_path, (x_radius, y_radius, z_scale, z_base)
+
+
+def test_z_scale(z_scale, n_envs=NUM_ENVS, n_steps=TEST_STEPS):
+    """测试给定 z_scale 下的物理稳定性。"""
+    import tempfile, cv2
+
+    # 生成测试 hfield
+    out_d = os.path.join(os.path.dirname(__file__), "..",
+                         "motrix_envs", "src", "motrix_envs",
+                         "locomotion", "go1", "xmls", "assets")
+    os.makedirs(out_d, exist_ok=True)
+
+    # 简单斜坡地形
+    nx, ny = 40, 40  # 小尺寸快速生成
+    canvas = np.zeros((ny, nx), dtype=np.uint16)
+    for i in range(ny):
+        for j in range(nx):
+            h = int((i / ny) * 65535)  # y 方向斜坡
+            canvas[i, j] = h
+    # 中央平台
+    p = nx // 6
+    canvas[ny // 2 - p:ny // 2 + p, nx // 2 - p:nx // 2 + p] = 0
+
+    png_path = os.path.join(out_d, "zscale_test.png")
+    cv2.imwrite(png_path, canvas)
+
+    total_x = nx * 0.1  # HS=0.1, 粗略
+    total_y = ny * 0.1
+
+    # 构建 scene XML
+    xml_path = os.path.join(os.path.dirname(__file__), "..",
+                            "motrix_envs", "src", "motrix_envs",
+                            "locomotion", "go1", "xmls", "scene_zscale_test.xml")
+    xml = f"""<mujoco model="zscale test">
+  <include file="go1_motor_actuator.xml" />
+  <include file="materials.xml" />
+  <statistic center="0 0 0.3" extent="2" />
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
+    <rgba haze="0.15 0.25 0.35 1" />
+    <global azimuth="120" elevation="-20" />
+  </visual>
+  <asset>
+    <hfield name="test_hf" file="assets/zscale_test.png"
+            size="{total_x/2:.1f} {total_y/2:.1f} {z_scale:.3f} 0.001" />
+  </asset>
+  <worldbody>
+    <light pos="0 0 2" dir="0 0 -1" directional="true" />
+    <geom name="floor" pos="0 0 0" type="hfield" hfield="test_hf"
+          material="motphys-ground" contype="1" conaffinity="0" priority="1" friction="0.6" />
+  </worldbody>
+  <sensor>
+    <contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
+    <contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
+    <contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
+    <contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
+  </sensor>
+</mujoco>"""
+    with open(xml_path, "w") as f:
+        f.write(xml)
+
+    import motrixsim as mtx
+    t0 = time.time()
+
+    try:
+        model = mtx.load_model(xml_path)
+        data = mtx.SceneData(model, batch=[n_envs])
+        data.reset(model)
+        body = model.get_body(0)
+
+        # 随机初始位置（在平台上）
+        init_pos = model.compute_init_dof_pos()
+        init_pos = np.tile(init_pos, (n_envs, 1))
+        init_pos[:, 0] += np.random.uniform(-0.5, 0.5, n_envs)
+        init_pos[:, 1] += np.random.uniform(-0.5, 0.5, n_envs)
+        init_pos[:, 2] = 0.5  # 从 0.5m 掉落
+        data.set_dof_pos(init_pos.astype(np.float32), model)
+        model.forward_kinematic(data)
+
+        heights = np.zeros((n_steps, n_envs), dtype=np.float32)
+        fall_count = 0
+        for step in range(n_steps):
+            model.step(data)
+            h = body.get_pose(data)[:, 2]
+            heights[step] = h
+            # 摔倒检测：base_z < 0.15 (趴了)
+            fall_count += np.sum(h < 0.15)
+
+        mean_h = float(np.mean(heights[-20:]))  # 最后 20 步平均
+        total_falls = fall_count
+        dt = time.time() - t0
+        return mean_h, total_falls, dt
+    except Exception as e:
+        return None, str(e), 0
+
+
+if __name__ == "__main__":
+    print(f"{'z_scale':>8s}  {'mean_base_z':>12s}  {'falls':>8s}  {'time':>8s}  verdict")
+    print("-" * 65)
+
+    for zs in [0.3, 0.5, 0.54, 0.8, 1.0, 1.5, 2.0, 3.0]:
+        mean_h, falls, dt = test_z_scale(zs, n_envs=64, n_steps=50)
+        if mean_h is None:
+            print(f"{zs:8.3f}  {'ERROR':>12s}  {str(falls)[:20]:>8s}")
+            continue
+        ok = "✅ 稳定" if mean_h > 0.25 and falls < 10 else "⚠ 不稳" if mean_h > 0.15 else "❌ 崩溃"
+        print(f"{zs:8.3f}  {mean_h:12.4f}  {falls:8d}  {dt:7.1f}s  {ok}")

scripts/view_dreamwaq.py

@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+"""DreamWaQ 地形可视化。
+
+键盘：
+  R=重置  H=高度采样点  T=遍历出生点调试  Esc=退出
+
+用法:
+    uv run scripts/view_dreamwaq.py                    # 金字塔地形
+    uv run scripts/view_dreamwaq.py --flat --num-envs 1
+"""
+import argparse, os, sys, time
+import numpy as np
+
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+
+if "--flat" in sys.argv:
+    os.environ["DREAMWAQ_TERRAIN"] = "flat"
+elif "--flat-stairs" in sys.argv:
+    os.environ["DREAMWAQ_TERRAIN"] = "flat_stairs"
+elif "--stairs" in sys.argv:
+    os.environ["DREAMWAQ_TERRAIN"] = "stairs"
+else:
+    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
+from motrix_envs.np.renderer import NpRenderer
+from motrix_envs.math import quaternion
+from motrixsim.render import RenderClosedError
+
+# 地形类型名称（与 gen_dreamwaq_terrain.py 的 PROPORTIONS 对应）
+PROPORTIONS = [0.1, 0.1, 0.35, 0.35, 0.1]
+CUM = [sum(PROPORTIONS[:i + 1]) for i in range(len(PROPORTIONS))]
+TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"]
+
+
+def _cell_origin(row, col, border_m=5.0, cell_m=8.0, num_rows=10, num_cols=20):
+    """计算 cell (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 _type_name(col):
+    """根据列索引返回地形类型名称。"""
+    choice = col / 20 + 0.001
+    for i, cum in enumerate(CUM):
+        if choice < cum:
+            return TYPE_NAMES[i]
+    return TYPE_NAMES[-1]
+
+
+def _spawn_at(env, cx, cy, spawn_z=None):
+    """在指定世界坐标 spawn 单个机器人。"""
+    state = env._state
+    data = state.data
+    init_pos = env._init_dof_pos.copy().reshape(1, -1)
+    init_pos[0, 0] = cx
+    init_pos[0, 1] = cy
+    # 计算地形高度
+    terrain_z = float(env._sample_terrain_height(
+        np.array([[cx, cy]], dtype=np.float32), radius=0.35)[0])
+    if spawn_z is None:
+        spawn_z = terrain_z + 0.45  # 默认 clearance
+    init_pos[0, 2] = spawn_z
+    data.reset(env._model)
+    data.set_dof_pos(init_pos, env._model)
+    env._model.forward_kinematic(data)
+    # 重置 info
+    state.info["commands"][0] = np.array([0.0, 0.0, 0.0], dtype=np.float32)
+    state.info["steps"][0] = 0
+    state.info["obs_history"][0] = 0.0
+    return terrain_z, spawn_z
+
+
+def main():
+    p = argparse.ArgumentParser(description="DreamWaQ 地形可视化")
+    p.add_argument("--num-envs", type=int, default=1)
+    p.add_argument("--flat", action="store_true")
+    p.add_argument("--flat-stairs", action="store_true")
+    p.add_argument("--stairs", action="store_true")
+    p.add_argument("--level", type=int, default=None)
+    p.add_argument("--no-stand", action="store_true")
+    p.add_argument("--vx", type=float, default=0.5)
+    args = p.parse_args()
+
+    env = env_registry.make("go1-dreamwaq-walk", num_envs=max(args.num_envs, 1))
+    if args.level is not None:
+        env._force_level = args.level
+
+    env.init_state()
+    n = env._num_envs
+    cmd = np.array([args.vx, 0.0, 0.0], dtype=np.float32)
+
+    try:
+        renderer = NpRenderer(env)
+    except Exception as e:
+        print(f"[ERROR] 渲染器创建失败: {e}")
+        renderer = None
+
+    terrain_name = os.environ.get("DREAMWAQ_TERRAIN", "pyramid")
+    print(f"[View] {n} 机器人 | 地形={terrain_name}")
+    print(f"[View] R=重置 H=高度点 T=遍历出生点 Esc=退出")
+
+    show_heights = False
+    traverse_mode = False
+    traverse_row = 0
+    traverse_col = 0
+    traverse_pending = False  # 刚切换 cell, 等待稳定
+    traverse_settle = 0
+    step_count = 0
+
+    # 地形信息
+    num_rows = env._num_rows
+    num_cols = env._num_cols
+    cell_m = env._cell_size
+
+    def enter_traverse():
+        nonlocal traverse_mode, traverse_row, traverse_col, traverse_pending, traverse_settle
+        traverse_mode = True
+        traverse_row = 0
+        traverse_col = 0
+        traverse_pending = True
+        traverse_settle = 0
+        print(f"\n[T] 遍历模式: {num_rows}行 × {num_cols}列")
+        print(f"[T] 按 T 前进, R 退出遍历\n")
+
+    def exit_traverse():
+        nonlocal traverse_mode, traverse_pending
+        traverse_mode = False
+        traverse_pending = False
+        env.init_state()
+        print("[T] 退出遍历模式\n")
+
+    def advance_traverse():
+        nonlocal traverse_row, traverse_col, traverse_pending, traverse_settle
+        traverse_col += 1
+        if traverse_col >= num_cols:
+            traverse_col = 0
+            traverse_row += 1
+        if traverse_row >= num_rows:
+            print("[T] 遍历完成! 按 R 退出")
+            traverse_row = num_rows - 1
+            traverse_col = num_cols - 1
+            return False
+        traverse_pending = True
+        traverse_settle = 0
+        return True
+
+    def do_traverse_spawn():
+        """在当前位置 spawn 并打印信息。"""
+        nonlocal traverse_settle
+        cx, cy = _cell_origin(traverse_row, traverse_col)
+        tname = _type_name(traverse_col)
+        terrain_z, spawn_z = _spawn_at(env, cx, cy)
+
+        # 标记 spawn 点（绿色球）
+        if renderer is not None:
+            g = renderer._render.gizmos
+            g.draw_sphere(0.15, (np.float32(cx), np.float32(cy),
+                                 np.float32(spawn_z)))
+
+        # 让机器人稳定几步
+        for _ in range(30):
+            env.step(np.zeros((1, 12), dtype=np.float32))
+            if renderer is not None:
+                renderer.render()
+                time.sleep(0.005)
+
+        base_z = env._body.get_pose(env._state.data)[0, 2]
+        contacts = env._state.info.get("contacts", np.zeros(4))
+        cf = env._state.info.get("privileged_obs",
+                                  np.zeros((1, 247)))[0, 45:57]
+        total_cf = np.sum(np.abs(cf))
+
+        print(f"  r{traverse_row}c{traverse_col:02d}  {tname:6s}  "
+              f"origin=({cx:+.0f},{cy:+.0f})  "
+              f"terrain_z={terrain_z:.3f}  spawn_z={spawn_z:.3f}  "
+              f"base_z={base_z:.3f}  cf={total_cf:.0f}N  "
+              f"feet={contacts.astype(int).tolist()}")
+
+        traverse_settle = 30
+
+    try:
+        while True:
+            # ── 键盘 ──
+            if renderer is not None:
+                try:
+                    inp = renderer._render.input
+                    if inp.is_key_just_pressed("r"):
+                        if traverse_mode:
+                            exit_traverse()
+                        else:
+                            env.init_state()
+                            step_count = 0
+                            print("[R] 重置")
+                    if inp.is_key_just_pressed("h"):
+                        show_heights = not show_heights
+                        print(f"[H] 高度点: {'开' if show_heights else '关'}")
+                    if inp.is_key_just_pressed("t"):
+                        if traverse_mode:
+                            advance_traverse()
+                        else:
+                            enter_traverse()
+                except Exception:
+                    pass
+
+            # ── 遍历模式 ──
+            if traverse_mode and traverse_pending:
+                do_traverse_spawn()
+                traverse_pending = False
+
+            # ── 正常模式动作 ──
+            if not traverse_mode or traverse_settle > 0:
+                if args.no_stand:
+                    env._state.info["commands"][:] = cmd
+                else:
+                    act = 0.3 * (np.random.rand(n, 12).astype(np.float32) - 0.5)
+                    env.step(act)
+                if traverse_settle > 0:
+                    traverse_settle -= 1
+
+                if step_count % 100 == 0 and step_count > 0:
+                    bz = env._body.get_pose(env._state.data)[0, 2]
+                    print(f"[{step_count}] base_z={bz:.3f}")
+
+            # ── 高度点 ──
+            if show_heights and renderer is not None:
+                state = env._state
+                pose = env._body.get_pose(state.data)
+                bp = pose[0, :3]
+                yaw = quaternion.get_yaw(pose[0:1, 3:7])[0]
+                cos_y, sin_y = np.cos(yaw), np.sin(yaw)
+                for gy in env._hy:
+                    for gx in env._hx:
+                        wx = bp[0] + cos_y * gx - sin_y * gy
+                        wy = bp[1] + sin_y * gx + cos_y * gy
+                        try:
+                            wz = float(env._sample_terrain_height(
+                                np.array([[wx, wy]]))[0])
+                            g = renderer._render.gizmos
+                            g.draw_sphere(0.02, (np.float32(wx), np.float32(wy),
+                                                 np.float32(wz)))
+                        except Exception:
+                            pass
+
+            if renderer is not None:
+                renderer.render()
+                time.sleep(0.01)
+
+            step_count += 1
+
+    except (KeyboardInterrupt, RenderClosedError):
+        pass
+    try:
+        if renderer is not None:
+            renderer.close()
+    except Exception:
+        pass
+    print("[View] 结束")
+
+
+if __name__ == "__main__":
+    main()