[论文解读] Mastering Atari with Discrete World Models

Danijar Hafner, Timothy Lillicrap+2 more|arXiv (Cornell University)|Oct 5, 2020

Reinforcement Learning in Robotics参考文献 74被引用 23

一句话总结

DreamerV2 是一种基于模型的强化学习智能体，通过在独立训练的世界模型内纯化地学习行为，利用离散潜在表征，在 Atari 55 个游戏基准上实现了人类水平的性能。在相同的计算资源和实际运行时间下，其性能优于顶尖的单 GPU 模型无关智能体（如 Rainbow 和 IQN），证明了精确的世界模型能够实现从像素中样本高效、高性能的决策。

ABSTRACT

Intelligent agents need to generalize from past experience to achieve goals in complex environments. World models facilitate such generalization and allow learning behaviors from imagined outcomes to increase sample-efficiency. While learning world models from image inputs has recently become feasible for some tasks, modeling Atari games accurately enough to derive successful behaviors has remained an open challenge for many years. We introduce DreamerV2, a reinforcement learning agent that learns behaviors purely from predictions in the compact latent space of a powerful world model. The world model uses discrete representations and is trained separately from the policy. DreamerV2 constitutes the first agent that achieves human-level performance on the Atari benchmark of 55 tasks by learning behaviors inside a separately trained world model. With the same computational budget and wall-clock time, Dreamer V2 reaches 200M frames and surpasses the final performance of the top single-GPU agents IQN and Rainbow. DreamerV2 is also applicable to tasks with continuous actions, where it learns an accurate world model of a complex humanoid robot and solves stand-up and walking from only pixel inputs.

研究动机与目标

• 开发一种基于模型的强化学习智能体，仅通过在学习到的世界模型内进行想象的轨迹 rollout，实现在 Atari 基准上的人类水平性能。
• 证明在策略训练之外独立训练的世界模型，可在复杂视觉环境中支持高精度规划与泛化能力。
• 仅使用单张 GPU 和 10 天训练时间，实现与最先进模型无关智能体（如 Rainbow 和 IQN）相当的性能。
• 验证离散潜在空间表征是否能提升视觉控制任务中世界模型的准确性和样本效率。
• 建立一个可复现、高效且可扩展的基于模型强化学习框架，适用于离散与连续控制任务。

提出的方法

• 世界模型使用离散变分自编码器（dVAE）从原始像素输入中学习紧凑的类别型潜在表征。
• 智能体在离散潜在空间中学习一个转移模型，根据当前状态和动作预测未来状态。
• 在世界模型生成的想象轨迹上，使用价值函数和通过 KL 平衡损失优化的策略网络，独立训练一个演员-评论家网络。
• 策略通过内在好奇心与外在奖励预测的结合进行训练，世界模型提供规划的环境。
• 在策略训练前，先在大规模环境交互数据集上预训练世界模型，从而实现样本高效的策略学习。
• 算法采用一种新颖的 KL 平衡机制，以稳定训练并提升离散潜在空间中的表征质量。

实验结果

研究问题

• RQ1 在复杂视觉环境（如 Atari）中，能否通过与策略分离训练的世界模型实现样本高效、高性能的决策？
• RQ2 与连续潜在空间相比，使用离散潜在表征是否能提升世界模型的准确性和泛化能力？
• RQ3 基于模型的智能体能否在策略训练期间不与环境直接交互，而在完整 Atari 55 个游戏基准上实现人类水平性能？
• RQ4 在相同计算约束下，基于模型的智能体与最先进模型无关智能体（如 Rainbow 和 IQN）的性能相比如何？
• RQ5 世界模型在图像和奖励预测方面的能力，在多大程度上促进了策略的泛化与样本效率？

主要发现

• DreamerV2 在完整的 Atari 55 个游戏基准上实现了人类水平性能，成为首个仅通过在世界模型内进行想象 rollout 实现该目标的基于模型智能体。
• 在相同的计算预算和实际运行时间下，DreamerV2 在 2 亿帧评估中超越了顶尖单 GPU 模型无关智能体 Rainbow 和 IQN 的最终性能。
• 使用离散潜在表征显著提升了世界模型的准确性，并在潜在空间中实现了稳定、高质量的规划。
• DreamerV2 表明，通过图像重建和奖励预测信号训练的世界模型，能够学习到对多样化控制任务有用的通用表征。
• 该智能体在连续控制任务中也表现出色，仅使用像素输入和单张 GPU，即成功解决了人形机器人复杂的运动行为（如站立和行走）。
• 通过消融研究，实证验证了离散潜在空间和 KL 平衡机制带来的性能增益，确认了其对最终智能体性能的显著影响。