[论文解读] Transformer-based World Models Are Happy With 100k Interactions

Jan Robine, Marc Höftmann+2 more|arXiv (Cornell University)|Mar 13, 2023

Reinforcement Learning in Robotics被引用 4

一句话总结

该论文提出了一种基于Transformer的世界模型（TWM），仅使用100,000次环境交互即在Atari 100k基准上实现了最先进性能。通过使用Transformer-XL架构自回归地建模观测、动作和预测奖励，并将预测奖励反馈至模型，TWM实现了样本高效的策略学习，其性能优于以往的无模型和基于模型的方法，同时保持了较低的推理计算成本。

ABSTRACT

Deep neural networks have been successful in many reinforcement learning settings. However, compared to human learners they are overly data hungry. To build a sample-efficient world model, we apply a transformer to real-world episodes in an autoregressive manner: not only the compact latent states and the taken actions but also the experienced or predicted rewards are fed into the transformer, so that it can attend flexibly to all three modalities at different time steps. The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state. By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient. Our transformer-based world model (TWM) generates meaningful, new experience, which is used to train a policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark.

研究动机与目标

• 解决深度强化学习的样本效率低下问题，后者通常需要数十亿次环境交互，远超人类学习者所需。
• 通过自注意力机制实现对过去状态的直接、长距离注意力，克服递归世界模型的局限性。
• 通过在潜在想象中使用生成式世界模型合成新经验来训练策略，提升样本效率。
• 通过新颖的损失函数和采样策略稳定训练，提升泛化能力和超参数鲁棒性。
• 证明基于Transformer的世界模型可在极少真实环境交互（具体为Atari上的100,000次交互）下实现优异性能。

提出的方法

• 使用CNN将高维图像观测编码为潜在状态，再将其嵌入Transformer输入序列。
• 将潜在状态、动作和预测奖励作为输入标记输入Transformer，使模型能够灵活地跨三种模态进行注意力计算。
• 采用带相对位置编码的Transformer-XL架构，高效建模长期依赖关系，并降低推理阶段的计算成本。
• 利用多层感知机（MLPs）基于Transformer隐藏状态预测下一潜在状态、奖励和折扣因子。
• 引入平衡KL散度损失，微调世界模型目标函数中熵项与交叉熵项之间的权衡。
• 应用截断熵损失以稳定训练过程中的策略熵，降低对超参数选择的敏感性。
• 实施平衡采样策略，优先选择近期经验，以纠正训练数据集中的数据分布偏移。

实验结果

研究问题

• RQ1 在仅使用100,000次环境交互进行训练的情况下，基于Transformer的世界模型是否能在Atari 100k基准上超越现有的无模型和基于模型的强化学习方法？
• RQ2 将预测奖励反馈至Transformer输入对世界模型质量及下游策略性能有何影响？
• RQ3 与均匀采样相比，所提出的平衡采样策略在多大程度上提升了训练稳定性和性能？
• RQ4 截断熵损失是否能稳定策略训练并降低在多种Atari游戏中对超参数的敏感性？
• RQ5 使用Transformer-XL架构是否能在保持低推理成本的同时，实现比基于RNN的世界模型更优的长期依赖建模能力？

主要发现

• 所提出的TWM在Atari 100k基准上实现了人类归一化得分的最先进水平，优于以往的无模型和基于模型方法。
• 在消融研究中显示，将预测奖励作为条件输入至Transformer可显著提升大多数游戏的性能。
• 平衡采样策略通过将关注点转向更近期经验，减轻了数据分布偏移，从而提升了策略性能。
• 截断熵损失稳定了训练过程中的策略熵，简化了超参数调优，并提升了跨游戏的泛化能力。
• 平衡KL散度损失实现了对熵与交叉熵权衡的精细控制，改善了世界模型的学习效果。
• 该模型仅通过100,000次真实环境交互即实现高性能，证明了其在样本效率和推理计算效率方面的优势，原因在于策略推理过程中无需使用Transformer。