世界模型实践路线¶
新手学世界模型,不要一开始复现 Genie 或 Cosmos。先从小环境学会“预测未来 -> 用预测选动作”,再逐步加图像、隐空间和视频。
1. 总路线¶
目标是一步步把复杂度加上去,而不是直接跳到大模型。
2. 阶段一:低维动力学模型¶
选择一个简单环境:
- GridWorld
- CartPole
- Pendulum
- MountainCar
训练一个模型:
\[
\hat{s}_{t+1} = f_\theta(s_t,a_t)
\]
你要做的事:
- 收集随机策略数据。
- 训练 MLP 转移模型。
- 预测一步未来。
- 预测多步未来。
- 画出真实轨迹和预测轨迹的偏差。
重点观察:一步预测看起来很好,多步预测可能快速漂移。
2.1 数据怎么收集?¶
先用随机策略或简单专家策略收集:
存成表格:
| 字段 | 示例 |
|---|---|
state |
CartPole 的位置、速度、角度、角速度 |
action |
左/右,或连续力 |
next_state |
下一步状态 |
reward |
环境奖励 |
done |
是否结束 |
2.2 最小训练循环¶
for batch in dataloader:
state, action, next_state = batch
pred_next_state = model(state, action)
loss = mse(pred_next_state, next_state)
optimizer.zero_grad()
loss.backward()
optimizer.step()
先别急着接规划。你应该先画出:
如果单步都不准,MPC 只会把错误放大。
3. 阶段二:用模型做 MPC¶
有了转移模型后,用随机 shooting 或 CEM 搜索动作。
最简单版本:
for each candidate_action_sequence:
state = current_state
total_reward = 0
for action in candidate_action_sequence:
state = model(state, action)
total_reward += reward_fn(state, action)
choose sequence with highest total_reward
execute first action
你要理解:
- horizon 太短会短视。
- horizon 太长误差会累积。
- 候选动作数量越多,规划越准但越慢。
- 每次只执行第一步能修正模型误差。
3.1 调参顺序¶
建议按这个顺序调:
- 固定 horizon = 5,先让动作能朝目标方向移动。
- 增加候选动作数量,看成功率是否提升。
- 增加 horizon,看是否出现模型误差导致的怪动作。
- 加动作限幅,避免搜索到训练分布外动作。
- 加 ensemble 或不确定性惩罚。
3.2 你应该保存的图¶
- 不同 horizon 下的成功率。
- 预测轨迹 vs 真实轨迹。
- 候选动作回报分布。
- 模型不确定性随 rollout 步数变化。
4. 阶段三:加入图像输入¶
把状态从低维变量换成图像:
可以先做一个简单 autoencoder:
- encoder 压缩图像。
- decoder 重建图像。
- dynamics 预测 latent。
这一阶段重点不是画面多漂亮,而是 latent 是否能支持控制。
4.1 Autoencoder 的检查¶
你至少要检查:
- 重建图是否保留物体位置。
- latent 维度太小是否丢目标。
- latent 维度太大是否学到无关纹理。
- 在 latent 中做最近邻检索,相似状态是否靠近。
4.2 Latent Dynamics 的训练¶
训练两部分:
loss 可以先用:
\[
\mathcal{L}
=\|o_t-\hat{o}_t\|^2
+\lambda\|z_{t+1}-\hat{z}_{t+1}\|^2
\]
入门时先别追求复杂 RSSM。先确认 latent 能被预测、能用于规划。
5. 阶段四:理解 Dreamer 风格训练¶
Dreamer 风格实践可以按这个理解:
- 从环境收集真实轨迹。
- 训练 RSSM 世界模型。
- 从真实 latent 出发,在模型里想象未来。
- 在 imagined trajectory 上训练 actor 和 critic。
- 用 actor 回到真实环境收集新数据。
你不一定要从零写完整 Dreamer,但要能看懂代码里的几类 loss:
- reconstruction loss
- reward loss
- continuation loss
- KL loss
- actor loss
- critic loss
5.1 看 Dreamer 代码时按模块找¶
| 模块名 | 你要找什么 |
|---|---|
| replay buffer | 真实环境序列如何存 |
| encoder | 图像如何变成 embedding |
| rssm/world model | prior、posterior、transition |
| decoder heads | image/reward/continue 预测 |
| imagination | imagined rollout 从哪里开始 |
| actor | 根据 latent 输出动作 |
| critic | 估计 imagined return |
不要从训练脚本第一行硬读。先找这些模块,会清楚很多。
6. 阶段五:视频世界模型小实验¶
做一个小视频预测任务:
- 输入前 4 帧。
- 输入动作或方向键。
- 预测下一帧或未来 4 帧。
可选环境:
- 小型 grid game。
- Atari-like toy environment。
- 简化机械臂仿真。
你要比较:
- 无动作条件 vs 有动作条件。
- 像素预测 vs latent 预测。
- 一步预测 vs 多步自回归预测。
6.1 视频实验最小指标¶
| 指标 | 看什么 |
|---|---|
| one-step loss | 下一帧是否基本对 |
| rollout drift | 连续生成是否越来越偏 |
| action sensitivity | 换动作后未来是否改变 |
| object consistency | 物体身份是否保持 |
| rule consistency | 游戏规则/墙体/碰撞是否保持 |
6.2 最小对照实验¶
做三组模型:
- 无动作模型:只输入过去帧。
- 有动作模型:输入过去帧 + 动作。
- 有动作 + latent 模型:先编码再预测 latent。
如果第 2 组明显优于第 1 组,说明动作条件真的被模型用上了。
7. 阶段六:读大模型路线¶
当你已经做过小闭环,再读大模型会清楚很多:
| 工作 | 重点看什么 |
|---|---|
| World Models | VAE/RNN/Controller 三模块 |
| PlaNet | RSSM 和 CEM 规划 |
| DreamerV3 | imagination + actor-critic |
| TD-MPC2 | decoder-free latent planning |
| I-JEPA/V-JEPA | embedding prediction |
| Genie | latent action 和交互式环境 |
| GAIA-1/UniSim | 自动驾驶闭环模拟 |
| Cosmos | physical AI 世界基础模型 |
8. 四周学习安排¶
第 1 周:RL 与低维模型¶
- 学 MDP/POMDP。
- 训练低维状态转移模型。
- 做多步 rollout 可视化。
第 2 周:MPC 与不确定性¶
- 实现随机 shooting 或 CEM。
- 调 horizon 和候选数量。
- 加 ensemble 观察不确定性。
第 3 周:图像和 latent dynamics¶
- 训练 autoencoder。
- 在 latent 中预测未来。
- 阅读 World Models 和 PlaNet。
第 4 周:Dreamer 与视频世界模型¶
- 阅读 DreamerV3。
- 跑一个已有 Dreamer/TD-MPC demo。
- 做动作条件视频预测小实验。
- 阅读 Genie/Cosmos 的方法部分。
9. 常见问题排查¶
| 现象 | 常见原因 |
|---|---|
| 一步预测准,多步预测崩 | 误差累积、训练只优化 one-step |
| 模型预测画面清楚但控制差 | 表征保留了纹理,没保留任务状态 |
| MPC 找到奇怪动作 | 模型漏洞、分布外动作、不确定性没惩罚 |
| actor 在模型里很强,真实环境很弱 | 策略利用了模型错误 |
| 视频长期漂移 | 自回归误差累积、记忆不足 |
| latent 塌缩 | 表征损失设计不当,缺少正则或有效预测目标 |
9.1 排查顺序¶
遇到训练失败,按这个顺序查:
- 数据:状态、动作、next_state 是否对齐。
- 单位:动作单位、时间步长、归一化是否正确。
- 单步预测:one-step 是否能过拟合小数据。
- 多步 rollout:是否从第几步开始漂移。
- 规划器:是否搜索到了训练分布外动作。
- 闭环执行:真实观测是否被及时用于重规划。
先过拟合小数据
用 100 条 transition 训练到很低 loss。如果做不到,别急着改算法,先查数据管线和模型输入输出。
10. 入门完成标准¶
当你能做到下面这些,就算真正入门世界模型了:
- 能解释世界模型和 model-free RL 的区别。
- 能训练一个低维转移模型。
- 能用模型做简单 MPC。
- 能解释 latent dynamics、RSSM、imagination。
- 能说清 JEPA 为什么不预测像素。
- 能区分视频生成、视频预测和交互式世界模型。
- 知道世界模型用于机器人/自动驾驶时必须加安全约束。