Research map / Reinforcement learning / Visual and representation RL
Efficient World Models with Context-Aware Tokenization
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.
Proposes Δ-IRIS, a world-model agent whose discrete autoencoder encodes stochastic frame deltas conditioned on past frames and actions, and whose transformer interleaves continuous I-tokens to summarize the current state, achieving SOTA on Crafter with ~10x faster training than IRIS.
Problem
Transformer-based world models for model-based RL need long token sequences to simulate visually complex environments, making imagination slow and scaling to complex environments and long horizons computationally prohibitive.
Approach
Condition the discrete autoencoder on past frames and actions so frames are encoded as a few stochastic Δ-tokens, and interleave continuous I-tokens (CNN frame embeddings, not discretized) in the autoregressive transformer's input so it can summarize the current world state without attending to all past Δ-tokens.
Key results
- On Crafter at 10M frames, Δ-IRIS returns 16.1 vs DreamerV3 XL 15.1 and solves on average 17 of 22 tasks.
- Δ-IRIS trains about 10x faster than IRIS (64 tokens): 20 FPS vs 2 FPS.
- On Atari 100k, Δ-IRIS reaches mean HNS 1.39 vs IRIS 1.05 and DreamerV3 1.10, with 11 superhuman games.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add continuous I-tokens to the transformer input (ablation) | Δ-IRIS · Crafter @1M, @5M, and @10M | return: 6.6 (0.2) -> 7.7 (0.5); 10.4 (0.5) -> 15.4 (0.4); 12.6 (0.8) -> 16.1 (0.1) | Helped |
| add continuous I-tokens to transformer conditioning (ablation) | Δ-IRIS · Held-out world-model test set | Next token loss (CE): 1.73 -> 1.57 | Helped |
| add continuous I-tokens to transformer conditioning (ablation) | Δ-IRIS · Held-out world-model test set | Reward loss (CE): 0.135 -> 0.108 | Helped |
| use context-conditioned Δ-token autoencoding rather than independent frame tokenization | Δ-IRIS versus IRIS · Held-out world-model test set | L2 loss: 0.001715 -> 0.000185 | Helped |
| remove the auxiliary max-pixel reconstruction loss (ablation) | Δ-IRIS autoencoder · Held-out world-model test set | L2 loss: 0.000185 -> 0.000178 | Helped |
| remove the auxiliary max-pixel reconstruction loss (ablation) | Δ-IRIS autoencoder · Held-out world-model test set | Max-pixel loss: 0.018 -> 0.031 | Hurt |
| replace the CNN autoencoder with a transformer autoencoder (ablation) | Δ-IRIS autoencoder · Continual-learning setup | training speed | Hurt |
| reduce autoencoder conditioning from four time steps to one (ablation) | Δ-IRIS autoencoder · Crafter | agent performance | No clear effect |
Surprising findings
- Δ-IRIS does not use its world model representations for policy learning, and DreamerV3 is better at the smallest frame budgets (1M frames).
- Predicting Δ-tokens autoregressively is much harder than predicting absolute image tokens, which is why I-tokens were needed.
Limitations
- Uses the same number of Δ-tokens at every time step, not adapting to context uncertainty.
- Does not leverage world model representations for policy learning, leaving the scarce-data regime underexplored.
- Hard exploration: fails to discover crafting tools requiring a table and furnace, missing 4 of 22 Crafter tasks.
Experimental setup
Datasets: Crafter Atari 100k
Benchmarks: Crafter Atari 100k
Metrics: return Crafter score (geometric mean of success rates) human-normalized score interquartile mean L2 loss FPS number of parameters
Baselines: DreamerV3 IRIS SimPLe STORM Δ-IRIS w/o I-tokens
Builds on: IRIS DreamerV3 MPEG I-frames VQ-VAE
Closest papers by meaning
- RLVR-World: Training World Models with Reinforcement Learning (NeurIPS 2025 · similarity 0.87)
- Reinforcement World Model Learning for LLM-based Agents (NeurIPS 2026 · similarity 0.87)
- Improving Token-Based World Models with Parallel Observation Prediction (ICML 2024 · similarity 0.86)
- PoE-World: Compositional World Modeling with Products of Programmatic Experts (NeurIPS 2025 · similarity 0.86)
- Improving Transformer World Models for Data-Efficient RL (ICML 2025 · similarity 0.86)
- Locality Sensitive Sparse Encoding for Learning World Models Online (ICLR 2024 · similarity 0.85)
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.