Improving Token-Based World Models with Parallel Observation Prediction

📄 arXiv: 2402.05643v5 📥 PDF

作者: Lior Cohen, Kaixin Wang, Bingyi Kang, Shie Mannor

分类: cs.LG, cs.AI

发布日期: 2024-02-08 (更新: 2024-05-29)

🔗 代码/项目: GITHUB

💡 一句话要点

提出并行观察预测机制以提升基于令牌的世界模型性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 基于令牌的世界模型 并行观察预测 保留网络 强化学习 Atari基准测试

📋 核心要点

  1. 现有的基于令牌的世界模型在想象过程中存在逐个生成观察的瓶颈，导致训练时间长和GPU利用率低。
  2. 本文提出了一种并行观察预测机制（POP），通过增强保留网络以适应强化学习环境，从而提高生成效率。
  3. 实验结果表明，REM代理在Atari 100K基准测试中表现优异，训练时间少于12小时，且在12个游戏中达到了超人类水平。

📝 摘要（中文）

受Transformer在离散符号序列应用成功的启发，最近提出了基于令牌的世界模型（TBWMs）作为一种样本高效的方法。在TBWMs中，世界模型将代理经验视为类似语言的令牌序列，每个观察结果构成一个子序列。然而，在想象过程中，逐个生成下一个观察结果导致了严重的瓶颈，造成训练时间长、GPU利用率低和表示能力有限。为了解决这一瓶颈，本文提出了一种新颖的并行观察预测（POP）机制。POP通过一种新颖的前向模式增强了保留网络（RetNet），并在一种新型TBWM代理REM（Retentive Environment Model）中实现，展示了相比于先前TBWMs快15.4倍的想象速度。REM在Atari 100K基准测试的26个游戏中有12个达到了超人类表现，同时训练时间少于12小时。代码可在https://github.com/leor-c/REM获取。

🔬 方法详解

问题定义：本文旨在解决基于令牌的世界模型在想象过程中逐个生成观察所带来的效率瓶颈，导致训练时间过长和GPU资源利用不充分的问题。

核心思路：提出并行观察预测（POP）机制，通过增强保留网络（RetNet）以实现观察的并行生成，从而加速训练过程并提升模型的表示能力。

技术框架：整体架构包括一个增强的保留网络，结合POP机制进行观察的并行生成，形成一个新的TBWM代理REM。该框架在强化学习环境中进行训练和评估。

关键创新：最重要的技术创新在于并行观察预测机制的引入，使得观察生成不再是线性序列，而是可以同时生成多个观察，显著提高了效率。

关键设计：在网络结构上，POP机制通过特定的前向模式与保留网络结合，优化了损失函数以适应并行生成的需求，确保了生成质量与速度的平衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，REM代理在想象速度上相比于传统TBWMs提升了15.4倍，并在Atari 100K基准测试中，12个游戏达到了超人类表现，训练时间少于12小时，展现了显著的性能优势。

🎯 应用场景

该研究的潜在应用领域包括强化学习、机器人控制和游戏AI等。通过提升基于令牌的世界模型的效率，能够更快地训练智能体，从而在复杂环境中实现更高效的决策与控制，具有重要的实际价值和未来影响。

📄 摘要（原文）

Motivated by the success of Transformers when applied to sequences of discrete symbols, token-based world models (TBWMs) were recently proposed as sample-efficient methods. In TBWMs, the world model consumes agent experience as a language-like sequence of tokens, where each observation constitutes a sub-sequence. However, during imagination, the sequential token-by-token generation of next observations results in a severe bottleneck, leading to long training times, poor GPU utilization, and limited representations. To resolve this bottleneck, we devise a novel Parallel Observation Prediction (POP) mechanism. POP augments a Retentive Network (RetNet) with a novel forward mode tailored to our reinforcement learning setting. We incorporate POP in a novel TBWM agent named REM (Retentive Environment Model), showcasing a 15.4x faster imagination compared to prior TBWMs. REM attains superhuman performance on 12 out of 26 games of the Atari 100K benchmark, while training in less than 12 hours. Our code is available at \url{https://github.com/leor-c/REM}.