mlsanityData to 7 Oct 2026

Improving Token-Based World Models with Parallel Observation Prediction

arXiv 2402.05643 ↗First posted 8 Feb 2024cs.LG cs.AIcited by 11

On the map: World models → Latent dynamics models and model-based RL

Abstract (arXiv)

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 urlhttps://github.com/leor-c/REM.

Key facts. Cited by 11 papers in the map (data through 7 Oct 2026). Placed in Latent dynamics models and model-based RL, one of 11 directions in World models. Builds on 29 of its 39 references that are in the map. First posted on arXiv 8 Feb 2024.

Builds on

29 of 39 references.

Closest ancestors · 8

Same direction, newest first

  1. 2023STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning14 Oct
  2. 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
  3. 2023Mastering Diverse Domains through World Models10 Jan
  4. 2022Transformers are Sample-Efficient World Models1 Sep
  5. 2021Mastering Atari Games with Limited Data30 Oct
  6. 2020Mastering Atari with Discrete World Models5 Oct
  7. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec
  8. 2018Recurrent World Models Facilitate Policy Evolution4 Sep

Other references · 21

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2017Proximal Policy Optimization Algorithmscited by 14,441
  5. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 12,865
  6. 2016Layer Normalizationcited by 5,579
  7. 2017Neural Discrete Representation Learningcited by 4,650
  8. 2020Taming Transformers for High-Resolution Image Synthesiscited by 3,229
  9. 2023Sparks of Artificial General Intelligence: Early experiments with GPT-4cited by 2,579
  10. 2019Transformer-XL: Attentive Language Models Beyond a Fixed-Length Contextcited by 1,830
  11. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
  12. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
9 more
  1. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,221
  2. 2019Dota 2 with Large Scale Deep Reinforcement Learningcited by 1,084
  3. 2022A Generalist Agentcited by 802
  4. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 781
  5. 2023Retentive Network: A Successor to Transformer for Large Language Modelscited by 594
  6. 2022Perceiver-Actor: A Multi-Task Transformer for Robotic Manipulationcited by 545
  7. 2019Model-Based Reinforcement Learning for Ataricited by 495
  8. 2019Stabilizing Transformers for Reinforcement Learningcited by 262
  9. 2022MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesiscited by 222

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 1; 2025: 8; 2026: 2.

11 papers.

  1. 2025Improving Transformer World Models for Data-Efficient RLcited by 13
  2. 2025Dyn-O: Building Structured World Models with Object-Centric Representationscited by 12
  3. 2024Sable: a Performant, Efficient and Scalable Sequence Model for MARLcited by 10
  4. 2025DMWM: Dual-Mind World Model with Long-Term Imaginationcited by 10
  5. 2025Object-Centric World Models from Few-Shot Annotations for Sample-Efficient Reinforcement Learningcited by 9
  6. 2025Simulus: Combining Improvements in Sample-Efficient World Model Agentscited by 6
  7. 2025EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modelingcited by 3
  8. 2025DyMoDreamer: World Modeling with Dynamic Modulationcited by 2
  9. 2026Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Modelscited by 1
  10. 2026OrbiSim: World Models as Differentiable Physics Engines for Embodied Intelligencecited by 1
  11. 2025A Survey of Retentive Networkcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026Identifiable Token Correspondence for World Models
  2. 2024Masked Generative Priors Improve World Models Sequence Modelling Capabilities
  3. 2026Token-World: World Modeling in Vision-Language Model Token Space for Robot Manipulation
  4. 2026Improving Weak World Models Behind Strong Agents in Atari Pong
  5. 2023Facing Off World Model Backbones: RNNs, Transformers, and S4
  6. 2024Efficient World Models with Context-Aware Tokenization
  7. 2024Open-World Reinforcement Learning over Long Short-Term Imagination
  8. 2024Learning to Play Atari in a World of Tokens