Files
Motrixlab/docs/dreamwaq_alignment.md

190 lines
11 KiB
Markdown
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# DreamWaQ 复现与对齐记录
> 目标:让本地 `go1-dreamwaq-walk` 完全对齐并复现 [Manaro-Alpha/DreamWaQ](https://github.com/Manaro-Alpha/DreamWaQ.git)(基于 IsaacGym + rsl_rl 的实现),迁移到 MotrixLabMotrixSim + 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 地形高度为 01.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](../motrix_envs/src/motrix_envs/locomotion/go1/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](../motrix_envs/src/motrix_envs/locomotion/go1/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](../motrix_envs/src/motrix_envs/locomotion/go1/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=4200Hz 物理 / 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.2std 0.143.2)。但部署管线直接喂**原始观测**给 actor输入尺度完全错误 → 策略输出垃圾动作 → 倒地。
- **验证**:加上归一化后,站立从 **0/16 → 16/16 upright**
- **修复**
- [export_dreamwaq_onnx.py](../scripts/export_dreamwaq_onnx.py):从 checkpoint 提取 `state_preprocessor` 的 mean/var前 64 维 = code+obs烘焙进 ONNX 模型actor 前做归一化+clip
- [play_dreamwaq.py](../scripts/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](../motrix_rl/src/motrix_rl/tasks/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](../motrix_envs/src/motrix_envs/locomotion/go1/cfg.py) | PD 增益、关节角度、噪声尺度 |
| [motrix_envs/.../go1/dreamwaq.py](../motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py) | **地形生成修复**、ctrl 频率、命令范围、奖励、域随机化、观测噪声 |
| [motrix_envs/.../go1/walk_np.py](../motrix_envs/src/motrix_envs/locomotion/go1/walk_np.py) | 观测顺序、**feet_air_time 修复** |
| [motrix_rl/.../skrl/jax/train/dreamwaq_ppo.py](../motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py) | CENet 解码目标、VAE 缓冲、正交初始化 |
| [motrix_rl/.../tasks/go1_dreamwaq.py](../motrix_rl/src/motrix_rl/tasks/go1_dreamwaq.py) | **SKRL 字段名修复**、超参 |
| [scripts/gen_dreamwaq_terrain.py](../scripts/gen_dreamwaq_terrain.py) | 地形高度尺度 |
| [scripts/play_dreamwaq.py](../scripts/play_dreamwaq.py) | **GPU 预分配修复** |
| [scripts/dreamwaq_sim2sim_mujoco.py](../scripts/dreamwaq_sim2sim_mujoco.py) | PD/角度/decimation 与训练对齐 |
| [scripts/export_dreamwaq_onnx.py](../scripts/export_dreamwaq_onnx.py) | CENet 编码器 + Actor → ONNX |
---
## 6. 使用方法
### 训练
```bash
uv run scripts/train_dreamwaq.py --num-envs 2048 --timesteps 100M --seed 42
```
> 用 2048 envs 避免 4096 在 ~70% 处的 OOMexit 144。检查点存于 `runs/go1-dreamwaq-walk/skrl/`。
### 原生查看MotrixSimGPU 空闲时)
```bash
uv run scripts/play_dreamwaq.py --num-envs 9
```
### 导出 ONNX
```bash
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
```bash
uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain dreamwaq
# 控件: W/S 前后 Q/E 左右 A/D 旋转 Space 停 R 重置 Esc 退出
```
### TensorBoard
```bash
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。