Research map / Multi-agent systems and games / Human-AI and goal-conditioned
Accurate and Efficient World Modeling with Masked Latent Transformers
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 EMERALD, a Transformer world model with a spatial categorical latent state and MaskGIT parallel token prediction for training Dreamer-style actor-critic agents in latent space, reaching state-of-the-art on Crafter.
Problem
Compressed vector latent spaces in Dreamer-style world models lose visual details crucial for agents on complex visual environments like Crafter, while accurate pixel-space world models (Δ-IRIS, DIAMOND) hurt training efficiency by forcing agents to learn from reconstructed images.
Approach
Use a spatial (4x4x32x32) categorical latent state combined with a temporal masked-attention Transformer (TSSM) and a spatial MaskGIT predictor decoded in a few parallel steps, so the agent learns in latent space while benefiting from accurate, efficiently decoded imagined trajectories.
Key results
- EMERALD reaches a Crafter achievement score of 58.1% after 10M steps over 5 seeds, above human experts (50.5%) and Δ-IRIS (42.5%).
- EMERALD mean return on Crafter is 16.8 ± 0.6 versus 16.1 for Δ-IRIS and 14.0 for DreamerV3 (M).
- On Atari 100k, EMERALD achieves 134% normed mean and 51% normed median, below DIAMOND (146%) and Δ-IRIS (139%) on mean.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use spatial latents ($4 \times 4 \times 32$) and MaskGIT instead of vector latents without MaskGIT (ablation) | DreamerV3 RSSM · Crafter | Score (%): 34.9 -> 40.4 | Helped |
| Use spatial latents ($4 \times 4 \times 32$) and MaskGIT instead of vector latents without MaskGIT (ablation) | DreamerV3 RSSM · Crafter | Pixel $L_{2}$ loss: 0.000522 -> 0.000231 | Helped |
| Use spatial latents ($4 \times 4 \times 32$) with MaskGIT instead of vector latents without MaskGIT (ablation) | EMERALD TSSM · Crafter | Score (%): 31.8 -> 58.1 | Helped |
| Use spatial latents ($4 \times 4 \times 32$) with MaskGIT instead of vector latents without MaskGIT (ablation) | EMERALD TSSM · Crafter | Pixel $L_{2}$ loss: 0.000890 -> 0.000241 | Helped |
| Use 3 MaskGIT decoding steps instead of the linear dynamics head (ablation) | EMERALD · Crafter | Score (%): 51.6 -> 58.1 | Helped |
| Use 3 MaskGIT decoding steps instead of 1 (ablation) | EMERALD · Crafter | Score (%): 53.8 -> 58.1 | Helped |
| Use 3 MaskGIT decoding steps instead of 8 (ablation) | EMERALD · Crafter | Score (%): 55.1 -> 58.1 | Helped |
| Increase MaskGIT decoding steps from 3 to 8 (ablation) | EMERALD · Crafter | FPS: 27 -> 23 | Hurt |
| Use EMERALD rather than Δ-IRIS | Crafter agents trained for 10M environment frames · Crafter | Score (%): 42.5 -> 58.1 | Helped |
| Use EMERALD rather than DIAMOND | Atari 100k agents · Atari 100k (26 games) | Normed Mean (%): 146 -> 134 | Hurt |
| Use EMERALD rather than Δ-IRIS | Atari 100k agents · Atari 100k (26 games) | Normed Mean (%): 139 -> 134 | Hurt |
| Train with EMERALD rather than Δ-IRIS or DIAMOND | Atari 100k Breakout · Breakout | Training time to 400k environment frames (hours): 27 hours -> 17 hours; 75 hours -> 17 hours | Helped |
Surprising findings
- Removing MaskGIT and using a linear head with the spatial latent reduces Crafter score from 58.1% to 31.8% for the TSSM variant, showing the MaskGIT predictor matters a lot.
- On Atari 100k, EMERALD does not beat Δ-IRIS or DIAMOND in aggregate performance, despite its Crafter gains.
Limitations
- Atari 100k aggregate performance is comparable to but not above Δ-IRIS and DIAMOND.
- Results on Crafter are averaged over only 5 seeds and training is limited to 10M steps.
- Evaluated only on Crafter and Atari 100k; not tested on Minecraft-like or other complex environments.
Experimental setup
Datasets: Crafter Atari 100k
Benchmarks: Crafter Atari 100k
Metrics: achievement score (%) mean episode return FPS parameters human-normalized mean and median pixel L2 loss token prediction accuracy
Baselines: DreamerV3 (M) DreamerV3 (XL) Δ-IRIS DIAMOND STORM IRIS TWM SimPLe
Builds on: DreamerV3 IRIS TSSM MaskGIT DreamerV2
Closest papers by meaning
- Improving Transformer World Models for Data-Efficient RL (ICML 2025 · similarity 0.87)
- Learning Transformer-based World Models with Contrastive Predictive Coding (ICLR 2025 · similarity 0.86)
- Zero-shot World Models via Search in Memory (NeurIPS 2025 · similarity 0.85)
- Learning Latent Dynamic Robust Representations for World Models (ICML 2024 · similarity 0.85)
- Copilot4D: Learning Unsupervised World Models for Autonomous Driving via Discrete Diffusion (ICLR 2024 · similarity 0.84)
- WALL-E: World Alignment by NeuroSymbolic Learning improves World Model-based LLM Agents (NeurIPS 2025 · similarity 0.84)
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.