# Improving Transformer World Models for Data-Efficient RL

arXiv: https://arxiv.org/abs/2502.01591 · HTML: https://mlsanity.com/paper/2502.01591/

## Key facts

- Cited by 13 papers in the map (data through 14 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.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

## Builds on (most cited first, 30 of 41)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [SAM 2: Segment Anything in Images and Videos](https://mlsanity.com/paper/2408.00714/) (2024)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Diffusion for World Modeling: Visual Details Matter in Atari](https://mlsanity.com/paper/2405.12399/) (2024)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Benchmarking the Spectrum of Agent Capabilities](https://mlsanity.com/paper/2109.06780/) (2021)
- [Bigger, Better, Faster: Human-level Atari with human-level efficiency](https://mlsanity.com/paper/2305.19452/) (2023)
- [Stop Regressing: Training Value Functions via Classification for Scalable Deep RL](https://mlsanity.com/paper/2403.03950/) (2024)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [When to use parametric models in reinforcement learning?](https://mlsanity.com/paper/1906.05243/) (2019)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)

## Built on by (most cited first, 13 of 13)

- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling](https://mlsanity.com/paper/2603.04553/) (2026)
- [One Life to Learn: Inferring Symbolic World Models for Stochastic Environments from Unguided Exploration](https://mlsanity.com/paper/2510.12088/) (2025)
- [Simulus: Combining Improvements in Sample-Efficient World Model Agents](https://mlsanity.com/paper/2502.11537/) (2025)
- [Reimagination with Test-time Observation Interventions: Distractor-Robust World Model Predictions for Visual Model Predictive Control](https://mlsanity.com/paper/2506.16565/) (2025)
- [Deep RL Needs Deep Behavior Analysis: Exploring Implicit Planning by Model-Free Agents in Open-Ended Environments](https://mlsanity.com/paper/2506.06981/) (2025)
- [Curious Causality-Seeking Agents Learn Meta Causal World](https://mlsanity.com/paper/2506.23068/) (2025)
- [Morphology-Conditioned World Model for Cross-Embodiment Quadrupedal Locomotion](https://mlsanity.com/paper/2604.08780/) (2026)
- [Preemptive Solving of Future Problems: Multitask Preplay in Humans and Machines](https://mlsanity.com/paper/2507.05561/) (2025)
- [From Observations to Events: Event-Aware World Model for Reinforcement Learning](https://mlsanity.com/paper/2601.19336/) (2026)
- [CrafterDojo: A Suite of Foundation Models for Building Open-Ended Embodied Agents in Crafter](https://mlsanity.com/paper/2508.13530/) (2025)
- [Soft Geometric Inductive Bias for Object Centric Dynamics](https://mlsanity.com/paper/2512.15493/) (2025)
- [MA-WAM: Multi-Agent World-Action Model for Test-Time Planning](https://mlsanity.com/paper/2609.31281/) (2026)

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