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

2026-06-30 14:53:12 +08:00
parent c569f4d6d9
commit 422421d263
17 changed files with 2056 additions and 17 deletions
--- a/docs/dreamwaq_replication.md
+++ 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
+```
--- a/docs/dreamwaq_usage.md
+++ 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
+```