Research map / Architectures and efficiency / Attention and transformers
Learning to Play Atari in a World of Tokens
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces DART, a model-based RL agent that uses VQ-VAE discrete tokens, a GPT-style transformer decoder for world modeling, and a ViT-style transformer encoder policy with a memory token, reaching SOTA median score among no-look-ahead methods on Atari 100k.
Problem
Model-based RL on Atari 100k is sample-inefficient, and continuous latent representations struggle to model discrete object properties and long-range dependencies; partial observability also requires memory.
Approach
Discrete VQ-VAE tokens model both the world (GPT-style causal transformer decoder predicting next tokens, rewards, terminations) and the policy (ViT-style transformer encoder with CLS and learned MEM tokens acting as recurrent memory), replacing CNN/LSTM policy networks used in IRIS.
Key results
- DART reaches a median human-normalized score of 0.790 on Atari 100k, the best among no-look-ahead-search methods, versus 0.466 for DreamerV3 and 0.289 for IRIS.
- DART is superhuman on 9 of 26 Atari games, matching DreamerV3 and IRIS, and achieves IQM 0.575 and optimality gap 0.458.
- On Crafter at 1M steps, DART obtains return 12.2 ± 1.67 versus 11.7 for DreamerV3 and 9.23 for IRIS.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DART instead of DreamerV3 | Atari 100k agents · Atari 100k | Median human-normalized score: 0.466 -> 0.790 | Helped |
| Use DART instead of IRIS | Atari 100k agents · Atari 100k | Median human-normalized score: 0.289 -> 0.790 | Helped |
| Add learnable positional encoding (ablation) | DART · Boxing, Amidar, Road Runner, Seaquest, KungFuMaster | score: 3.86 -> 83.0; 77.1 -> 125.7; 1030.0 -> 7772.5; 64.2 -> 895.8; 1028.0 -> 23744.3 | Helped |
| Add random exploration during data collection (ablation) | DART · Boxing, Amidar, Road Runner, Seaquest, KungFuMaster | score: 58.67 -> 83.0; 92.75 -> 125.7; 3597.1 -> 7772.5; 753.93 -> 895.8; 15464.7 -> 23744.3 | Helped |
| Add the memory token instead of masking it (ablation) | DART · Boxing, Amidar, Road Runner, Seaquest, KungFuMaster | score: 81.45 -> 83.0; 113.69 -> 125.7; 8021.0 -> 7772.5; 704.8 -> 895.8; 20378.0 -> 23744.3 | Mixed |
| Retain all observation tokens instead of randomly masking 25% (ablation) | DART · Boxing, Amidar, Road Runner, Seaquest, KungFuMaster | score: 77.79 -> 83.0; 102.47 -> 125.7; 7354.0 -> 7772.5; 491.4 -> 895.8; 16436.0 -> 23744.3 | Helped |
| Retain all observation tokens instead of randomly masking 50% (ablation) | DART · Boxing, Amidar, Road Runner, Seaquest, KungFuMaster | score: 51.14 -> 83.0; 56.37 -> 125.7; 2730.0 -> 7772.5; 207.8 -> 895.8; 9760.0 -> 23744.3 | Helped |
| Retain all observation tokens instead of randomly masking 75% (ablation) | DART · Boxing, Amidar, Road Runner, Seaquest, KungFuMaster | score: 15.64 -> 83.0; 52.22 -> 125.7; 988.0 -> 7772.5; 104.0 -> 895.8; 4676.2 -> 23744.3 | Helped |
| Retain all observation tokens rather than attend solely to memory (ablation) | DART · Boxing, Amidar, Road Runner, Seaquest, KungFuMaster | score: -11.91 -> 83.0; 30.43 -> 125.7; 961.0 -> 7772.5; 142.0 -> 895.8; 1571.8 -> 23744.3 | Helped |
Surprising findings
- Mean human-normalized score of DART (1.022) is slightly below DreamerV3 (1.120) and IRIS (1.046), so the median gain does not carry over to the mean.
- Masking the memory token slightly improved RoadRunner performance, suggesting the memory token is redundant in some games.
Limitations
- Designed only for discrete action spaces; continuous-control extension is only sketched in the appendix.
- Does not use look-ahead search; integration with MCTS is described but not evaluated.
- Mean score is below DreamerV3 and IRIS, and performance is weaker on games like Breakout and CrazyClimber.
Experimental setup
Datasets: Atari 100k
Benchmarks: Atari 100k Crafter
Metrics: median human-normalized score mean human-normalized score interquartile mean (IQM) optimality gap superhuman game count
Baselines: IRIS DreamerV3 TWM SPR STORM SimPLe CURL DrQ
Builds on: IRIS DreamerV2 VQ-VAE GPT ViT
Closest papers by meaning
- EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling (NeurIPS 2025 · similarity 0.86)
- In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.86)
- Diffusion for World Modeling: Visual Details Matter in Atari (NeurIPS 2024 · similarity 0.86)
- Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining (ICLR 2025 · similarity 0.86)
- Improving Token-Based World Models with Parallel Observation Prediction (ICML 2024 · similarity 0.86)
- Investigating Pre-Training Objectives for Generalization in Vision-Based Reinforcement Learning (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.