ICML 2018 · 理论与优化

World Models

学会在"梦境"中驯化智能体:压缩时空表示 + 虚拟环境训练
David Ha (Google Brain) & Jürgen Schmidhuber (NNAISENSE, IDSIA) · 2018

本文提出将智能体分解为感知(VAE)、记忆(MDN-RNN)与控制(CMA-ES 线性层)三个模块。 世界模型在无监督条件下快速学习环境的压缩时空表示;控制器则可完全在世界模型生成的"梦境"中训练, 再将策略迁移回真实环境,首次在 CarRacing-v0VizDoom: Take Cover 两个任务上达到 SOTA。

ICML 2018 VAE + MDN-RNN + CMA-ES model-based RL 📄 arXiv:1803.10122 Project / Interactive Demo
world models 世界模型 VAE MDN-RNN CMA-ES model-based RL 梦境学习 dream environment

01 动机

现实中的 model-free 强化学习方法通常只使用参数量较少的小型网络, 因为信用分配问题(credit assignment problem)使得在百万参数量的大模型上进行直接 RL 训练极为困难。 人类依靠内部"心理模型"(mental model)做出快速决策,无需在脑中显式规划所有未来情景—— 本文受此启发,构建能预测未来的世界模型,并以此为"舞台"训练极简的控制策略。

"We explore building generative neural network models of popular reinforcement learning environments. Our world model can be trained quickly in an unsupervised manner to learn a compressed spatial and temporal representation of the environment. By using features extracted from the world model as inputs to an agent, we can train a very compact and simple policy that can solve the required task. We can even train our agent entirely inside of its own hallucinated dream generated by its world model, and transfer this policy back into the actual environment."
World Models overview
世界模型的核心思路:在真实环境中收集数据训练 RNN-based 世界模型,之后以该模型完整替代真实环境,让智能体在"虚拟梦境"中学习策略,最终迁移回真实世界。图中展示了 CarRacing 和 VizDoom 两个实验的任务示意。
906±21CarRacing-v0 最终得分(满分 900 为 "solved")
1092±556VizDoom Take Cover 最佳生存步数(τ=1.15)
867Controller 参数量(线性层,极度紧凑)
10,000初始随机 rollouts 用于训练世界模型

02 方法

智能体由三个模块组成:V(VAE)将每帧图像压缩为低维隐向量 zM(MDN-RNN)对未来 z 建立概率预测模型; C(Controller)是一个只有数百参数的线性层,以 [z, h] 为输入直接映射到动作。 V 和 M 在无监督条件下训练,C 则用进化策略(CMA-ES)优化,互不干扰信用分配。

Agent architecture overview
智能体总览图:V(Vision)负责单帧空间压缩,M(Memory)负责时序预测,C(Controller)做出动作决策。三者协同工作:V 编码当前帧得到 z,M 维护隐状态 h,C 以 [z, h] 为输入输出动作 a。

V 模型:Variational Autoencoder (VAE)

环境每步给出高维像素观测,V 模型的任务是将 64×64 RGB 图像编码为低维隐向量 z。 使用卷积 VAE(4 层卷积 + 4 层反卷积),在 CarRacing 中 Nz=32, 在 VizDoom 中 Nz=64。VAE 的 Gaussian prior 不仅限制了信息容量, 也使得 M 模型生成的"想象帧"更具鲁棒性,不易因输入超出分布而崩溃。

VAE diagram
VAE(Variational Autoencoder)流程图:输入图像经卷积编码器得到均值 μ 和方差 σ,从 Gaussian prior N(μ, σI) 采样得到隐向量 z,再经反卷积解码重建图像。训练目标为重建误差(L² loss)+ KL 散度。

M 模型:MDN-RNN(Mixture Density Network + RNN)

M 模型的目标是对时序动态建模,预测下一帧的隐向量分布: P(zt+1 | at, zt, ht)。 使用 LSTM 作为 RNN 主体,输出层为 Mixture Density Network(混合高斯分布), 通过调节温度参数 τ 控制采样随机性。MDN 的离散模式(discrete modes)对于 建模具有随机离散事件的环境(如怪物是否发射火球)尤为重要。 在 VizDoom 实验中,M 额外预测 done 信号,使得世界模型成为完整的 RL 环境接口。

MDN-RNN diagram
MDN-RNN 架构图:RNN 接收当前 z 和动作 a,隐状态 h 经 Mixture Density Network 输出层产生混合高斯分布参数(权重 π,均值 μ,方差 σ),用于采样预测下一帧隐向量 z。

C 模型:Controller(控制器)

控制器是单层线性模型:at = Wc [zt ht] + bc。 CarRacing 中 C 只有 867 个参数,VizDoom 中为 1,088 个。 参数量如此之少,使得进化策略 CMA-ES 能够在多核 CPU 上高效并行搜索最优策略, 完全回避了深度 RL 中的信用分配难题。

Full agent flow diagram
完整智能体数据流:原始观测 → V 编码 → zt → C 输出动作 at → 环境更新 → M 以 (zt, at) 更新隐状态 ht+1。所有操作均在一步完成,无需显式规划未来轨迹。

在梦境中训练(Dream Environment)

VizDoom 实验中,M 被封装为符合 OpenAI Gym 接口的虚拟环境(gym.Env)。 智能体在此"梦境"中完全以隐空间 z 为观测进行训练,无需调用真实游戏引擎。 训练完成后,将策略直接迁移回真实 VizDoom 场景。整个训练流程:

  1. 随机 rollout 收集 10,000 条轨迹数据。
  2. 训练 VAE 将帧编码为 z ∈ ℝ64,并转换数据集。
  3. 训练 MDN-RNN 建模 P(zt+1, dt+1 | at, zt, ht)。
  4. 在虚拟梦境环境中用 CMA-ES 优化 C 的参数。
  5. 将学到的策略部署回真实环境评估。

03 实验

论文在两个环境中验证:CarRacing-v0(连续控制,随机生成赛道) 和 VizDoom: Take Cover(躲避火球,生存时间越长越好)。 两个环境均从 10,000 条随机 rollout 出发,无需人工设计特征或奖励塑形。

实验一:CarRacing-v0

任务定义为平均分 ≥ 900(100 次连续试验)即视为 "solved"。V 模型在不含奖励信息的情况下训练, M 模型仅做帧预测,仅 C 接触奖励信号。结果表明,仅使用 V 时得分 632±251,加入 M 后升至 906±21, 首次在公开结果中"解决"该任务。

方法平均得分
DQN343 ± 18
A3C (continuous)591 ± 45
A3C (discrete)652 ± 10
ceobillionaire (Gym Leaderboard)838 ± 11
V model only632 ± 251
V model + hidden layer788 ± 141
Full World Model (V + M + C)906 ± 21

实验二:VizDoom Take Cover — 在梦境中训练,迁移至真实

任务要求智能体在 2100 步内尽量存活,得分超过 750 视为通过。 关键消融:温度参数 τ 控制虚拟环境随机性。低温(τ=0.1/0.5)导致模式坍塌, 怪物永不发射火球,策略在真实环境中几乎失效(得分 ~196); τ=1.15 获得最佳迁移效果(实际得分 1092±556);τ 过高(1.30)则梦境过难,学习效果下降。

温度 τ梦境得分(Virtual)真实得分(Actual)
0.102086 ± 140193 ± 58
0.502060 ± 277196 ± 50
1.001145 ± 690868 ± 511
1.15918 ± 5461092 ± 556
1.30732 ± 269753 ± 139
Random PolicyN/A210 ± 108
Gym Leaderboard BestN/A820 ± 58
VizDoom policy transfer
在梦境 RNN 环境中训练的策略被成功迁移回真实 VizDoom 场景。智能体仅在隐空间中接受训练,从未直接处理真实游戏帧,但在真实环境中表现超越 OpenAI Gym 排行榜最佳成绩。

模型规模对比

世界模型(V + M)承载了绝大部分复杂度,而 Controller 刻意保持极简以便 CMA-ES 优化。 CarRacing 中三者参数量如下:

模块CarRacing 参数量VizDoom 参数量
VAE (V)4,348,5474,446,915
MDN-RNN (M)422,3681,678,785
Controller (C)8671,088

04 局限性

说明:以下局限性均来自论文 Discussion 节作者的明确陈述,未做推断性扩展。
对抗性策略利用世界模型缺陷(Cheating the World Model)

由于世界模型只是真实环境的近似,Controller 可以发现"漏洞": 在 VizDoom 实验初期,智能体学会了一种对抗性策略,移动方式使怪物在整个 rollout 过程中从不发射火球,在虚拟环境中得分极高,但在真实环境中完全失效。 论文通过提高温度参数 τ 缓解此问题,但并未根本解决。

VAE 编码与任务相关性不足(Task-irrelevant Feature Encoding)

VAE 在无监督条件下训练,无法知晓哪些特征对任务重要。 论文指出:"it reproduced unimportant detailed brick tile patterns on the side walls in the Doom environment, but failed to reproduce task-relevant tiles on the road in the Car Racing environment." 将 VAE 与奖励预测联合训练或许有助于对焦任务相关区域,但代价是失去跨任务复用能力。

世界模型容量有限,存在灾难性遗忘(Limited Capacity & Catastrophic Forgetting)

LSTM-based 世界模型的权重存储容量有限,无法随数据量线性扩展。 作者指出可能需要引入更大容量模型(Mixture of Experts、Hypernetworks 等) 或外部记忆模块来应对更复杂的环境,同时也面临传统神经网络的灾难性遗忘问题。

缺乏层次规划与抽象推理(No Hierarchical Planning)

当前方法逐步模拟未来(step-by-step),未能实现人类式的层次化规划与抽象推理。 论文承认这一点,并指向 Learning to Think 框架作为更通用但未在此实现的方向。