Improving Token-Based World Models with Parallel Observation Prediction
This addresses a bottleneck in sample-efficient reinforcement learning for AI agents, though it is an incremental improvement over existing token-based world models.
The paper tackled the slow sequential token generation in token-based world models by introducing a Parallel Observation Prediction (POP) mechanism, resulting in a 15.4x faster imagination and superhuman performance on 12 out of 26 Atari 100K games while training in under 12 hours.
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}.