Research map / Reinforcement learning / Visual and representation RL
Improving Transformer World Models for Data-Efficient RL
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.
Improves transformer-world-model MBRL on Craftax-classic via Dyna-with-warmup, patch nearest-neighbor tokenization, and block teacher forcing, plus a stronger PPO+GRU model-free baseline, reaching SOTA reward of 69.66% at 1M steps.
Problem
Sample-efficient model-based RL in a visually complex, partially observable, procedurally generated open-world survival game (Craftax-classic / Crafter) where existing MBRL and MFRL agents underperform.
Approach
Combines real-environment PPO updates with imagined rollouts after a warmup (Dyna with warmup), tokenizes each 7x7 image patch independently with a growing non-parametric nearest-neighbor codebook, and trains the transformer world model with block teacher forcing that predicts all next-step tokens in parallel from history.
Key results
- Best MBRL agent reaches 69.66% reward and 31.77% score on Craftax-classic at 1M steps, exceeding previous SOTA reward of 53.20% (DreamerV3) and reported human-expert reward.
- Best MFRL agent (CNN+GRU, larger model) reaches 55.49% reward and 16.77% score, beating DreamerV3's 53.20% reward at a training time of about 15 minutes on one A100.
- Ablation ladder: Dyna 43.36%, patch factorization 58.92%, nearest-neighbor tokenizer 64.96%, block teacher forcing 67.42% (fast variant, about twice as fast).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add Dyna training on real and imagined trajectories (ablation) | Craftax-classic MBRL · Craftax-classic | reward (%): 31.93 ± 2.22 -> 43.36 ± 1.84 | Helped |
| Factorize the VQ-VAE over observation patches (ablation) | Craftax-classic MBRL with Dyna · Craftax-classic | reward (%): 43.36 ± 1.84 -> 58.92 ± 1.03 | Helped |
| Replace patch VQ-VAE with nearest-neighbor tokenizer (ablation) | Craftax-classic MBRL with patch factorization · Craftax-classic | reward (%): 58.92 ± 1.03 -> 64.96 ± 1.13 | Helped |
| Add block teacher forcing (ablation) | Craftax-classic MBRL with patch nearest-neighbor tokenizer · Craftax-classic | reward (%): 64.96% -> 67.42% | Helped |
| Train the policy on both real and imagined trajectories rather than only imagined trajectories (ablation) | Craftax-classic MBRL · Craftax-classic | reward (%): 55.02% -> 67.42% | Helped |
| Warm up the world model before starting imagination training (ablation) | Craftax-classic MBRL with Dyna · Craftax-classic | reward (%): 33.54% -> 67.42% | Helped |
| Increase TWM training steps to 4k (ablation) | Craftax-classic M5 · Craftax-classic | reward (%): 67.42% -> 69.66% | Helped |
| Use 5×5 or 9×9 patches instead of ground-truth 7×7 patches (ablation) | Craftax-classic NNT · Craftax-classic | reward | Mixed |
| Train the TWM to reconstruct continuous patches instead of quantizing them (ablation) | Craftax-classic MBRL · Craftax-classic | reward | Hurt |
| Remove any individual rung of the MBRL improvement ladder (ablation) | Craftax-classic best MBRL · Craftax-classic | reward | Hurt |
| Increase imagined policy updates linearly from 0 to 300 (ablation) | Craftax-classic MBRL · Craftax-classic | reward (%): 65.71% (± 1.11) | Helped |
| Remove the RNN from the MFRL policy (ablation) | Craftax-classic MFRL · Craftax-classic | performance | Hurt |
Surprising findings
- Naively increasing MFRL model size harms performance, while a larger model combined with a low-dimensional GRU helps.
- Training the policy on imagined rollouts from step 0 (no warmup) drops reward from 67.42% to 33.54%, and removing real-environment policy updates drops reward to 55.02%.
Limitations
- Experiments restricted to grid-world environments (Craftax, MinAtar, OpenSpiel); generalization beyond them is untested.
- Patch tokenization is sensitive to patch size and NNT may create a large codebook with appearance variation.
- Best MBRL result requires about 2 days of training (4k TWM updates).
Experimental setup
Datasets: Craftax-classic Craftax (full) MinAtar OpenSpiel (Bargaining, Leduc Poker, Tic-Tac-Toe)
Benchmarks: Craftax-classic Craftax MinAtar OpenSpiel
Metrics: reward (%) score (%) symbol accuracy of rollouts wall-clock time
Baselines: DreamerV3 IRIS Delta-IRIS Curious Replay Previous best MFRL (Crafter baseline) Artificial Dopamine
Builds on: IRIS DreamerV3 PPO Dyna purejaxrl
Closest papers by meaning
- Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining (ICLR 2025 · similarity 0.87)
- RLVR-World: Training World Models with Reinforcement Learning (NeurIPS 2025 · similarity 0.87)
- Accurate and Efficient World Modeling with Masked Latent Transformers (ICML 2025 · similarity 0.87)
- Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement (NeurIPS 2024 · similarity 0.86)
- Zero-shot World Models via Search in Memory (NeurIPS 2025 · similarity 0.86)
- Efficient World Models with Context-Aware Tokenization (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.