用离散世界模型掌握 Atari

官方源标题： Mastering Atari with Discrete World Models
作者： Danijar Hafner、Timothy Lillicrap、Mohammad Norouzi、Jimmy Ba
发表： ICLR 2021；arXiv:2010.02193（v4，2022-02-12）
主来源： arXiv 摘要与论文正文

核心结论

DreamerV2 的关键变化，是先从真实游戏画面里学出一个紧凑的离散世界模型，再把大量策略训练搬进这个模型的“想象”中。它因此成为论文所称第一个仅在独立训练的世界模型里学习行为、却在 55 个 Atari 游戏上达到人类水平的智能体。摘要；论文 §1、§5

三个最值得记住的结论

1. 先学世界，再练动作。 真实交互负责更新经验库；世界模型负责预测状态、奖励和游戏是否继续；策略与价值网络只在预测出的轨迹上训练。
2. “离散化”不是表面改动。 相比 DreamerV1 的高斯潜变量，DreamerV2 使用 32 个、每个 32 类的离散变量，并用直通梯度训练；消融实验显示，去掉离散状态或 KL 平衡都会明显掉分。
3. 强结果不等于通用智能。 论文对每个游戏分别训练一个智能体，训练一个 2 亿步实验约需单张 V100 GPU 不到 10 天；它还会漏掉只占一个像素、却决定游戏成败的物体。

问题

像素世界为什么难以“想远一点”

直接预测未来画面很昂贵，也容易让误差一步步放大。更麻烦的是，画面里的大量像素对动作并不重要，模型却仍要花容量重建它们。潜空间世界模型的思路，是把画面压成较小的内部状态，只预测行动真正需要的变化；这样既能减少预测负担，也能在一张 GPU 上同时展开很多轨迹。论文 §1、§2.1

旧方法卡在了哪里

在这篇论文之前，图像世界模型已能解决一些连续控制任务，但在竞争激烈的 Atari 基准上还追不上强大的无模型方法。SimPLe 直接预测像素并只在 36 个较容易的游戏子集上评测；MuZero 能在 Atari 上取得强结果，但论文指出其工程和计算预算远高于单卡研究条件。DreamerV2 要回答的问题更具体：一个单 GPU、单环境实例的世界模型智能体，能否在标准 55 游戏、2 亿步设置下与 Rainbow、IQN 等方法正面比较？论文 §1、§4，表 3

方法

第一步：把画面压成可预测的离散状态

世界模型从经验库抽取画面、动作、奖励和终止信息。编码器读取当前画面；循环状态空间模型（RSSM）保留过去信息。每一步都有两个内部状态：后验状态看得到当前画面，先验状态只能根据过去状态和动作猜下一步。模型同时学习重建画面、预测奖励，以及预测游戏是否继续。论文 §2.1，图 2

第二步：离散状态与 KL 平衡各解决一个问题

DreamerV2 把随机状态表示为 32 个分类变量，每个变量有 32 个类别。一次采样可展平成 1024 位稀疏向量，其中 32 位激活。作者确认离散表示更有效，但没有断言唯一原因；多模态变化、稀疏性和优化性质都只是待验证的解释。论文 §2.1、§3.2

KL 项既要教先验预测后验，又会反过来限制后验。如果先验还很差，过早把后验拉向它，内部表示也会一起变差。KL 平衡通过分别停止两边的部分梯度，让训练更偏向“先把先验学准”；论文在 Atari 中使用 0.8 的平衡系数。论文 §2.1；附录 D，表 D.1

第三步：冻结世界模型，在里面批量想象

策略学习时，世界模型保持冻结。演员网络选择动作，模型预测下一状态、奖励和是否继续，评论家网络估计长期回报；梯度只更新演员和评论家。因为不必生成像素，单张 GPU 可同时模拟 2500 条潜空间轨迹。论文 §2.2，图 3

证据与局限

55 个游戏上的结果胜过四个单卡基线

实验采用 55 个 Atari 游戏、粘滞动作、动作重复 4、每局最多 108,000 步、单任务训练和单环境实例。论文把预算写作 2 亿环境步或帧；动作重复后，智能体实际接收 5000 万次画面输入。训练期间，世界模型又提供了 4680 亿个想象状态，约为真实输入数的 10,000 倍。论文 §3，实验设置

DreamerV2 在论文报告的四种汇总指标上都领先。作者更推荐最后一列：先用人类世界纪录归一化，再把超过纪录的得分截到 1，避免少数特别容易刷高分的游戏支配均值。论文 §3.1，表 1

消融证明两个核心改动确实有用

在较早版本的 DreamerV2 上，完整模型的截断世界纪录均值为 0.25；去掉离散潜变量后降到 0.19，去掉 KL 平衡后降到 0.16，切断图像重建梯度后只剩 0.01。按单游戏计，离散状态相对高斯状态在 42 个任务更好、8 个更差、5 个持平；KL 平衡在 44 个任务更好、6 个更差、5 个持平。这支持“两个改动有效”，但没有证明作者列出的机制猜想就是原因。论文 §3.2，表 2、图 5

结果不能回答哪些问题

这不是一个同时掌握 55 个游戏的通用智能体，而是每个游戏训练一个独立模型。完整改动清单也没有做全面消融：作者估算，每个改动若跑 55 个任务、5 个种子、每次 10 天，将超过 60,000 GPU 小时。论文还指出 Video Pinball 是明显弱项，并推测一像素大小的球没有得到足够的重建学习信号；这只是作者假设，不是已验证因果结论。论文 §3.1；附录 C

论文也没有证明世界模型在 Atari 之外普遍优于无模型方法。附录展示了仅从像素解决 Humanoid 站起和行走的结果，说明方法能处理连续动作；但这只是额外任务，不等于跨环境、跨任务迁移已经完成。论文附录 A；§5

实际意义

真正可复用的是“学习环境，再复用想象”

如果真实交互昂贵，而环境又能从历史数据中学到足够准确的动态，DreamerV2 展示了一条有吸引力的路线：用有限真实经历训练世界模型，再把策略搜索放进便宜、可并行的内部模拟中。这一点对机器人、交互式系统和仿真训练有启发，但属于从论文机制推导出的实践解释，不是论文已经完成的产品验证。

上线前先问三个验证问题

• 模型会不会漏掉小而关键的信号？ 不要只看整体重建质量，要专门测试稀有事件、细小目标和终止条件。
• 想象越多，偏差会不会越大？ 需要比较模型预测与真实回放，并监控策略是否利用模型漏洞。
• 计算换来的收益值不值？ 论文的单任务、单卡约 10 天是研究级成本基线；实际系统还要把数据、调参与失败重跑算进去。

术语与核查索引

• 世界模型： 根据过去状态与动作，预测未来状态、奖励和是否继续的模型。
• 潜状态： 画面的紧凑内部表示；DreamerV2 的随机部分由多个分类变量组成。
• 后验 / 先验： 后验看当前画面来编码状态；先验不看当前画面，只凭历史和动作预测。
• 直通梯度： 前向过程使用离散采样，反向过程借用概率的梯度近似更新。
• KL 平衡： 调整 KL 损失两侧的学习力度，优先把预测用的先验学准。

核查入口： 贡献与结论看摘要、§1、§5；完整算法看 §2；Atari 设置与主结果看 §3、表 1；关键消融看 §3.2、表 2；改动与算力限制看附录 C；超参数看附录 D。所有章节均可从论文主来源进入。