用对比编码提升Transformer世界模型,让智能体在游戏上表现超越此前方法。
Learning Transformer-based World Models with Contrastive Predictive Coding
- 用动作条件的对比预测编码,让Transformer学习更长时序特征。
- 在Atari 100k上达162%人类基准分,刷新无前瞻搜索方法纪录。
- 适合研究模型预测与强化学习融合的学者,尤其关注高效训练的场景。
DreamerV3算法通过基于循环神经网络(RNN)的世界模型,在多个环境域中取得了优异性能。随着基于模型的强化学习的成功以及Transformer架构在训练效率和可扩展性上的优势,近期工作如STORM尝试用掩码自注意力机制替代RNN构建基于Transformer的世界模型。然而,尽管训练效率提高,其性能仍不及Dreamer算法,难以学习到具备竞争力的Transformer世界模型。本文指出,先前方法采用的下一状态预测目标不足以充分挖掘Transformer的表征能力。为此,我们提出TWISTER(基于Transformer的世界模型,具有对比表示),通过引入动作条件的对比预测编码,学习高层时间特征表示,显著提升智能体性能。TWISTER在Atari 100k基准上达到162%的人类标准化平均得分,成为不使用前瞻搜索的最先进方法中的新纪录。
原文摘要 · Abstract (English)
The DreamerV3 algorithm recently obtained remarkable performance across diverse environment domains by learning an accurate world model based on Recurrent Neural Networks (RNNs). Following the success of model-based reinforcement learning algorithms and the rapid adoption of the Transformer architecture for its superior training efficiency and favorable scaling properties, recent works such as STORM have proposed replacing RNN-based world models with Transformer-based world models using masked self-attention. However, despite the improved training efficiency of these methods, their impact on performance remains limited compared to the Dreamer algorithm, struggling to learn competitive Transformer-based world models. In this work, we show that the next state prediction objective adopted in previous approaches is insufficient to fully exploit the representation capabilities of Transformers. We propose to extend world model predictions to longer time horizons by introducing TWISTER (Transformer-based World model wIth contraSTivE Representations), a world model using action-conditioned Contrastive Predictive Coding to learn high-level temporal feature representations and improve the agent performance. TWISTER achieves a human-normalized mean score of 162% on the Atari 100k benchmark, setting a new record among state-of-the-art methods that do not employ look-ahead search.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。