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通过为强化学习场景定制的前向模式对Retentive Network（RetNet）进行增强。我们将POP集成到名为REM（Retentive Environment Model）的新型TBWM智能体中，相较于现有TBWMs实现了15.4倍的想象生成加速。在Atari 100K基准测试的26个游戏中，REM在12小时内完成训练并在其中12个任务上达到超人类表现。我们的代码开源在\url{https://github.com/leor-c/REM}。