Improving Token-Based World Models with Parallel Observation Prediction

TL;DR

This work tackles the inefficiency of token-based world models in reinforcement learning by addressing the imagination bottleneck that arises from autoregressively generating next observations. It introduces Parallel Observation Prediction (POP), a forward-mode extension to RetNet that enables generating entire next-observation token sequences in parallel, significantly accelerating imagination. The authors instantiate REM (Retentive Environment Model), a TBWM agent that uses a VQ-VAE tokenizer, a token-space world model, and an actor-critic controller trained entirely in imagination, achieving $15.4\times$ faster imagination and superhuman performance on 12 of 26 Atari 100K games, in under 12 hours. The approach demonstrates that TBWMs can be both fast and competitive, with broader potential for longer sequences and frame-level generation in future work.

Abstract

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

Figures (15)

• Figure 1: Top: comparison between the run times of token-based world model agents (IRIS and REM) during the world model training and imagination (actor-critic training). Bottom: interquantile mean (IQM) human-normalized score comparison between REM and state-of-the-art baselines on the Atari 100K benchmark with 95% stratified bootstrap confidence intervals Agarwal2021rliable. A line separates token-based methods from other baselines.
• Figure 2: An overview of REM's training cycle. Each epoch has 4 steps: experience collection (1), tokenizer training (2), world model training (3), and controller training in imagination (4). Orange color represents component(s) that undergo training. Blue squares denote token inputs, where light blue is used for observation tokens and dark blue for actions. Replay buffer data at steps 3 and 4 contains observations, actions, rewards, and termination signals.
• Figure 3: The "chunkwise" computation mode. Long sequences can be split into smaller "chunks" for enhanced training efficiency. Previous chunks are summarized by the recurrent state $\mathbf{S}$. Blue squares represents tokens, while circles denote output vectors. Crucially, RetNet's chunkwise mode does not natively support both a batched generation of tokens at imagination and an efficient world model training. These are achieved by our POP extension.
• Figure 4: A single imagination step. Starting from a recurrent state $\mathbf{S}_{t}$, initially obtained from real experience, $\mathcal{M}$ computes all next-observation tokens $\hat{\mathbf{z}}_{t+1}$ in parallel using the prediction tokens $\mathbf{u}$ as inputs. Then, the agent observes $\hat{\mathbf{z}}_{t+1}$ and picks an action $a_{t+1}$. Finally, $\mathcal{M}$ takes $\mathbf{S}_{t}, \hat{\mathbf{z}}_{t+1},$ and $a_{t+1}$ and outputs $\mathbf{S}_{t+1}, \hat{r}_{t+1}, \hat{d}_{t+1}$. Dashed arrows emphasize sampling operations.
• Figure 5: When appending $\mathbf{u}$ after an observation-action block, the sequence is no longer a prefix of the observation-action token trajectory. Thus, the recurrent state only summarizes observation and action tokens (top trajectory).