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}.

翻译：受Transformer在离散符号序列处理中成功应用的启发，基于令牌的世界模型（TBWMs）近期被提出作为一种样本高效的方法。在TBWMs中，世界模型将智能体的经验视为类似语言的令牌序列，每个观测构成一个子序列。然而，在想象过程中，逐令牌顺序生成下一个观测会导致严重的计算瓶颈，表现为训练时间长、GPU利用率低以及表征能力受限。为解决此瓶颈，我们设计了一种新颖的并行观测预测（POP）机制。POP为保留网络（RetNet）配备了一种针对我们强化学习场景定制的前向模式。我们将POP集成到名为REM（保留环境模型）的新型TBWM智能体中，相比先前TBWMs实现了15.4倍的想象速度提升。在Atari 100K基准测试的26个游戏中，REM在12小时内完成训练，并在其中12个游戏中达到超人性能。我们的代码开源在\url{https://github.com/leor-c/REM}。