Research map / Multi-agent systems and games / Human-AI and goal-conditioned
Diffusion for World Modeling: Visual Details Matter in Atari
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 DIAMOND, an RL agent trained entirely within a diffusion world model (EDM-based), reaching a mean HNS of 1.46 on Atari 100k and showing the diffusion model can serve as a playable CS:GO game engine.
Problem
World models for RL usually compress observations into discrete latents, which can lose visual details important for decision-making; the paper asks whether diffusion models can serve as stable, sample-efficient world models.
Approach
Use a continuous-time EDM diffusion model (with its preconditioning and signal/noise-mixing training target) as the next-observation predictor of a world model, sampled with very few Euler denoising steps (3), so that imagined trajectories remain stable over long horizons and preserve visual details; the RL agent is trained with REINFORCE and lambda-returns entirely in imagination.
Key results
- DIAMOND achieves a mean HNS of 1.459 and IQM of 0.641 on Atari 100k, a new best for agents trained entirely within a world model, superhuman on 11 games.
- Compared with IRIS (mean HNS 1.046) using 16 NFE per frame, DIAMOND uses 3 NFE per frame with 13M parameters versus IRIS's 30M.
- With a single denoising step, DIAMOND's mean HNS drops on the 10-game subset from 3.052 to 1.962 (single seed), notably on Boxing and RoadRunner.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the DIAMOND diffusion world model instead of STORM | Atari 100k world-model agents · Atari 100k | mean HNS: 1.266 -> 1.459 | Helped |
| Use the DIAMOND diffusion world model instead of STORM | Atari 100k world-model agents · Atari 100k | IQM HNS: 0.636 -> 0.641 | Helped |
| Use the DIAMOND diffusion world model instead of BBF | Atari 100k agents · Atari 100k | mean HNS: 2.247 -> 1.459 | Hurt |
| Use the DIAMOND diffusion world model instead of BBF | Atari 100k agents · Atari 100k | IQM HNS: 1.139 -> 0.641 | Hurt |
| Use 3 denoising steps instead of 1 (ablation) | DIAMOND · Atari 100k, 10 highest-performing games | Mean HNS: 1.962 -> 3.052 | Helped |
| Use 3 denoising steps instead of 1 (ablation) | DIAMOND · Atari 100k games in the 10-game ablation | game return: Amidar 191.8 -> 225.8; Assault 782.5 -> 1526.4; Asterix 6687.0 -> 3698.5; Boxing 41.9 -> 86.9; Breakout 50.8 -> 132.5; C | Mixed |
| Use EDM instead of DDPM (ablation) | Diffusion world model on Breakout · Breakout, autoregressive trajectories up to 1000 timesteps | trajectory stability | Helped |
Surprising findings
- DDPM-based world models suffer severe compounding error with few denoising steps, whereas EDM remains stable even with a single step.
- Single-step EDM sampling yields blurry predictions in multimodal situations (e.g., Boxing), requiring multi-step denoising for crisp frames.
Limitations
- Evaluation focused on discrete-action environments; continuous control not explored.
- Memory is limited to frame stacking of the last 4 observations, limiting long-term memory.
- Reward/termination prediction is kept separate from the diffusion model.
Experimental setup
Datasets: Atari 100k benchmark CS:GO Online dataset (Dust II, 5.5M frames) CS:GO Clean dataset (190k frames) commaai motorway driving dataset
Benchmarks: Atari 100k
Metrics: mean human-normalized score IQM superhuman game count per-game return
Baselines: IRIS DreamerV3 STORM TWM SimPLe EfficientZero BBF
Builds on: EDM DDPM IRIS DreamerV3 Atari 100k
Closest papers by meaning
- EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling (NeurIPS 2025 · similarity 0.90)
- Learning 3D Persistent Embodied World Models (NeurIPS 2025 · similarity 0.87)
- Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective (NeurIPS 2025 · similarity 0.87)
- Diffusion Models Are Real-Time Game Engines (ICLR 2025 · similarity 0.87)
- Zero-shot World Models via Search in Memory (NeurIPS 2025 · similarity 0.87)
- Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL (NeurIPS 2026 · similarity 0.87)
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.