[论文解读] TransDreamer: Reinforcement Learning with Transformer World Models

Chang Chen, Yifu Wu+2 more|arXiv (Cornell University)|Feb 19, 2022

Reinforcement Learning in Robotics被引用 8

一句话总结

TransDreamer 引入了一种完全基于 Transformer 的基于模型强化学习智能体，用一种新型的 Transformer 状态空间模型（TSSM）替代了 Dreamer 中的循环世界模型，从而实现了更优的长程记忆与规划能力。它在需要长期推理的复杂 2D 和 3D 视觉控制任务上优于 Dreamer，同时在短期依赖性较强的简单任务上性能与 Dreamer 相当。

ABSTRACT

The Dreamer agent provides various benefits of Model-Based Reinforcement Learning (MBRL) such as sample efficiency, reusable knowledge, and safe planning. However, its world model and policy networks inherit the limitations of recurrent neural networks and thus an important question is how an MBRL framework can benefit from the recent advances of transformers and what the challenges are in doing so. In this paper, we propose a transformer-based MBRL agent, called TransDreamer. We first introduce the Transformer State-Space Model, a world model that leverages a transformer for dynamics predictions. We then share this world model with a transformer-based policy network and obtain stability in training a transformer-based RL agent. In experiments, we apply the proposed model to 2D visual RL and 3D first-person visual RL tasks both requiring long-range memory access for memory-based reasoning. We show that the proposed model outperforms Dreamer in these complex tasks.

研究动机与目标

• 探究基于 Transformer 的世界模型是否能提升基于模型强化学习中的样本效率与长期推理能力。
• 解决在 MBRL 中训练稳定基于 Transformer 的策略的挑战，此前研究表明从稀疏奖励中直接训练纯 Transformer 策略存在困难。
• 开发一种兼容 Transformer 的随机、动作条件化世界模型，支持并行训练与高效的想象能力。
• 评估改进后的世界模型是否能提升需要复杂长程记忆交互任务的性能。
• 验证所提出的框架在简单任务上的强性能表现，这些任务预期 RNN 基模型能够成功。

提出的方法

• 提出 Transformer 状态空间模型（TSSM），一种随机、动作条件化的世界模型，利用自注意力机制捕捉潜在空间中的长程依赖关系。
• 调整 TSSM 以支持先验与后验推理，用于随机动力学预测，从而在想象中实现准确的轨迹滚动。
• 在策略与世界模型组件之间共享 TSSM 世界模型，实现具有统一潜在表征的一体化训练。
• 采用变分推理目标（ELBO）训练 TSSM，结合随机潜在状态与确定性隐藏状态，并通过多头自注意力机制进行更新。
• 使用基于 Transformer 的演员-评论家策略，基于 TSSM 生成的想象轨迹进行训练，实现样本高效的强化学习。
• 将 TSSM 集成到 Dreamer 框架中，替代 RSSM，同时保留整体 MBRL 训练流程。

实验结果

研究问题

• RQ1 基于 Transformer 的世界模型是否能在需要复杂记忆交互的长时程视觉控制任务中超越循环模型？
• RQ2 构建稳定、端到端的基于 Transformer 的 MBRL 智能体所面临的关键架构与训练挑战是什么？
• RQ3 基于 Transformer 的世界模型是否能提升想象轨迹的质量，特别是在图像生成与奖励预测方面？
• RQ4 在具有短期依赖关系的任务上，基于 Transformer 的 MBRL 智能体与基于循环模型的基线（Dreamer）相比表现如何？
• RQ5 TSSM 世界模型是否能在多样化视觉环境中保持训练稳定性与计算效率的同时实现良好泛化？

主要发现

• 在 5-Ball Dense 环境中，TransDreamer 超越了 Dreamer，该 2D 视觉控制任务复杂，需要长期记忆与精确的球序推理能力。
• TransDreamer 生成的想象帧更清晰、更准确，并能正确预测在收集小球事件时的奖励峰值，而 Dreamer 生成的图像模糊且奖励预测错误。
• TSSM 世界模型在后续想象步骤中对未来的帧预测与奖励估计均优于 Dreamer 的 RSSM，尤其在长程依赖关系起关键作用时表现更优。
• 在短期记忆任务（如 DMC Cheetah Run 和 Atari 游戏）中，TransDreamer 达到了与 Dreamer 相当的最终性能，尽管收敛速度较慢。
• TransDreamer 展现出更高的奖励预测准确率与更优的视觉想象保真度，表明其世界模型比 Dreamer 的 RNN 基基线更鲁棒、更精确。
• TSSM 实现了完全基于 Transformer 的强化学习智能体的稳定训练，证明当采用恰当的架构设计时，基于 Transformer 的 MBRL 是可行且高效的。