World Model 从零搭建指南：强化学习与想象空间的三大模块与实操避坑

world model 这几年在 AI 圈里出镜率越来越高，强化学习、视频预测、多模态大模型都在提这个概念。我最早被它吸引，是因为被无模型强化学习的采样效率折磨得够呛：想训练一个会过迷宫、会开车、会抓物体的智能体，环境里动辄跑上百万步，真实世界里的机器人根本试不起。World Model 的思路完全不同，它更像是把环境“搬进”模型的脑子里，先脑补、后执行，很多策略在虚拟预演里就能完成评估。

这篇内容不是学术综述，而是我实际做 World Model 项目时的实操笔记。我会先用生活化的方式讲清楚它解决了什么问题，再拆解它的三大核心模块，接着给出一套从零训练的标准流程和关键参数，最后把我踩过的坑、排查思路整理成速查表。无论你是刚接触强化学习的新手，还是已经在做视频预测、模型预测控制的老手，这套笔记应该都能帮你少走不少弯路。

1. 一个能“脑补”出来的环境

1.1 World Model 到底解决了什么痛点

先把痛点放桌上。传统强化学习里，智能体只能在真实环境里“试错”，每一步都要和环境交互才能拿到状态转移和奖励。这种在线采样的方式有个致命问题：训练成本极高。一个简单的 2D 赛车游戏，想要学到基本的漂移过弯，动辄需要几十万步交互；换成机械臂抓取这种真实物理任务，现实世界根本不允许你跑几千次实验。

World Model 的核心思路就是：既然智能体要的是环境的规律，那不如先把规律本身学出来。具体来说，我们会训练一个模型，让它学会“环境对未来会怎么演化”的预测能力。这个模型输入当前状态和智能体的动作，输出对下一秒甚至未来几十步的预测。一旦这个模型够准，智能体就可以在模型的“想象空间”里疯狂试错，完全不需要触碰真实环境。训练一个能在脑子里跑几十万步的虚拟世界，比在现实里跑几十万步便宜太多了。

我用一个生活化的类比来解释：考驾照的时候，学员先在模拟器里练习科目二，等模拟器里已经不会压线、不会熄火了，再上真车。真车驾驶次数有限、撞坏了要花钱，但模拟器可以无限次重开。World Model 就是那台模拟器，而强化学习智能体就是那个学员。

这里要先澄清一个容易混淆的边界：深度学习里的生成模型也能生成未来帧，视频预测模型也能预测下一帧，但它们不一定算 World Model。真正的 World Model 强调两件事，第一，它要能压缩高维观测，比如把一帧 64×64×3 的图像压缩成几十维的向量；第二，它要在压缩后的空间里预测长期未来。只会在像素上做几帧外推的模型，并没有建立“世界”的概念，它更像是个视频插值器。

1.2 World Model 和常见模型的边界在哪里

我在做项目初期，经常被一个问题困扰：是不是只要能用神经网络预测下一步就算 World Model？答案是否定的。传统 model-based RL 的做法是用一个状态转移函数预测下一个状态，但这通常假设状态已经是低维向量，比如位置、速度、关节角。而真实环境给到智能体的往往是图像、点云这类高维观测，直接把像素丢给预测网络，计算量和噪声干扰都是灾难。

World Model 给出的解法是分两步走。先做一个“压缩器”，把高维像素变成紧凑的潜变量；再在这个潜空间里做动力学预测。这样做的好处是，预测目标从“生成一张清晰的未来图”变成了“预测几十个数字的分布”。前者要处理像素层面的细节，比如叶片颤动、光影变化，这些细节对任务往往不重要；后者只需要保留对决策有意义的语义信息，比如障碍物在哪里、车速是多少、前方有没有弯道。

另一个边界在于，World Model 不一定要做策略决策。它负责的只是“世界会怎么变”，至于智能体该怎么做，那是 policy 模块的事。很多初学者会把二者混为一谈，拿着一个训练好的 World Model 去问“为什么它不能直接帮我规划路径”，其实是因为没有接上决策模块。World Model 提供了想象空间，决策模块负责在想象空间里找最优动作，两者是上下游关系。

如果从概率图模型的角度看，World Model 在学的是一个条件概率分布 P(z_{t+1}, r_t | z_t, a_t)，其中 z_t 是压缩后的状态，a_t 是动作。它和时间序列模型、隐马尔可夫模型的区别在于，World Model 明确把“环境动态”和“观测生成”分开建模，并且都是在高维、部分可观测的环境中工作。这也是它能在机器人、自动驾驶、游戏 AI 这些场景里落地的根本原因——现实世界基本都是高维且部分可观测的。

2. 拆开看 World Model 的三大模块

2.1 视觉压缩模块：把高维画面变成向量

最早把 World Model 做成完整系统的，是 David Ha 和 Jürgen Schmidhuber 在 2018 年发表的论文。他们把整个模型拆成了三个模块：V（Vision）、M（Memory）、C（Controller）。我在实践里基本沿用了这个框架，因为它清晰到可以直接照抄。

V 模块即视觉压缩模型，通常用变分自编码器（VAE）实现。它的任务是读入一帧图像，把它编码成一个低维潜变量 z，然后再从 z 重建出图像。关键在于，这个 z 就是智能体对“当前世界状态”的抽象理解。我们不需要让 z 保留所有像素信息，只需要保留能重建出画面、且对预测下一步有帮助的信息。VAE 天然适合这个任务，因为它强迫潜变量服从某种先验分布，让 z 的每个维度都有明确的概率解释。

我实际训练时最常踩的坑是潜变量维度的选择。维度太小，比如只有 8 维，模型为了强行压缩信息会把很多关键细节抹掉，重建画面糊成一片；维度太大，比如 256 维，理论上表达能力更强，但动力学模型要预测的维度也变多了，反而更难训。我常用的起点是 32 维，对应一个 64×64 的灰度或 RGB 图像，实测下来性价比最高。

VAE 的训练目标包含两项：重建损失和 KL 散度。重建损失保证 z 能还原图像，KL 散度保证 z 的分布在结构上是良好的。这里有个容易被忽略的细节：KL 项的权重需要单独调。权重设得太大，模型会走捷径，把所有输入都编码成同一个 z，因为这样 KL 散度最低，但完全忽略了输入信息，这就是著名的“后验坍缩”问题。我的经验是先让 KL 权重保持很小，比如 0.01，等重建损失降到平台期，再慢慢调大权重。

2.2 记忆与动力学模块：让“脑内模拟”成为可能

有了压缩后的潜变量 z，下一步就要学习“世界的动态规律”，也就是 M 模块。M 模块接收当前隐状态 h 和当前动作 a，输出对下一个潜变量 z 的概率分布。最经典的实现是 MDN-RNN，即混合密度网络加循环神经网络。

为什么不用简单的 LSTM 直接输出一个确定性的 z？因为真实环境的演化往往是多模态的。举个例子，你开车到一个十字路口，既可能直行也可能转弯，单靠一个高斯分布很难描述这种不确定性，如果把两种可能平均成一个中间值，那既不是直行也不是转弯，预测就废了。MDN 的做法是让网络输出多个高斯分布的混合，每个高斯有自己的均值和方差，同时输出一个权重系数，表示每种情况的概率。这样模型可以说“有 60% 概率左转，40% 概率直行”，这是对真实未来更诚实的描述。

在训练 M 模块时，输入的是 VAE 编码出来的真实 z 序列，输出是对下一个 z 的预测。预测目标不需要是精确的 z_t，而是整个分布。训练用的损失是负对数似然，也就是让真实出现的 z 在预测分布下概率尽量高。这里我特别想提醒一件事：如果只用单步预测损失来训练 M，模型往往只能学会“预测下一步”，等到实际使用时要连续 rollout 几十步，误差会像滚雪球一样累积。我通常会在训练时随机截取长度为 10 到 100 的序列，让模型在多个时间尺度上做预测，虽然训练时间会长一些，但长期稳定性明显更好。

M 模块是整个 World Model 里最像“世界”的部分。它不关心像素长什么样，只关心世界在抽象的语义空间里如何演化。很多性能瓶颈也出在 M 模块：VAE 压缩得再好，如果 M 预测不准，想象空间就是空中楼阁。所以我在项目里会把大部分调试时间花在 M 上，而不是 V 上。

2.3 控制器：在想象中学会行动

最后是 C 模块，也就是策略网络。它的输入是当前潜变量 z 和记忆隐状态 h，输出是动作 a。在训练过程中，智能体并不和真实环境交互，而是在 M 模块构建的“梦境”里做决策。具体说，从初始状态出发，C 根据当前 z 和 h 输出动作，M 根据动作和当前状态预测下一时刻 z 的分布，智能体从这个分布中采样一个 z 作为新状态，然后重复这个过程。这样一条完整的想象轨迹，连同奖励一起，就是训练策略的数据。

策略的训练可以用常见的强化学习算法，比如 PPO，也可以用进化算法 CMA-ES。我第一次实现时用的是 CMA-ES，因为它几乎不需要调超参数，只需要定义奖励函数就行。在想象空间里跑几百条轨迹，评估累计奖励，然后用进化策略更新控制器参数，循环往复。这一步的速度非常快，因为纯神经网络推理要比环境模拟快几个数量级。

这里有一个初学者很难理解的点：C 模块输入的 z 是低维潜变量，不是原始图像，为什么部署到真实环境时也能直接用？因为 VAE 在训练时已经建立了像素到 z 的编码器，真实环境里先用 V 编码器把图像变成 z，再喂给 C，整条链路是一致的。所以只要 V 模块能处理真实分布的数据，C 就能正常工作。这也是 World Model 天然适合 sim-to-real 的原因之一。

三大模块之间的关系可以用下表概括：

这套架构最让我佩服的地方在于“分工明确”：V 负责压缩感知，M 负责预测动态，C 负责决策推理，每一块都能独立调试、独立替换。如果换了一个新环境，我通常只需要重新训练 V 和 M，C 可以继承大部分基础策略，这在实际项目里省了很多时间。

3. 从零训练一个 World Model 的实操流程

3.1 环境选择和数据收集

要练手，不建议一上来就选太复杂的环境。DQN 时代大家爱用 Atari，但 Atari 的画面噪声和任务复杂度对 World Model 来说都不算友好。我推荐从自定义的 2D 迷宫或者 OpenAI Gym 里的 CarRacing 开始，这两种环境状态清晰、奖励明确、可视化方便，适合验证模型是否真的学到了规律。

数据收集阶段的核心原则是“覆盖面要广”。用一个随机策略或者带探索噪声的策略，在环境里跑足够多的步数，把图像、动作、奖励全部存下来。这个阶段不需要任何智能策略，因为目标是让 VAE 见过足够多样的画面，同时让 RNN 见过足够多样的状态转移模式。我以前犯过错误：为了省时间，只收集了 2 万步数据，结果 VAE 在常见场景里表现很好，但一旦遇到训练集中少见的转角画面就彻底崩掉。一般 5 万到 10 万步是一个比较稳妥的起点，数据量取决于环境的复杂程度，宁可多存一点，也不要后面回头补数据。

收集到的数据需要保存成统一的格式，我习惯存成 HDF5 文件，里面有图像数组、动作数组、奖励数组。图像统一 resize 到 64×64，方便 VAE 处理。注意动作空间需要归一化到 [-1, 1]，这会直接影响 M 模块的训练稳定性。

3.2 VAE 训练的关键参数

VAE 训练可以照搬我常用的这组参数，实测下来稳定性很好：

• 输入图像：64×64×3，像素归一化到 [0, 1]
• 编码器：4 层卷积，通道数 32/64/128/256
• 潜变量维度：32
• KL 权重：0.01（先小后调）
• 优化器：Adam，学习率 3e-4
• Batch size：128
• Epoch：50 到 100，视重建损失是否收敛而定

训练过程中我会关注两个指标：重建损失和 KL 散度。如果重建损失下降很慢，而 KL 散度几乎为 0，说明模型陷入了后验坍缩，此时把 KL 权重调小是最快的解法。如果重建损失很低，但 KL 散度非常大，说明潜变量维度利用不够充分，可以考虑把维度再加大一点。

另外，我对 VAE 的 decoder 结构也特意做了调整。与其用一层反卷积直接输出整张图，不如在 decoder 中间加一层小尺寸的 feature map，这样重建图像的边缘细节会更清楚。这个小改动不复杂，但对后续 M 模块的训练很有帮助，因为 z 向量里如果能保留更细的空间信息，M 做预测的时候就不容易丢失关键线索。

3.3 MDN-RNN 训练和想象策略训练

VAE 训好后，先把所有训练图像编码成 z 序列，再拿 z 序列训练 M 模块。M 模块我用的配置是：LSTM 隐层 256，混合高斯组分数量设为 5。训练时随机从数据里截取长度 50 的序列，每一步输入 (h, a, z)，预测目标设为下一个 z 的分布。损失函数是负对数似然，这里要注意，z 的分布是 VAE 编码器输出的是均值和对数方差，M 模型预测的也是 z 的分布，但两者要独立处理，不能混为一谈。

M 训练完成后，进入“梦境训练”阶段。这个阶段的代码其实非常简洁：

1. 取一个初始状态 z_0，初始隐状态 h_0 为零向量；
2. 控制器 C 根据 (z_t, h_t) 输出动作 a_t；
3. M 模型根据 (h_t, a_t, z_t) 输出下一时刻 z 的分布，从中采样出 z_{t+1}；
4. 把 (z_{t+1}, a_t) 反馈给控制器，继续循环；
5. 每一步从环境奖励模型中拿到奖励，累计起来作为这条想象轨迹的总回报。

控制器我用的是 PPO，因为它对超参数的敏感性较低，而且能直接利用想象轨迹做多轮更新。PPO 里我特别关注的是 GAE（广义优势估计）的 lambda 值，一般设成 0.95。这个值太小，优势估计会偏短视；太大，方差又会上升。在想象空间里训练还有个好处：每轮可以抽几千条平行轨迹，并行度很高，一张中端 GPU 就能在几分钟内跑完一轮策略迭代。

整个流程训完，我自己习惯做一个验证：把训练好的控制器接回真实环境，让它跑几个 episode，把真实轨迹和想象轨迹做对比。如果真实轨迹的累计奖励和想象轨迹相差很大，说明 VAE 或 M 模块对真实环境的拟合还不够好，需要补数据或者调参。这个验证不能省，它是判断 World Model 是否真正学到了环境规律的重要依据。

4. 当代 World Model 的进化方向：从 Dreamer 到扩散模型

4.1 Dreamer 系列：把想象训练做到极致

MDN-RNN 那套架构虽然经典，但在复杂环境中训练容易不稳定，因为它把“压缩”和“动态预测”分成了两个独立的训练阶段。后来的 Dreamer 系列把这两者统一到一个端到端的框架里，通过 RSSM（循环状态空间模型）学习潜变量动态，再用一个单独的价值网络在预测轨迹上做策略优化。

RSSM 最大的改进是引入两种潜变量：确定性变量和随机变量。确定性变量由 RNN 维护，负责记住长期依赖；随机变量负责描述当前时刻的不确定性。这两个变量结合后，世界模型在长期记忆和随机性建模上都要比 MDN-RNN 平衡。我在一个仿真机械臂环境里试过，Dreamer V2 的训练步数只有传统 World Model 的三分之一，就能达到同样的成功率。

Dreamer 系列另一个值得学习的地方是它对奖励的建模。它除了预测 next latent，还单独训练一个奖励预测器，用同一个潜变量作为输入，输出期望奖励的分布。这样在想象轨迹里，每一步的奖励都能算出来，策略优化就不需要依赖环境真实奖励了。这个技巧让整个训练流程更加端到端。

我个人的感受是：如果做研究，建议先把 MDN-RNN 这套经典实现吃透，因为它简单、透明、好调试；如果做应用，建议直接上 Dreamer 系，比如 Dreamer V3，它在稀疏奖励、多任务迁移方面表现更稳，社区复现也比较成熟。

4.2 扩散 World Model：未来不是单点而是分布

传统的 World Model 在预测下一状态时，输出的是某个参数化分布，比如高斯混合，但这类分布的表达能力有限。现实世界的未来充满多模态，比如一个行人可能向左走也可能向右走，真正诚实的模型应该输出一个多峰的、甚至无法用参数化分布描述的概率。这正是扩散模型擅长的地方。

扩散 World Model 的思路是：直接用扩散模型来建模未来潜变量或者未来帧的分布。它的训练方式和图像扩散模型类似，前向过程逐步加噪，逆向过程学习去噪，只不过条件中加入了当前状态和动作。在采样未来时，从随机噪声开始，用学习到的去噪网络逐步还原出未来帧，这样输出的未来不是被平均掉的模糊结果，而是真实可能发生的多种情况之一。

这个方向的优势非常明显：它彻底摆脱了高斯分布的限制，能表达任意复杂的未来分布。代价是采样速度极慢，一次未来帧的生成往往需要几十步去噪过程，比一次 RNN 推理慢了几个数量级。所以在实际项目里，我一般只在需要精准预测罕见事件时才用扩散 World Model，比如自动驾驶里预测行人突然横穿马路这类场景，普通 RNN 很容易把这种小概率事件平滑掉，扩散模型反而保留住了这些长尾情况。

4.3 World Model 的下一个战场

现在业界对 World Model 的关注点越来越多地落在“通用世界模型”上。所谓的通用，是指不针对某一个具体环境，而是希望模型从大量多模态数据中学会通用的物理规律和交互规则。比如视频数据天然就包含了物体运动、遮挡、碰撞等物理信息，如果一个大模型能从海量视频里学到这些规律，它就可以像人的常识一样，在没有真实交互过的情况下预测物体的行为。

这个方向有一个著名的思路叫“联合嵌入预测”，核心做法是让模型预测未来观测在表示空间中的 embedding，而不是直接重建未来帧。它不要求模型重构出像素级的未来，只要求模型对“未来状态”有正确的语义理解。这种设计的好处是任务难度大幅降低，模型的泛化能力更强，但代价是解释性变差——我们不知道模型内部到底理解和预测了什么。

从我这些年做模型的经验看，World Model 的下一个突破口大概率不是模型结构的大改，而是训练数据的大规模扩充。结构上我们已经有 RNN、Transformer、扩散模型这些强大的组件，缺的往往是足够多样、足够接近真实物理规律的数据。谁能在数据这条路上走得远，谁的 World Model 就能先把“世界”这两个字真正立起来。

5. 训练 World Model 踩过的坑和排查列表

5.1 重建模糊与潜变量坍缩

这是新手最容易遇到的头号问题。症状非常典型：VAE 重建出来的图像模糊得像是隔了一层毛玻璃，不同输入图像重建出来的结果几乎一模一样。出现这种情况，不用怀疑，基本就是后验坍缩，也就是 decoder 忽略了输入，直接用一个通用的先验去生成图像。

排查顺序我建议这样走：第一步看 KL 权重是不是太大，先降到 0.0001 甚至 0，看看重建损失能不能降下来；第二步看潜变量维度是不是太窄，32 维不行就试 64 维；第三步检查编码器最后一层的激活函数，很多坍缩问题是因为编码器输出全落在激活函数的饱和区。这里要注意，KL 权重调小后，潜变量空间会变得不规则，但我们后面还要用 M 模块在 z 空间里做预测，所以也不能只顾重建不管 KL，需要在两个损失之间找到平衡。

我在实际项目里找到的平衡点是：先用很小的 KL 权重把重建损失压到目标值，再逐步把 KL 权重涨到想要的数值。整个过程类似课程学习，前几十个 epoch 重点学编码，后面再开始规整潜空间。这个“先重建、后规整”的顺序，帮我避开了大多数坍缩问题。

5.2 长序列预测漂移

M 模块在短时间步内预测得很准，但一旦想象的序列超过 20 步，预测的状态就开始飘，甚至直接飞出合理范围。这个问题的根源是误差累积：单步预测误差虽然小，但每一步都会作为下一步的输入，误差呈指数级增长。这就像你闭着眼睛走路，每一步方向偏一点点，走几十步后就完全找不到北了。

解决思路有三个层面。第一，训练 M 模块时，不要只做单步预测，而是随机预测多步，让模型习惯“从自身预测中继续预测”。第二，RNN 的输出分布不能只取均值，必须在每次 rollout 时从预测分布中重新采样，这样能避免模型输出一个四平八稳但完全无效的中间状态。第三，在损失函数里加入隐状态的约束项，比如限制 h 的范数，防止长期 rollout 后隐状态爆炸。

还有一个我从炼丹中发现的细节：训练 M 时不要总使用真实 z 作为输入，也就是不要一直用 teacher forcing。适当用一部分模型自己预测出来的 z 作为后续输入，让模型在自己犯的错误上学习修正。比例可以从 10% 慢慢涨到 50%，这个办法对长序列稳定性提升非常明显。

5.3 想象与真实分布的偏移

这是 World Model 落地时绕不过去的一道坎：在想象里跑得很好的策略，拿到真实环境里就失灵了。原因很直接，World Model 是在有限数据上训练的，它对真实世界理解的只是“数据里的世界”，而不是“真实的世界”。一旦真实环境出现训练数据里没有的状态，模型就会开始胡言乱语。

应对方法里，最实用的就是域随机化。在采集训练数据时，刻意变化环境的颜色、光照、物理参数，比如迷宫的墙色随机改、摩擦力随机改。这样训练出的 World Model 对分布偏移会有更强的容忍度。另一个思路是让控制器学会“不确定性感知”，控制器如果发现当前状态对应的预测分布方差特别大，就刻意采取保守动作，而不是冒险冲锋。

我在做机械臂仿真到真机迁移时，还发现一个容易被忽略的地方：动作延迟。仿真环境里动作立即生效，但真实物理世界有通信延迟、执行器惯性，导致真实轨迹和想象轨迹对不上。解决办法是在训练时给动作加一点随机延迟，让 World Model 提前适应这种时间差。这个改动几乎没有成本，但对 sim-to-real 的稳定性贡献极大。

5.4 常见问题速查表

我把实际开发中遇到的典型问题整理成下面这张表，方便排查：

这张表不是标准答案，但它覆盖了我自己项目里九成以上的翻车现场。每次你最怀疑模型结构有问题的时候，先别急着换网络，把上面这些参数挨个捋一遍，往往能省下好几个晚上的调试时间。

最后分享一个我坚持了很久的小习惯：每次训练到一个阶段，我会让想象空间自己跑一段，录制一条视频，看看里面有没有完整的物体轨迹变化，而不仅仅是盯着曲线。曲线上的数字只能告诉你“模型在学没学”，视频才能告诉你“模型学到的东西能不能用”。在 World Model 这类生成式模型里，视觉化的直觉观察往往比所有指标加起来都更能发现问题。我见过不少项目，loss 一路下降，测试指标也像模像样，但把想象视频打开一看，里面全是一堆有纹理的随机噪点，这就是典型的“学了个寂寞”。所以，多让模型说，少只看数字，这条经验对我来说一直很受用。