diff --git a/docs/dreamwaq_replication.md b/docs/dreamwaq_replication.md
new file mode 100644
index 0000000..b40252c
--- /dev/null
+++ b/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
+```
diff --git a/docs/dreamwaq_usage.md b/docs/dreamwaq_usage.md
new file mode 100644
index 0000000..d2047af
--- /dev/null
+++ b/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
+```
diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py
new file mode 100644
index 0000000..1cdac59
--- /dev/null
+++ b/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)
diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/dreamwaq_terrain.png b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/dreamwaq_terrain.png
new file mode 100644
index 0000000..c6ed826
--- /dev/null
+++ b/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
diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_flat.xml b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_flat.xml
new file mode 100644
index 0000000..ce79e45
--- /dev/null
+++ b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_flat.xml
@@ -0,0 +1,25 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_terrain.xml b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_terrain.xml
new file mode 100644
index 0000000..310efa6
--- /dev/null
+++ b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_terrain.xml
@@ -0,0 +1,34 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/motrix_rl/src/motrix_rl/rslrl/cfg.py b/motrix_rl/src/motrix_rl/rslrl/cfg.py
index 3b3910a..cdc5649 100644
--- a/motrix_rl/src/motrix_rl/rslrl/cfg.py
+++ b/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
diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/models/__init__.py b/motrix_rl/src/motrix_rl/rslrl/torch/models/__init__.py
new file mode 100644
index 0000000..5654ccc
--- /dev/null
+++ b/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 自定义模型模块。"""
diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py b/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py
new file mode 100644
index 0000000..ce310e5
--- /dev/null
+++ b/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())}")
diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py b/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py
new file mode 100644
index 0000000..2cb8f05
--- /dev/null
+++ b/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
diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/train/ppo.py b/motrix_rl/src/motrix_rl/rslrl/torch/train/ppo.py
index 0f7238a..25fa51d 100644
--- a/motrix_rl/src/motrix_rl/rslrl/torch/train/ppo.py
+++ b/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
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ # Determine device（可通过 MOTRIX_DEVICE=cpu 强制 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)
diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py b/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py
index 919c2ef..25c7d12 100644
--- a/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py
+++ b/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
- obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+ # 构建多键 TensorDict（policy + obs_history + privileged_obs）
+ 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)
-
- # Create TensorDict for observations
- obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+ # 构建多键 TensorDict
+ obs = TensorDict(self._build_obs_dict(state),
+ 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({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+ obs = TensorDict(self._build_obs_dict(self._state),
+ batch_size=[self._num_envs], device=self._device)
return obs
def render(self) -> None:
diff --git a/motrix_rl/src/motrix_rl/tasks/go1.py b/motrix_rl/src/motrix_rl/tasks/go1.py
index b100268..565cacb 100644
--- a/motrix_rl/src/motrix_rl/tasks/go1.py
+++ b/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"
diff --git a/scripts/gen_dreamwaq_terrain.py b/scripts/gen_dreamwaq_terrain.py
new file mode 100644
index 0000000..1b89f9c
--- /dev/null
+++ b/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"""
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+"""
+ 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()
diff --git a/scripts/test_mesh_collision.py b/scripts/test_mesh_collision.py
new file mode 100644
index 0000000..2ffa443
--- /dev/null
+++ b/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("""
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+""")
+
+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}')
diff --git a/scripts/test_zscale_limit.py b/scripts/test_zscale_limit.py
new file mode 100644
index 0000000..7c48212
--- /dev/null
+++ b/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"""
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+"""
+ 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}")
diff --git a/scripts/view_dreamwaq.py b/scripts/view_dreamwaq.py
new file mode 100644
index 0000000..3354250
--- /dev/null
+++ b/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()