mlsanityData to 14 Oct 2026

Learning to Play Atari in a World of Tokens

arXiv 2406.01361 ↗First posted 3 Jun 2024cs.LGcited by 6

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

also in 1 more
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/.

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.

Builds on

58 of 73 references.

Closest ancestors · 12

Same direction, newest first

  1. 2023STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning14 Oct
  2. 2023Facing Off World Model Backbones: RNNs, Transformers, and S45 Jul
  3. 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
  4. 2023Mastering Diverse Domains through World Models10 Jan
  5. 2022Transformers are Sample-Efficient World Models1 Sep
  6. 2022TransDreamer: Reinforcement Learning with Transformer World Models19 Feb
  7. 2021Mastering Atari Games with Limited Data30 Oct
  8. 2020Mastering Atari with Discrete World Models5 Oct
  9. 2020Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction29 Jul
  10. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec
  11. 2018World Models27 Mar
  12. 2018Model-Ensemble Trust-Region Policy Optimization28 Feb

Other references · 46

  1. 2017Attention Is All You Needcited by 60,361
  2. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  3. 2017Neural Discrete Representation Learningcited by 4,712
  4. 2020Taming Transformers for High-Resolution Image Synthesiscited by 3,253
  5. 2019Transformer-XL: Attentive Language Models Beyond a Fixed-Length Contextcited by 1,850
  6. 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,204
  7. 2021Transformers in Vision: A Surveycited by 1,013
  8. 2019An Introduction to Variational Autoencoderscited by 1,000
  9. 2020On Layer Normalization in the Transformer Architecturecited by 922
  10. 2023Vision Transformers Need Registerscited by 871
  11. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 807
  12. 2020CURL: Contrastive Unsupervised Representations for Reinforcement Learningcited by 760
34 more
  1. 2019RLBench: The Robot Learning Benchmark & Learning Environmentcited by 689
  2. 2020Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixelscited by 585
  3. 2019Model-Based Reinforcement Learning for Ataricited by 502
  4. 2019Are Transformers universal approximators of sequence-to-sequence functions?cited by 406
  5. 2020Data-Efficient Reinforcement Learning with Self-Predictive Representationscited by 318
  6. 2024Reinforcement learningcited by 247
  7. 2019MineRL: A Large-Scale Dataset of Minecraft Demonstrationscited by 215
  8. 2022Recurrent Memory Transformercited by 203
  9. 2021Benchmarking the Spectrum of Agent Capabilitiescited by 200
  10. 2022Masked World Models for Visual Controlcited by 170
  11. 2018Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansioncited by 168
  12. 2019Is a Good Representation Sufficient for Sample Efficient Reinforcement Learning?cited by 140
  13. 2018Probabilistic Recurrent State-Space Modelscited by 64
  14. 2023When Do Transformers Shine in RL? Decoupling Memory from Credit Assignmentcited by 63
  15. 2020On the role of planning in model-based deep reinforcement learningcited by 57
  16. 2023An Investigation into Pre-Training Object-Centric Representations for Reinforcement Learningcited by 49
  17. 2021Offline Reinforcement Learning with Reverse Model-based Imaginationcited by 47
  18. 2019Learning to Combat Compounding-Error in Model-Based Reinforcement Learningcited by 39
  19. 2021Discrete Representations Strengthen Vision Transformer Robustnesscited by 35
  20. 2022Sample Efficient Deep Reinforcement Learning via Uncertainty Estimationcited by 27
  21. 2023E(2)-Equivariant Vision Transformercited by 27
  22. 2018Efficient Model-Based Deep Reinforcement Learning with Variational State Tabulationcited by 20
  23. 2022Planning for Sample Efficient Imitation Learningcited by 20
  24. 2022Temporal Disentanglement of Representations for Improved Generalisation in Reinforcement Learningcited by 19
  25. 2020Weakly-Supervised Reinforcement Learning for Controllable Behaviorcited by 17
  26. 2022When to Update Your Model: Constrained Model-based Reinforcement Learningcited by 14
  27. 2021High-Accuracy Model-Based Reinforcement Learning, a Surveycited by 13
  28. 2022Temporal Latent Bottleneck: Synthesis of Fast and Slow Processing Mechanisms in Sequence Learningcited by 13
  29. 2022Discrete and continuous representations and processing in deep learning: Looking forwardcited by 12
  30. 2021Learning Temporally-Consistent Representations for Data-Efficient Reinforcement Learningcited by 10
  31. 2024Is Next Token Prediction Sufficient for GPT? Exploration on Code Logic Comprehensioncited by 6
  32. 2021Maximum Entropy Model-based Reinforcement Learningcited by 5
  33. 2018The Dreaming Variational Autoencoder for Reinforcement Learning Environmentscited by 4
  34. 2022Deep Reinforcement Learning with Vector Quantized Encodingcited by 2

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 3; 2026: 3.

6 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2022Multi-Game Decision Transformers
  2. 2021Decision Transformer: Reinforcement Learning via Sequence Modeling
  3. 2026Mixture-of-World Models: Scaling Multi-Task Reinforcement Learning with Modular Latent Dynamics
  4. 2025EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling
  5. 2020Smaller World Models for Reinforcement Learning
  6. 2013Playing Atari with Deep Reinforcement Learning
  7. 2024Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining
  8. 2026Improving Weak World Models Behind Strong Agents in Atari Pong