# Learning to Play Atari in a World of Tokens

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

## Key facts

- Cited by 6 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 58 of its 73 references that are in the map.
- First posted on arXiv 3 Jun 2024.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [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 Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [RLBench: The Robot Learning Benchmark & Learning Environment](https://mlsanity.com/paper/1909.12271/) (2019)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Reinforcement learning](https://mlsanity.com/paper/2405.10369/) (2024)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Model-Ensemble Trust-Region Policy Optimization](https://mlsanity.com/paper/1802.10592/) (2018)
- [MineRL: A Large-Scale Dataset of Minecraft Demonstrations](https://mlsanity.com/paper/1907.13440/) (2019)
- [Recurrent Memory Transformer](https://mlsanity.com/paper/2207.06881/) (2022)
- [Benchmarking the Spectrum of Agent Capabilities](https://mlsanity.com/paper/2109.06780/) (2021)
- [Masked World Models for Visual Control](https://mlsanity.com/paper/2206.14244/) (2022)
- [Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansion](https://mlsanity.com/paper/1807.01675/) (2018)

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

- [Improving Transformer World Models for Data-Efficient RL](https://mlsanity.com/paper/2502.01591/) (2025)
- [Hadamax Encoding: Elevating Performance in Model-Free Atari](https://mlsanity.com/paper/2505.15345/) (2025)
- [Curious Causality-Seeking Agents Learn Meta Causal World](https://mlsanity.com/paper/2506.23068/) (2025)
- [Probing the Impact of Scale on Data-Efficient, Generalist Transformer World Models for Atari](https://mlsanity.com/paper/2605.08578/) (2026)
- [Understanding and Stabilizing Deep Q-Learning via Controlled Bootstrapping and Regulated Value Dynamics](https://mlsanity.com/paper/2608.16182/) (2026)
- [Beyond Static Graph World Models: Learning Stochastic Latent Dynamics over Evolving Topologies](https://mlsanity.com/paper/2609.28670/) (2026)

## Abstract (arXiv)

Model-based reinforcement learning agents utilizing transformers have shown improved sample efficiency due to their ability to model extended context, resulting in more accurate world models. However, for complex reasoning and planning tasks, these methods primarily rely on continuous representations. This complicates modeling of discrete properties of the real world such as disjoint object classes between which interpolation is not plausible. In this work, we introduce discrete abstract representations for transformer-based learning (DART), a sample-efficient method utilizing discrete representations for modeling both the world and learning behavior. We incorporate a transformer-decoder for auto-regressive world modeling and a transformer-encoder for learning behavior by attending to task-relevant cues in the discrete representation of the world model. For handling partial observability, we aggregate information from past time steps as memory tokens. DART outperforms previous state-of-the-art methods that do not use look-ahead search on the Atari 100k sample efficiency benchmark with a median human-normalized score of 0.790 and beats humans in 9 out of 26 games. We release our code at https://pranaval.github.io/DART/.