Motrixlab/docs/dreamwaq_alignment.md

DreamWaQ 复现与对齐记录

目标：让本地 go1-dreamwaq-walk 完全对齐并复现 Manaro-Alpha/DreamWaQ（基于 IsaacGym + rsl_rl 的实现），迁移到 MotrixLab（MotrixSim + SKRL-JAX）。

最后更新：2026-06-25

---

1. 架构总览

CENet (VAE):  history(5×45=225) → [128,64] → latent(16) + vel_est(3) = code(19)
Decoder:      code(19) → [64,128] → obs_recon(45)
Actor:        code(19) + obs(45) = 64 → [512,256,128] → action(12)
Critic:       privileged_obs(235) → [512,256,128] → value(1)

观测(45): ang_vel(3) + gravity(3) + commands(3) + dof_pos(12) + dof_vel(12) + actions(12)
特权观测(235): obs(45) + base_vel(3) + heights(187, 17×11 网格)
增广观测(254, 训练用): [code(19) | obs(45) | base_vel(3) | heights(187)]

---

2. 🔴 真正让机器人站不起来的 Bug（最关键）

排查发现并非训练不充分，而是两个环境 bug 让约一半的训练数据是垃圾：

Bug A — 机器人生成在地形地下

• 现象：reset 时机器人固定生成在绝对高度 z=0.40，但 DreamWaQ 地形高度为 0–1.6m。在高地形格子上，机器人生成在地表以下最多 1.2m。
• 影响：31/64 的环境在生成瞬间就卡进地形 → 爆炸/立即终止/疯狂抖动 → dof_acc 高达 150 万 → 奖励被 only_positive_rewards 裁剪为 0。一半训练样本无效。
• 修复：新增 _sample_terrain_height(xy)，按生成点的局部地形高度设置 z = terrain_z + 0.40。修复后 0/64 在地下。
• 位置：dreamwaq.py reset() + _sample_terrain_height()

Bug B — feet_air_time 奖励恒为 0

• 现象：腾空时间在 update_feet_air_time 中被 *= ~contacts 清零，早于 reward 函数读取 first_contact = (air_time > 0) * contacts。落地瞬间 air_time 已是 0 → first_contact 永远为 0 → 步态奖励完全失效。
• 影响：机器人没有迈步激励，倾向于站着不动/拖行，学不会正常步态。
• 修复：按上游顺序重构 —— 先用 contact_filt = contact OR last_contact 算出 first_contact 和落地时的 air_time，再清零；reward 函数改为读取预计算值（无状态，对 DreamWaQ 的 _get_reward 二次调用安全）。
• 位置：walk_np.py update_feet_air_time() + _reward_feet_air_time()

Bug C — 初始地形难度过高

• 现象：首次 reset 在 0-9 级随机生成（含最难地形）。上游用 max_init_terrain_level=5。
• 修复：首次生成限制在 0-5 级，由课程学习逐步提升。
• 位置：dreamwaq.py reset()

---

3. 对齐改动（与上游一致）

控制 / 物理（cfg.py + dreamwaq.py）

参数	改动前	改动后（上游）
stiffness (Kp)	80	28.0
damping (Kd)	1.0	0.7
action_scale	0.05	0.25
sim_dt / ctrl_dt	0.01 / 0.01	0.005 / 0.02（decimation=4，200Hz 物理 / 50Hz 控制）

默认关节角度（cfg.py，非对称）

hips:   ±0.1 (FL/RL: +0.1, FR/RR: -0.1)
thighs: 前 0.8 / 后 1.0
calves: -1.5 (全部)

观测顺序（walk_np.py）

改为上游顺序 [ang_vel, gravity, commands, dof_pos, dof_vel, actions]，并加入观测噪声（add_noise=True, noise_level=1.0）。

命令范围（dreamwaq.py）

lin_vel_x/y/yaw 全部 [-1, 1]（之前 vel_x 上限是 2.0）。

奖励函数对齐（dreamwaq.py）

• smoothness：基于 joint_pos_target（= action×scale + default）的二阶差分，而非原始 action
• base_height：square(base_z - mean(measured_heights) - 0.30)，扣除局部地形高度
• power_distribution：var(torque×vel) 跨全部 12 关节，而非按腿分组
• 12 项奖励尺度与上游 Go1 base config 精确一致

域随机化（dreamwaq.py）

motor_strength / Kp_factor / Kd_factor ∈ [0.9, 1.1]，每 4 秒重采样（rand_interval_s=4），在 _compute_torques 中应用。

训练（dreamwaq_ppo.py + go1_dreamwaq.py）

• 关键 bug 修复：SKRL 字段名 entropy_coef→entropy_loss_scale、clip_param→ratio_clip、gamma→discount_factor。改前 entropy 被静默设为 0.0（字段名不存在）。
• CENet 解码目标：重构当前观测（历史最后一帧），而非预测下一帧（匹配上游 decode_target = obs_batch）
• 正交权重初始化 orthogonal(sqrt(2))，应用于 CENet/Actor/Critic 所有 Dense 层
• PPO 超参：rollouts=24, lr=1e-3, entropy=0.01, gamma=0.99, lam=0.95, clip=0.2, epochs=5, mini_batches=4

地形（gen_dreamwaq_terrain.py）

降低高度尺度（最大 2.73m → 1.66m），更适合 0.33m 高的机器人。生成后更新 XML 的 hfield size。

---

3.5 🔴🔴 部署缺少观测归一化（"站不起来"的真正根因）

这是导致导出/play/sim2sim 中机器人完全站不起来的根本原因，独立于训练质量。

• 现象：训练奖励正常（mean 35 / max 924），但导出 ONNX / play_dreamwaq.py / MuJoCo sim2sim 中机器人塌成一团，actor 输出动作幅度高达 12（关节目标偏移 ~3 rad，直接打到限位）。
• 根因：SKRL PPO 通过 _add_runtime_config 启用了 RunningStandardScaler 状态预处理器——训练时观测被 clip((x - running_mean) / (sqrt(var) + 1e-8), -5, 5) 归一化（running_mean 幅度达 ±3.2，std 0.14–3.2）。但部署管线直接喂原始观测给 actor，输入尺度完全错误 → 策略输出垃圾动作 → 倒地。
• 验证：加上归一化后，站立从 0/16 → 16/16 upright。
• 修复：
  • export_dreamwaq_onnx.py：从 checkpoint 提取 state_preprocessor 的 mean/var（前 64 维 = code+obs），烘焙进 ONNX 模型（actor 前做归一化+clip）
  • play_dreamwaq.py：policy_forward 中应用同样的归一化
  • sim2sim 用 ONNX，重新导出后自动修复
• 注意：RunningStandardScaler 是 SKRL 的附加项，上游 rsl_rl 没有（上游只用固定 obs_scales）。任何脱离 SKRL agent 的部署都必须复现这个归一化，否则策略失效。

---

3.6 🔴 训练不收敛：初始探索过猛（IsaacGym→MotrixSim reality gap）

修复部署归一化后发现策略根本没学会站立——在 MotrixSim 训练环境里 0/64 站立/行走，比"零动作"还差。

• 诊断：零动作（纯 PD 保持默认姿态）奖励 1.27/step，27/64 站立；训练后的策略奖励 0.035/step，0/64 站立。训练把策略训得比啥都不做还差。

• 根因：initial_log_std=0.0（std=1）× action_scale=0.25 = 每步 ±0.25 rad 随机抖动 @50Hz。实测：

  探索 std	抖动	站立率	平均奖励
  1.0（上游）	0.25rad	27/64	0.025
  0.37	0.09rad	29/64	0.512
  0.2	0.05rad	23/64	0.730

  机器人在所有 std 下都能站（站立率相近），但 std=1 时剧烈抖动产生的 dof_acc/action_rate 惩罚把站立的正奖励完全淹没（净奖励 0.025 ≈ 噪声）→ PPO 没有指向站立的梯度 → 收敛到塌倒局部最优。

• 为何上游能用：DreamWaQ 的 std=1 是为 IsaacGym 调的；MotrixSim 接触/动力学不同（你能用的 go1 平地任务用的是 5× 更温和的 action_scale=0.05）。

• 修复：go1_dreamwaq.py 设 initial_log_std=-1.0（std=0.37，抖动 0.09rad），让机器人早期保持直立、获得清晰的站立奖励信号后再 bootstrap 行走。这是 reality-gap 的必要偏离（上游用 0.0）。

---

4. 🛠 原生 play 渲染崩溃修复

• 现象：play_dreamwaq.py 启动 MotrixSim 渲染器时崩溃 Couldn't get swap chain texture，而其他 go1 任务的 play 正常。
• 根因：JAX 默认预分配 75% 显存（12GB 中 ~9GB），饿死了 Vulkan 渲染器。标准 play.py 因渲染器创建与 JAX 初始化间隔充足而不受影响。
• 修复：在 import jax 前设置 os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"，并优雅处理 RenderClosedError。
• 注意：训练占用 GPU 时无法同时跑原生渲染；请在训练结束后查看。

---

5. 修改的文件

文件	改动
motrix_envs/.../go1/cfg.py	PD 增益、关节角度、噪声尺度
motrix_envs/.../go1/dreamwaq.py	地形生成修复、ctrl 频率、命令范围、奖励、域随机化、观测噪声
motrix_envs/.../go1/walk_np.py	观测顺序、feet_air_time 修复
motrix_rl/.../skrl/jax/train/dreamwaq_ppo.py	CENet 解码目标、VAE 缓冲、正交初始化
motrix_rl/.../tasks/go1_dreamwaq.py	SKRL 字段名修复、超参
scripts/gen_dreamwaq_terrain.py	地形高度尺度
scripts/play_dreamwaq.py	GPU 预分配修复
scripts/dreamwaq_sim2sim_mujoco.py	PD/角度/decimation 与训练对齐
scripts/export_dreamwaq_onnx.py	CENet 编码器 + Actor → ONNX

---

6. 使用方法

训练

uv run scripts/train_dreamwaq.py --num-envs 2048 --timesteps 100M --seed 42

用 2048 envs 避免 4096 在 ~70% 处的 OOM（exit 144）。检查点存于 runs/go1-dreamwaq-walk/skrl/。

原生查看（MotrixSim，GPU 空闲时）

uv run scripts/play_dreamwaq.py --num-envs 9

导出 ONNX

uv run scripts/export_dreamwaq_onnx.py \
  --checkpoint runs/go1-dreamwaq-walk/skrl/<run>/checkpoints/best_agent.pickle \
  --vae runs/go1-dreamwaq-walk/skrl/vae_<N>.pkl \
  --output exports_go1_dreamwaq/policy.onnx

MuJoCo sim2sim

uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain dreamwaq
# 控件: W/S 前后  Q/E 左右  A/D 旋转  Space 停  R 重置  Esc 退出

TensorBoard

uv run tensorboard --logdir runs/go1-dreamwaq-walk

---

7. 待办 / 注意事项

• 检查点兼容性：改动前的检查点（含旧 PD/观测顺序/地形 bug）不兼容，应丢弃重训。
• 未对齐项（有意保留）：
  • 奖励未乘 dt（上游 _prepare_reward_function 会 scale *= dt）。这是全局均匀缩放，PPO 优势归一化会抵消，不影响行为；为避免干扰已调好的 value 学习而保留。
  • VAE 训练为 PPO 更新后独立一步，而非上游的联合优化（功能等价，更易维护）。
  • 地形几何为自定义生成器，非 IsaacGym terrain_utils（概念相似，具体形状不同）。
• OOM：4096 envs 曾在 70% 处被杀（exit 144），改用 2048。