10 KiB
DreamWaQ on MotrixSim — rsl_rl 复现交接文档
目标:在 MotrixLab(MotrixSim + numpy 物理)上忠实复现 Manaro-Alpha/DreamWaQ(原版基于 IsaacGym + 定制 rsl_rl-1.0.2),让 Go1 在地形上行走。
最后更新:2026-06-26。当前状态:rsl_rl 路线训练中,趋势健康(站立解决、行走学习中)。
0. 一句话总览
原版 DreamWaQ = ActorCritic_DWAQ(CENet+Actor+Critic 一体网络) + PPO 联合训练 VAE + 金字塔地形课程。我把上游定制的 rsl_rl-1.0.2 核心原样移植到 MotrixLab,写了个 numpy↔torch 环境适配器,并针对 MotrixSim 的两个 reality gap(hfield 大 z_scale 接触不稳、IsaacGym 超参偏激进)做了适配。
1. 两条路线(背景)
复现走过两条路线,rsl_rl 是当前主线:
| SKRL-JAX 路线(旧,参考) | rsl_rl 路线(当前主线) | |
|---|---|---|
| 代码 | motrix_rl/skrl/jax/train/dreamwaq_ppo.py |
motrix_rl/dwaq_rsl/(新模块) |
| 网络 | CENet 拼在 wrapper,Actor/Critic 在 SKRL | 上游 ActorCritic_DWAQ 一体网络 |
| VAE 训练 | 单独一步(CENet 不参与 PPO 梯度) | PPO 联合优化(和上游一致) |
| 结果 | 站立解决,但学不会走(机器人趴/站着不动) | 趋势健康,行走学习中 |
为什么切到 rsl_rl:SKRL 路线机器人学不会走,根因是 VAE 单独训练 —— CENet 的 code(latent+vel)早期是垃圾,actor 用不上。上游把 VAE 和 PPO 联合优化(一个 optimizer step),code 早期就有意义。切到 rsl_rl 后同样的 std=1.0 下 reward 从卡死的 0.025 变成快速上涨,这是关键转折。
SKRL 路线的详细 bug 修复记录见 dreamwaq_alignment.md(地下生成、feet_air_time、部署归一化等,这些 env 层 bug 两条路线都受益)。
2. 新建模块:motrix_rl/src/motrix_rl/dwaq_rsl/
把上游 rsl_rl-1.0.2 的核心原样移植(纯 torch,去掉 isaacgym 依赖,改相对 import):
| 文件 | 来源 | 作用 |
|---|---|---|
actor_critic_dwaq.py |
上游 modules/actor_critic_DWAQ.py 原样 |
CENet(encoder 225→128→64 + latent/vel 头)+ decoder + Actor[512,256,128] + Critic[512,256,128]。act(obs, obs_history) 把 CENet code(19)和 obs(45)拼成 64 维喂 actor |
ppo.py |
上游 algorithms/ppo.py 原样 |
PPO + VAE 联合 loss:loss = surrogate + value + autoenc,其中 autoenc = MSE(vel, prev_priv[:,45:48]) + MSE(decode, obs) + beta*KL |
rollout_storage.py |
上游 storage/rollout_storage.py 原样 |
经验存储,含 prev_critic_obs(VAE 的 vel_target 来源)、observation_history |
utils.py |
上游 utils/utils.py |
split_and_pad_trajectories |
runner.py |
上游 on_policy_runner.py 适配 |
去掉 rsl_rl 包 import + isaacgym 噪声,改用本地模块。learn() 主循环、save/load、tensorboard 日志 |
env_wrapper.py |
新写 | DwaqVecEnv:MotrixLab DreamWaQTask(numpy)→ 上游 VecEnv 接口(torch)。见下 |
DwaqVecEnv(env_wrapper.py)关键点
上游 runner 要求 env.step() 返回 7 元组:(obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, extras)。MotrixLab 的 DreamWaQTask 已提供 obs(45)、info["privileged_obs"](235)、info["obs_history"](N,5,45)。wrapper 做:
- numpy↔torch 转换(MotrixSim 物理在 CPU,网络在 GPU,每步传输)
obs_historyflatten:(N,5,45) → (N,225)- 维护
prev_privileged_obs(step 前的 privileged,VAE 用其 [45:48]=base_vel 作 vel_target) extras["time_outs"]= truncated- rewards/dones 是 (N,) 不是 (N,1)(上游约定,storage 内部 view 成 (N,1)) — 这是个踩过的坑
维度速查
obs=45, privileged=235(=obs45+base_vel3+heights187), obs_hist=225(=5*45)
actor_in=64(=code19+obs45), critic_in=235, cenet_in=225, cenet_out=19(=vel3+latent16)
3. 环境层改动:motrix_envs/.../go1/dreamwaq.py
(这些改动 rsl_rl 路线和 SKRL 路线共用同一个 DreamWaQTask)
- 去掉命令课程(关键):原来我加过一个命令课程(cmd_scale 0.3→1.0),但它让机器人趴着刷小命令分(小命令下趴着也能 tracking)。改回上游 full range [-1,1](
commands.curriculum=False),机器人被要求走,趴着不划算。见resample_commands。 - level 0 起步:
reset()里 init terrain level 全设 0(平地),靠地形难度课程逐级升。原来 0-5 随机会让机器人 spawn 在难地形秒摔。 - 中心 spawn:
_make_origins坐标中心化(hfield 中心在世界原点 (0,0)),spawn 在 cell 的 3m 平台中心 ±1m。 - 地形高度采样:
_sample_terrain_height/_get_heights从 hfield 采样(hf.height_matrix+hf.bound),用于 spawn 高度 + privileged obs 的 heights。 - 奖励对齐上游(smoothness 用 joint_pos_target + step mask、power_distribution 用
var(abs(power))、base_height 扣地形高度等)。
4. 地形:scripts/gen_dreamwaq_terrain.py + ⚠️ MotrixSim hfield 坑
生成器
移植了上游 IsaacGym terrain_utils 的金字塔地形(纯 numpy):pyramid_sloped_terrain / pyramid_stairs_terrain / discrete_obstacles_terrain / random_uniform_terrain,10 行难度 × 20 列类型网格,每格 3m 中心平台,输出 16-bit PNG hfield + XML <size> 行。
⚠️⚠️ 最大的坑:MotrixSim hfield 在大 z_scale 下接触不稳
实测结论(零动作站立 20 步摔倒数):
| hfield z_scale | 机器人 above-terrain | 20步摔倒 |
|---|---|---|
| 3.78(上游全高金字塔) | 0.12(趴) | 1267 |
| 1.66 | 0.14(趴) | 1049 |
| 0.54(矮金字塔) | 0.35(正常站立) | 120 |
| plane geom(对比) | 0.33 | 0 |
- MotrixSim 用 hfield(高度场),陡台阶垂直面只能近似成陡坡,机器人脚(sphere geom)打滑;z_scale 越大越不稳。
- 上游用 trimesh(三角网格),垂直面精确。调研确认 MotrixSim 不支持非凸 trimesh 碰撞(mesh 是凸包),接触参数(condim/solref/solimp)调优也无效。
- 解决:把金字塔楼梯/坡的高度缩小 ~4 倍(
make_terrain里step_height=0.012+0.020*diff、slope=0.10*diff),让 z_scale≈0.54。保留金字塔结构 + 难度课程,但高度适配 MotrixSim hfield。这是必要的 reality-gap 妥协(不完全忠实上游高度,但能训)。 scene_dreamwaq_terrain.xml的 hfieldsize="85 45 0.54 0.01":第1/2 是半径(全宽 170×90m),第3 是 z_scale,第4 是底部厚度。
重新生成地形后必须同步更新 XML 的 z_scale(gen 脚本会打印 size="..." 那行的值)。
5. 脚本(训练 / 导出 / 可视化)
| 脚本 | 作用 |
|---|---|
scripts/train_dreamwaq_rsl.py |
rsl_rl 训练。超参 dict 内嵌(对齐上游 Go1RoughCfgPPO)。--num-envs 2048 --iterations 3000 --init-noise-std 1.0 |
scripts/export_dreamwaq_rsl_onnx.py |
rsl 的 .pt checkpoint → ONNX。rsl 无 state 归一化(不像 SKRL),所以 ONNX 简单:obs(1,45)+obs_history(1,5,45)→action(1,12),用 mean CENet code |
scripts/gen_dreamwaq_terrain.py |
生成矮金字塔地形 PNG |
scripts/dreamwaq_sim2sim_mujoco.py |
MuJoCo sim2sim 可视化(用 ONNX)。PD/角度/decimation 已对齐(KP=28,KD=0.7,action_scale=0.25,decimation=4) |
命令
# 训练 (后台跑, print 会被 buffer, 用 tensorboard 看)
PYTHONUNBUFFERED=1 uv run scripts/train_dreamwaq_rsl.py --num-envs 2048 --iterations 3000 --seed 1 --init-noise-std 1.0
# 监控 (print 被 buffer, 看 tensorboard)
uv run tensorboard --logdir runs/go1-dreamwaq-walk/rsl_dwaq
# 导出最新 checkpoint
JAX_PLATFORMS=cpu uv run scripts/export_dreamwaq_rsl_onnx.py # 自动找最新 model_*.pt
# MuJoCo 可视化 (本地有显示的机器)
uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain flat # 平地最干净
6. 关键决策 + 为什么(reality gap 适配)
复现不是 1:1 照搬,MotrixSim vs IsaacGym 有真实差异,做了几处有意的适配(都在代码注释里标了原因):
- VAE 联合训练(用 rsl_rl 而非 SKRL):核心。SKRL 路线 VAE 单独训练→学不会走。
- 去命令课程:命令课程让机器人趴着刷小命令分;上游本来就 full range 无课程。
- 矮金字塔地形(z_scale 0.54):MotrixSim hfield 大 z_scale 接触不稳的妥协。
- level 0 平地起步:机器人先在平地学走,靠地形课程升级。
- 其余超参/网络/奖励严格对齐上游(PD 28/0.7、action_scale 0.25、decimation 4、init_noise_std 1.0、rollouts 24、lr 1e-3、entropy 0.01 等)。
7. 当前状态(2026-06-26)
训练中(runs/go1-dreamwaq-walk/rsl_dwaq/06-26_00-10-27),趋势健康:
| 指标 | it59 → it182 |
|---|---|
| mean_reward | 0.02 → 2.6 → 18.6 |
| episode 长度 | 12 → 101 → 143 |
| noise_std | 1.0 → 0.84 → 0.57(在学) |
| base_height | 稳定(不趴) |
机器人能站(above ~0.35)+ episode 越来越长 + std 持续降 = 在学走。需等训练跑完(3000 iter,~2.5h)再导出验证实际行走(站立率/前进速度跟踪/摔倒率,别只看奖励数字 — 之前吃过"奖励高但实际趴着"的亏)。
8. 已知问题 / 坑(接手必读)
- 训练 print 被 buffer:重定向到文件时 Python print 不实时 flush。用 tensorboard 监控,别等 print。(已加
PYTHONUNBUFFERED=1缓解) ps的 TIME 是 CPU 累计时间(跨核),不是 wall time。看 wall 用etime或 tfevents 时间戳。- VAE checkpoint 文件混在共享目录(SKRL 路线遗留):
runs/.../skrl/vae_*.pkl。rsl 路线 checkpoint 是rsl_dwaq/<run>/model_*.pt,别搞混。 - 物理在 CPU(MotrixSim numpy backend),网络在 GPU。每步 numpy↔torch 传输有开销,~16 it/min(2048 envs)。物理是瓶颈。
- MotrixSim hfield 不支持陡台阶(根本限制):如果要更忠实的高地形,得调研 MotrixSim 是否有 mesh 凸分解地形方案,或接受矮地形。
- 验证要看行为不看奖励:
only_positive_rewards+ 一堆"别动"惩罚,机器人容易学"趴着/站着不动"刷分。验证时务必测实际前进速度跟踪 + 摔倒率,不能只看 mean_reward。
9. 下一步
- 等当前训练跑完,导出 + MuJoCo 验证实际行走(前进/转向跟踪)。
- 若行走 OK:逐步提高地形难度(z_scale),找 MotrixSim hfield 能稳定接触的上限。
- 若行走不够好:检查地形课程是否真在升级(机器人有没有从 level 0 走到 level 1+)、调 init_noise_std。
- 长期:调研 MotrixSim mesh 地形(凸分解)以更忠实复现陡台阶。