文 / Google Research 学生研究员 Danijar Hafner
得益于深度强化学习 (RL)，人工智能体能够随着时间的推移不断改进其决策。传统的无模型方法与环境交互，通过大量试错来学习预测不同情况下的成功动作。相较之下，深度 RL 的最新进展已经使得 基于模型的方法 能够从图像输入中学习精确的 世界模型，并将其用于规划。这些世界模型可以 从更少的交互中学习，推动离线数据的泛化，实现前瞻性探索，并允许 在多个任务之间 重复使用知识。
尽管现有的世界模型（如 SimPLe）拥有诱人的优势，但其仍不够精确，无法在最具竞争力的强化学习基准上与性能最佳的无模型方法相媲美：迄今为止，在完善的 Atari 基准测试 中，我们需要使用 DQN、IQN 和 Rainbow 等无模型算法，才能达到人类的水准。因此，许多研究人员转而专注于开发 特定于任务 的规划方法，如 VPN 和 MuZero，这些方法通过预测预期任务奖励的总和进行学习。然而，这些方法均针对个别任务，目前尚不清楚其能在多大程度上推广到新任务或直接从无监督数据集中学习。与计算机视觉中无监督表征学习的最新突破 [1、2] 类似，世界模型旨在学习环境中比任何特定任务更普遍的模式，以便日后更有效地解决问题、完成任务。
今天，我们与 DeepMind 和多伦多大学合作，推出了首款基于世界模型的 RL 智能体 DreamerV2，其可在 Atari 基准测试中达到人类水准。DreamerV2 是第二代 Dreamer 智能体，仅在由像素级数据训练得到的世界模型中的隐空间内学习行为。DreamerV2 完全依赖于图像中的一般信息，并且即使其表征不受未来任务奖励的影响，也能精确预测这些奖励。DreamerV2 使用单个 GPU，在相同计算和样本量的情况下，性能优于顶级无模型算法。
经过 2 亿步后，进行游戏的模型在 55 款 Atari 游戏中的归一化得分的中位数。DreamerV2 的性能大大优于之前的世界模型。此外，在相同计算和样本量内，DreamerV2 的表现超过了顶级无模型智能体
DreamerV2 在 55 款 Atari 游戏（部分）中学习的行为。这些视频显示了来自环境的图像。对这些视频的预测则展示在下方的博文中
抽象的世界模型
如同其前身一样，DreamerV2 学习世界模型，并利用该模型完全从预测的轨迹中 训练 actor-critic（演员与评论家）行为。世界模型自动学习计算其图像的紧凑表征，从中发现物体位置等有用的概念，并学习这些概念如何根据不同的动作而改变。得益于此，智能体可以生成能够忽略无关细节的图像抽象，并且在单个 GPU 上即可实现大规模并行预测。在执行 2 亿个环境步骤的过程中，DreamerV2 会预测 4680 亿个紧凑状态以便学习其行为。
DreamerV2 基于循环状态空间模型 (RSSM) 构建，该模型针对 PlaNet 而推出，也被用于 DreamerV1。在训练过程中，编码器会将每幅图像转化为随机表征，并将其纳入世界模型的 循环状态 。这些表征具有随机性，导致智能体无法获得关于图像的完美信息，而只能提取到进行预测所需的信息，因此智能体不太会受到看不见的图像的影响。解码器从每个状态中重构相应的图像以便学习 一般表征。此外，该模型还会训练一个小型奖励网络，以便在规划过程中对结果进行排序。为了能够在不生成图像的情况下进行规划， 预测器 会学习在不获取要进行计算图像的情况下猜测随机表征。
DreamerV2 使用的世界模型的学习过程。该世界模型保持接收动作 (a1–a2) 的循环状态 (h1–h3)，并通过随机表征 (z1–z3) 纳入图像 (x1–x3) 的信息。预测器猜测表征为 (ẑ1–ẑ3)，而无需访问从中生成这些表征的图像
重要的是，DreamerV2 在 RSSM 中引入了两项新技术，从而为学习成功的策略提供了精确度大幅提升的世界模型。第一项技术是使用多个分类变量，而不是 PlaNet、DreamerV1 和文献 [1、2、3] 中许多其他世界模型所使用的高斯变量，来表示每幅图像。借此技术，世界模型能够根据离散概念来推理世界，并能够更精确地预测未来的表征。
编码器将每幅图像转换成 32 种分布，每种分布包含 32 个类别，其含义会随着世界模型的学习而自动确定。从这些分布中取样的独热向量会连接到一个 稀疏表征 ，并传递给循环状态。为了通过样本进行反向传播，我们使用了易于使用自动微分来实现的直通式梯度。如果使用分类变量来表示图像，则预测器可精确学习下一幅可能图像的独热向量上的分布情况。相较之下，使用高斯预测器的早期世界模型就无法精确匹配下一幅可能图像在多个高斯表征上的分布情况。
分类预测器可精确预测表示可能下一幅图像的多个分类，而高斯预测器则不够灵活，无法精确预测多个可能的高斯表征
DreamerV2 的第二项新技术是 KL 平衡。许多以前的世界模型均使用 ELBO 目标函数，该目标函数在保持随机表征（后验）接近其预测（先验）的同时，鼓励进行精准重构，以规范从每幅图像中提取的信息量，并促进泛化。该目标函数针对端到端优化，因此可以将随机表征及其预测这两者中的任意一者向另一者靠拢，以使两者更加相似。然而，当预测器尚不精确时，使表征向其预测靠拢可能会出现问题。有了 KL 平衡，预测向表征移动的速度比表征向预测移动的速度更快。这就会带来更精确的预测，这是成功规划的关键。
世界模型针对保持序列的长期视频预测。每个模型接收 5 帧画面作为输入（未显示），然后在只有动作的情况下即可预测未来 45 步。视频预测仅用于了解世界模型的质量。在规划期间，仅预测紧凑表征，而不预测图像
衡量 Atari 性能
DreamerV2 是能够在成熟且具有竞争力的 Atari 基准测试中，学习成功行为且达到人类水准的首款世界模型。我们选择了之前许多研究都曾使用过的 55 款游戏，并推荐以后的研究也使用这组游戏。按照标准评估协议，智能体可每个动作重复 4 次，并使用粘性动作（忽略一个动作而重复上一个动作的几率为 25%），从而可与环境交互 2 亿次。我们将 DreamerV2 与性能最佳的无模型智能体 IQN 和 Rainbow，以及在 多巴胺框架 中实现的著名 C51和 DQN 智能体等模型进行对比。
对 55 款游戏的得分进行汇总时，存在不同的标准。理想情况下，新算法在所有条件下均会表现得更好。在所有这四种汇总方法中，在使用相同计算量的情况下，DreamerV2 性能确实优于所有与之比较的无模型算法。
在对 55 款 Atari 游戏得分进行汇总的四种方法中，DreamerV2 性能优于顶级无模型智能体。我们引入并推荐 修剪记录平均值 （最右图）作为内容丰富且鲁棒性强的性能指标
2.1 进行游戏的模型的得分中位数
进行游戏的模型的得分中位数。最常见的情况是，每款游戏的得分都根据为 DQN 论文评估的人类玩家的表现进行归一化，并报告所有游戏归一化得分的中位数。可惜的是，中位数会忽略许多更简单和更困难游戏的得分。
2.2 进行游戏的模型的得分平均值
平均值将所有游戏的得分纳入考量，但在很大程度上会受人类玩家表现不佳的少数游戏的影响。这会导致算法很容易在某些游戏（如 James Bond、Video Pinball）中取得较高的归一化分数，然后左右平均值。
2.3 记录平均值
记录平均值。之前的研究建议基于人类世界纪录进行归一化，但这样的指标仍然受到少数游戏的过度影响，在这些游戏中，人工智能体的得分很容易超过人类纪录。
2.4 进行游戏的模型的得分中位数
我们引入了一种新的指标，该指标可以根据世界纪录对得分进行归一化，并将其修剪为不超过纪录。这就产生了一个信息全面且稳健性强的指标，将所有游戏中的表现在大致相等的程度上进行考量。
虽然目前很多算法都超过了人类玩家的基线，但仍然远远落后于人类的世界记录。如上方最右图所示，DreamerV2 在所有游戏中平均领先人类纪录的 25%。如果把得分修剪为记录线，我们就可以集中精力开发在所有游戏中均能更接近人类世界纪录的方法，而不是仅仅关注在少数游戏中超过人类世界记录。
重要事项和非重要事项
为深入了解 DreamerV2 的重要组成部分，我们进行了广泛的消融研究。重要的是，我们发现，尽管之前的研究已广泛使用高斯表征，但与之相比，分类表征具有明显的优势。与大多数生成模型所使用的 KL 正则化器相比，KL 平衡具有更大优势。
通过防止图像重构或奖励预测梯度影响模型状态，我们研究了这两者对学习成功表征的重要性。我们发现，DreamerV2 完全依赖于来自高维输入图像的通用信息， 即使未使用奖励信息对其进行训练，其表征也能实现精确的奖励预测。这与无监督表征学习在计算机视觉领域所取得的成功相同。
DreamerV2 各种消融的 Atari 性能（修剪记录平均值）。分类表征、KL 平衡和对图像的学习是 DreamerV2 成功的关键。使用特定于狭义任务的奖励信息，不会对学习世界模型带来其他好处
结论
我们展示了如何学习强大的世界模型，以在具有竞争力的 Atari 基准测试中达到人类水准，并超越顶级无模型智能体。研究结果表明，世界模型是在强化学习问题上实现高性能的有力方法，并且可供从业者和研究人员使用。我们认为这表明世界模型成功在强化学习中复制了无监督表征学习在计算机视觉中所取得的成功。Github 上提供了 DreamerV2 的非官方实现，为未来的研究项目提供了富有成效的起点。我们认为，利用大型离线数据集、长期记忆、分层规划和定向探索的世界模型是开展未来研究的令人兴奋的途径。
致谢
该项目由 Timothy Lillicrap、Mohammad Norouzi 和 Jimmy Ba 合作完成。此外，我们还要感谢在整个项目期间，对我们的论文草稿发表意见并随时提供反馈的 Google Brain 团队全体成员。
原文：Mastering Atari with Discrete World Models
中文：TensorFlow 公众号