AI Research 技术调研

World Models (Recurrent World Models Facilitate Policy Evolution)

14分钟阅读

一句话定位

把智能体拆成「大世界模型（Vision + Memory）+ 极小控制器（Controller）」三件套：用 VAE 把每帧像素压成低维隐向量 $z$，用 MDN-RNN 建模 $P(z_{t+1}\mid a_t,z_t,h_t)$ 学时序动力学，再用进化算法 CMA-ES 只训练一个几百参数的线性控制器；关键实验证明可以让控制器完全在世界模型生成的「梦境」里训练、再零样本迁回真实环境（VizDoom Take Cover），并为深度 RL 语境明确命名了 “world model” 这一术语。是把 Schmidhuber 1990 年代 C–M 系统 + Learning to Think（2015）思想用现代 VAE/RNN 工具复现落地的奠基性工作（NeurIPS 2018 正式题名 Recurrent World Models Facilitate Policy Evolution）。

背景与定位

彼时深度 RL 的主流是 model-free（DQN/A3C），受信用分配问题约束，往往只能用 $10^3$–$10^6$ 参数的小网络快速迭代出策略，无法有效训练百万级权重的大模型。本文提出的破法是分工：把绝大部分容量/参数放进一个无监督训练的大「世界模型」（$\sim10^7$ 参数），让它只做一件事——压缩并预测观测序列；而把接触奖励、需要靠 RL/进化求解的「控制器」做到极小（几百参数），从而把信用分配这一难题局限在极小搜索空间里。

思想谱系上，作者明确表示本文不是综述，而是把 Schmidhuber 1990–2015 的一系列 RNN 世界模型 + 控制器工作（Making the World Differentiable 1990、1990b/1991a、On Learning to Think 2015）用现代概率建模工具（VAE + MDN-RNN + ES）“蒸馏”成一个可实验验证的简化框架，术语与记号沿用 Learning to Think。与后来的连续世界模型范式相比：本文的 M 是 LSTM+混合密度网络（离散混合模式），而非扩散/DiT；这条 latent-imagination 路线后来被 dreamer-v3 系列（潜空间想象 + actor-critic）系统化，也是 cosmos-predict 等「视频生成即世界模型」工作的思想前身。与 PILCO（GP 动力学）、Bayesian NN 动力学等低维状态方法不同，本文直接处理高维原始像素输入。

范式名：model-based RL / learned latent dynamics + policy-in-dream（在学到的动力学模型内部训练策略）。

模型架构

三组件 V（Vision）+ M（Memory）+ C（Controller），分开训练、各司其职（论文 Fig.4/Fig.8）：

V — Convolutional VAE（视觉压缩）

  • 输入每帧 resize 到 64×64×3（RGB，各通道 0–1 浮点）。
  • 编码器 = 4 层卷积，输出 $\mu,\sigma\in\mathbb{R}^{N_z}$；隐向量 $z\sim\mathcal{N}(\mu,\sigma^2 I)$（对角高斯先验）；解码器 = 4 层反卷积重建。
  • 每个卷积/反卷积 stride=2，除输出层外均用 ReLU（输出层需落在 0–1）。
  • 隐维 $N_z$：CarRacing = 32，Doom = 64。高斯先验限制每帧信息容量，也让世界模型对 M 产生的不真实 $z$ 更鲁棒。

M — MDN-RNN（时序记忆 / 动力学）

  • LSTM + Mixture Density Network 输出层，建模下一隐向量的概率密度 $P(z_{t+1}\mid a_t,z_t,h_t)$，输出高斯混合参数而非确定性预测（环境随机 → 概率建模）。
  • 5 个高斯混合分量；只输出对角协方差（factored Gaussian），不建模相关系数 $\rho$。
  • LSTM 隐单元：CarRacing = 256，Doom = 512。
  • 采样时用温度 $\tau$ 调控不确定性（借鉴 SketchRNN）——训练控制器时此参数是关键旋钮。
  • Doom 版 M 额外预测 done（$d_t$，是否死亡）：$P(z_{t+1},d_{t+1}\mid a_t,z_t,h_t)$，death 概率 >50% 即置 done=true（用截断而非 Bernoulli 采样，更稳）。

C — Controller（决策）

  • 极简单层线性：$a_t = W_c,[z_t,h_t] + b_c$，输入是 $z_t$ 与 M 隐状态 $h_t$ 的拼接。
  • CarRacing 用 LSTM 的输出向量 $h$；Doom 同时用 cell 向量 $c$ 与输出 $h$。
  • 用 tanh 把动作 clip 到区间：CarRacing 三连续动作（转向 −11、油门 01、刹车 01）；Doom 把离散动作映到 −11 连续空间再三等分（左/不动/右）。

参数量（论文 Tab.，是全文最”硬”的规模数字）：

模型CarRacingDoom (Take Cover)
VAE (V)4,348,5474,446,915
MDN-RNN (M)422,3681,678,785
Controller (C)8671,088

即世界模型 $\sim10^7$ 量级，控制器仅几百参数——正是 CMA-ES 适用的规模。

数据

  • CarRacing-v0 / VizDoom Take Cover 各采集 10,000 条随机策略 rollout（agent 随机动作探索，记录动作 $a_t$ 与观测帧）。这批数据用于训练 V 与 M；世界模型全程不接触真实奖励信号，只压缩并预测观测序列，奖励只给 C。
  • 无外部大规模语料——纯粹是两个 OpenAI Gym 环境自采的像素+动作序列（sim 数据），单环境单任务。
  • 训练 M 时用 teacher forcing：对每帧预存 $\mu,\sigma$，每次构 batch 时重新采样 $z\sim\mathcal{N}(\mu,\sigma)$ 输入，避免 MDN-RNN 过拟合到某个特定采样 $z$。
  • 迭代式数据收集（论文第 5 节，仅作为方向讨论、未在本文实验落地）：更难的环境需要「rollout 真实环境 → 存动作/观测 → 重训 M、在 M 内训 C → 回到采集」的循环，并可把 M 的预测损失翻号当好奇心/内在动机信号驱动探索。

训练方法

分三段、各自独立优化（论文 §3.2 步骤）：

  1. 采集 10,000 条随机 rollout。
  2. 训 V（VAE）：把帧编码进 $z$；损失 = 重建 L2 距离 + KL 损失；只训 1 个 epoch。
  3. 训 M（MDN-RNN）：在 V 预处理出的 $z$ + 记录动作上，最大化 $P(z_{t+1}\mid a_t,z_t,h_t)$ 的似然；训 20 个 epoch。
  4. 进化 C：用 CMA-ES（Covariance-Matrix Adaptation Evolution Strategy）搜 $W_c,b_c$ 最大化期望累积回报。种群 64，每个个体跑 16 次不同随机种子，适应度 = 16 次平均累积回报。CMA-ES 适用于至多几千参数的解空间，且只需最终累积回报、天然易并行。
  • V、M、C 分开训练比端到端更实际（作者也指出原则上可端到端）。
  • 在梦里训 C（VizDoom）：把训练好的 M 包一层 gym.Env 接口当成虚拟环境，C 只在纯 latent 空间里训练（无需 V 解码真实像素），学好后直接部署回真实 VizDoom。
  • 温度 $\tau$ 正则化对抗”钻空子”：C 会发现对抗性策略去利用 M 的缺陷（如让梦里怪物永不发射火球）。用 MDN 的随机性 + 调高 $\tau$ 注入更多不确定性，使梦境更难、更接近真实，抑制 C 利用模型瑕疵；$\tau=0.1$ 时会模式崩溃（怪物不发火球），学到的策略在真实环境里比随机策略还差。

Infra（训练 / 推理工程）

  • 硬件（来自 worldmodels.github.io 项目页披露，比 arXiv 更具体）：实验在 一块 P100 GPU + 64 核 CPU 的 Ubuntu Linux 虚拟机（Google Cloud Platform）上完成，框架 TensorFlow + OpenAI Gym。
  • V、M 各自训练 < 1 小时（单 GPU），且无需大量调参即可得到可用的 VAE / MDN-RNN。
  • C 的进化在单机多 CPU 核上并行跑多个 rollout（64 核机器，每进程配 16 次 rollout）。
  • CarRacing 收敛：约 1800 代后，最优个体在 1024 次随机 rollout 上取得平均 900.46 分（用 1024 而非 100 次评估，因为 64 核 × 16 次刚好一整代算力）。
  • 推理侧：世界模型作为可微循环计算图，梦境模拟可 GPU 加速；DoomRNN 只在 latent 空间运行、不渲染像素、不跑真实 Doom 引擎，比真实 VizDoom 更省算力。全部模型甚至能用 deeplearn.js 部署进浏览器做交互 demo（p5.js + distill.pub）。控制/推理 FPS、延迟等未披露具体数字。

评测 benchmark

CarRacing-v0（solve 门槛 = 100 次连续 trial 平均 ≥ 900；论文 Tab.1，100 次随机 trial）：

方法平均分
DQN (Prieur 2017)343 ± 18
A3C (continuous)591 ± 45
A3C (discrete)652 ± 10
ceobillionaire（Gym 榜首）838 ± 11
本文 V-only（$a_t=W_c z_t+b_c$）632 ± 251
本文 V-only + 隐层788 ± 141
本文 Full World Model（V+M）906 ± 21

→ Full World Model 首次解出该任务并刷新 SOTA（此前深度 RL 方法多在 591–652，榜首 838±11）；且本文直接吃原始 RGB 像素，不做边缘检测/堆帧等预处理。V-only 只看 $z_t$ 时开车会”抖”（错过急弯），加上 $h_t$（含未来分布信息）后驾驶明显更稳——说明 M 提供的时序特征关键。

VizDoom Take Cover（solve 门槛 = 100 次平均存活 > 750 步；每 rollout 上限 2100 步 ≈ 60 秒）：

  • 在梦境（DoomRNN，$\tau=1.15$）里训练，虚拟环境得分 $\sim$900 步；最优个体虚拟 1024 次平均 959。
  • 迁回真实 VizDoom：100 次平均 $\sim$1100 步（1092 ± 556），远超 750 门槛，也高于虚拟环境分（更难的”噩梦”里练出的策略在更干净的真实环境里更强）。Gym 榜首为 820 ± 58，随机策略 210 ± 108。

温度消融（论文 Tab.2，梦里训练后回真实环境 100 次平均）：

$\tau$虚拟分真实分
0.102086 ± 140193 ± 58
0.502060 ± 277196 ± 50
1.001145 ± 690868 ± 511
1.15918 ± 5461092 ± 556
1.30732 ± 269753 ± 139

→ $\tau$ 太低（0.1/0.5）时 C 在梦里刷到 $\sim$2086 但迁移几乎全崩（193），因为它钻了确定化模型的空子；$\tau=1.15$ 是虚拟难度与迁移性的最佳折中。

创新点与影响

  • 术语与范式奠基：在现代深度 RL 语境里明确提出并命名 “world model”（V+M+C 架构），把 Schmidhuber 1990s C–M 系统 + Learning to Think(2015) 用 VAE+MDN-RNN+ES 复现落地。
  • 首个「完全在梦里训练、再迁回真实」的成功演示（VizDoom Take Cover）：证明只要世界模型足够准，可以用它替代真实环境来训策略；并指出对抗性钻空子问题及其温度正则解法。
  • 大模型/小控制器分工 + 进化训练控制器：把信用分配困难局限在几百参数的小搜索空间，同时保留大世界模型的表达力；首次解出 CarRacing-v0。
  • 作者自陈的局限：① VAE 独立训练会编码任务无关细节（重建了 Doom 墙上无关砖纹，却漏了 CarRacing 路面关键 tile），与 M 联训预测奖励可缓解但会牺牲跨任务复用性；② 世界模型容量有限（LSTM 权重存不下所有历史，易灾难性遗忘）；③ 逐步模拟未来、无人类式分层规划/抽象推理（更一般的 Learning to Think 允许 C 调用 M 的”子程序”、忽略无用的 M，本文未做）；④ 简单任务一次迭代即可，复杂任务需迭代式探索——均留作 future work（并提到 One Big Net(2018) 把 C 与 M 合并、用行为回放防遗忘的后续方向）。

原始链接

一手源存档（sources/）

  • world-models-ha-schmidhuber—project-page — worldmodels.github.io 交互文章全文快照（含 P100+64核 CPU / TensorFlow infra 披露、附录 config、温度表）
  • arXiv 原文 PDF（https://arxiv.org/pdf/1803.10122，arXiv 原文 PDF，不入 git）

关系图谱

反向链接