mlsanityData to 7 Oct 2026

Improving Transformer World Models for Data-Efficient RL

arXiv 2502.01591 ↗First posted 3 Feb 2025cs.LG cs.AIcited by 13

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

Abstract (arXiv)

We present three improvements to the standard model-based RL paradigm based on transformers: (a) "Dyna with warmup", which trains the policy on real and imaginary data, but only starts using imaginary data after the world model has been sufficiently trained; (b) "nearest neighbor tokenizer" for image patches, which improves upon previous tokenization schemes, which are needed when using a transformer world model (TWM), by ensuring the code words are static after creation, thus providing a constant target for TWM learning; and (c) "block teacher forcing", which allows the TWM to reason jointly about the future tokens of the next timestep, instead of generating them sequentially. We then show that our method significantly improves upon prior methods in various environments. We mostly focus on the challenging Craftax-classic benchmark, where our method achieves a reward of 69.66% after only 1M environment steps, significantly outperforming DreamerV3, which achieves 53.2%, and exceeding human performance of 65.0% for the first time. We also show preliminary results on Craftax-full, MinAtar, and three different two-player games, to illustrate the generality of the approach.

Key facts. Cited by 13 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 41 of its 53 references that are in the map. First posted on arXiv 3 Feb 2025.

Builds on

41 of 53 references.

Closest ancestors · 15

Same direction, newest first

  1. 2024Efficient World Models with Context-Aware Tokenization27 Jun
  2. 2024Learning to Play Atari in a World of Tokens3 Jun
  3. 2024Diffusion for World Modeling: Visual Details Matter in Atari20 May
  4. 2024Improving Token-Based World Models with Parallel Observation Prediction8 Feb
  5. 2023TD-MPC2: Scalable, Robust World Models for Continuous Control25 Oct
  6. 2023STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning14 Oct
  7. 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
  8. 2023Mastering Diverse Domains through World Models10 Jan
  9. 2022Transformers are Sample-Efficient World Models1 Sep
  10. 2022TransDreamer: Reinforcement Learning with Transformer World Models19 Feb
  11. 2021Mastering Atari Games with Limited Data30 Oct
  12. 2020Mastering Atari with Discrete World Models5 Oct
3 more
  1. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec
  2. 2018Recurrent World Models Facilitate Policy Evolution4 Sep
  3. 2018World Models27 Mar

Other references · 26

  1. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  2. 2017Proximal Policy Optimization Algorithmscited by 14,441
  3. 2023DINOv2: Learning Robust Visual Features without Supervisioncited by 7,764
  4. 2015Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shiftcited by 6,979
  5. 2016Layer Normalizationcited by 5,579
  6. 2017Neural Discrete Representation Learningcited by 4,650
  7. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
  8. 2024SAM 2: Segment Anything in Images and Videoscited by 2,784
  9. 2015High-Dimensional Continuous Control Using Generalized Advantage Estimationcited by 1,916
  10. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
  11. 2018IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architecturescited by 866
  12. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 774
14 more
  1. 2023Retentive Network: A Successor to Transformer for Large Language Modelscited by 594
  2. 2019Model-Based Reinforcement Learning for Ataricited by 495
  3. 2021Benchmarking the Spectrum of Agent Capabilitiescited by 193
  4. 2023Bigger, Better, Faster: Human-level Atari with human-level efficiencycited by 145
  5. 2024Stop Regressing: Training Value Functions via Classification for Scalable Deep RLcited by 136
  6. 2019When to use parametric models in reinforcement learning?cited by 118
  7. 2022Discovered Policy Optimisationcited by 103
  8. 2024Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learningcited by 80
  9. 2024Bridging State and History Representations: Understanding Self-Predictive RLcited by 52
  10. 2022Investigating Compounding Prediction Errors in Learned Dynamics Modelscited by 47
  11. 2023Curious Replay for Model-based Adaptationcited by 16
  12. 2018The Effect of Planning Shape on Dyna-style Planning in High-dimensional State Spacescited by 15
  13. 2023Discovering Hierarchical Achievements in Reinforcement Learning via Contrastive Learningcited by 6
  14. 2022A Look at Value-Based Decision-Time vs. Background Planning Methods Across Different Settingscited by 5

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 9; 2026: 4.

13 papers.

Sort:

By area:

  1. 2025A Comprehensive Survey on World Models for Embodied AIcited by 65
  2. 2026Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modelingcited by 11
  3. 2025One Life to Learn: Inferring Symbolic World Models for Stochastic Environments from Unguided Explorationcited by 9
  4. 2025Simulus: Combining Improvements in Sample-Efficient World Model Agentscited by 6
  5. 2025Reimagination with Test-time Observation Interventions: Distractor-Robust World Model Predictions for Visual Model Predictive Controlcited by 6
  6. 2025Deep RL Needs Deep Behavior Analysis: Exploring Implicit Planning by Model-Free Agents in Open-Ended Environmentscited by 4
  7. 2025Curious Causality-Seeking Agents Learn Meta Causal Worldcited by 4
  8. 2026Morphology-Conditioned World Model for Cross-Embodiment Quadrupedal Locomotioncited by 3
  9. 2025Preemptive Solving of Future Problems: Multitask Preplay in Humans and Machinescited by 2
  10. 2026From Observations to Events: Event-Aware World Model for Reinforcement Learningcited by 2
  11. 2025CrafterDojo: A Suite of Foundation Models for Building Open-Ended Embodied Agents in Craftercited by 1
  12. 2025Soft Geometric Inductive Bias for Object Centric Dynamicscited by 0
1 more
  1. 2026MA-WAM: Multi-Agent World-Action Model for Test-Time Planningcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026Identifiable Token Correspondence for World Models
  2. 2025Similarity-Aware Token Pruning: Your VLM but Faster
  3. 2024Exploring Token Pruning in Vision State Space Models
  4. 2021Learned Token Pruning for Transformers
  5. 2025Equivalence of Context and Parameter Updates in Modern Transformer Blocks
  6. 2023Morphing Tokens Draw Strong Masked Image Models
  7. 2024A Survey on Transformer Compression
  8. 2024Do Efficient Transformers Really Save Computation?