cyy/Motrixlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
8x54zj-m
2026-06-30 14:53:12 +08:00
parent c569f4d6d9
commit 422421d263
17 changed files with 2056 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
docs/dreamwaq_replication.md Normal file
View File

@@ -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
```

78
docs/dreamwaq_usage.md Normal file
View File

@@ -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
```

625
motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py Normal file
View File

@@ -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)

3
motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/dreamwaq_terrain.png Normal file
View File

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

25
motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_flat.xml Normal file
View File

@@ -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>

34
motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_dreamwaq_terrain.xml Normal file
View File

@@ -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>

4
motrix_rl/src/motrix_rl/rslrl/cfg.py
View File

@@ -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

2
motrix_rl/src/motrix_rl/rslrl/torch/models/__init__.py Normal file
View File

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

162
motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py Normal file
View File

@@ -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())}")

205
motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py Normal file
View File

@@ -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

37
motrix_rl/src/motrix_rl/rslrl/torch/train/ppo.py
View File

@@ -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)

40
motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py
View File

@@ -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:

75
motrix_rl/src/motrix_rl/tasks/go1.py
View File

@@ -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"

223
scripts/gen_dreamwaq_terrain.py Normal file
View File

@@ -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()

86
scripts/test_mesh_collision.py Normal file
View File

@@ -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}')

148
scripts/test_zscale_limit.py Normal file
View File

@@ -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}")

269
scripts/view_dreamwaq.py Normal file
View File

@@ -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()