PAPER REPORTENAll readings ↗

Diffusion for World Modeling: Visual Details Matter in Atari

English reading report: Method, equations, original figures, experiments and reproducibility.

Authors: Eloi Alonso; Adam Jelley; Vincent Micheli; Anssi Kanervisto; Amos Storkey; Tim Pearce; François Fleuret

Affiliations: University of Geneva; University of Edinburgh; Microsoft Research

Source: NeurIPS 2024 · ref-3b350556ce83b51f84c8 ↗ · Project page ↗ · Catalog record

Reading: 457 / 558 · 6 original figures & tables · ~19 min ·

1. Paper overview

In one sentence: EDM diffusion preserves useful Atari visual details with a short sampling schedule, enabling effective imagined policy training while retaining separate reward and action networks. e02e03e04e08e09e13e18e19e20

At a glanceWhat to know
Research problem
Source description

Discrete visual bottlenecks may erase small objects or changes that determine good actions. The problem is to preserve those details while generating enough stable imagined experience for sample-efficient reinforcement learning; expensive image synthesis alone would not solve it. e02e03

Core mechanism
Source description

DIAMOND adapts EDM to autoregressive, action-conditioned world modeling and combines it with reinforcement learning entirely in imagination. e03e04e05

A key reported resultAtari 100k aggregate policy performance: Mean 1.459; IQM 0.641; 11 games above the human reference.

Mean and interquartile mean human-normalized score (HNS); higher is better.. 26 games; 100,000 real actions per game; five DIAMOND training seeds; agent optimization entirely in imagination.

STORM: mean 1.266, IQM 0.636. IRIS: mean 1.046, IQM 0.501. The mean leads the listed imagination-trained baselines; IQM is essentially tied with STORM. Figure 2 supplies bootstrap intervals without textual endpoints. This is not an overall Atari record: Appendix J gives higher means for BBF (2.247) and EfficientZero (1.943). e08e09e16

Reading caution
Source description

The authors identify discrete-control evaluation, short frame-stack memory and separate reward/termination modeling as limitations. CS:GO may forget state when visibility is lost and permit impossible successive jumps. e14

Core contributions

  • Source description

    DIAMOND adapts EDM to autoregressive, action-conditioned world modeling and combines it with reinforcement learning entirely in imagination. e03e04e05

  • Reader analysis

    Static-data comparisons and denoising diagnostics support the authors’ visual-detail explanation for improved Atari performance, but do not isolate visual fidelity as its sole cause. e09e10e11e12

Figure 1. One image is denoised before the policy advances the imagined trajectory. Original paper, p. 2 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read upward within one column and then move right. The vertical arrows follow denoising time $\tau$ backward from $\mathcal{T}$ to zero; they do not represent additional environment actions. The bottom stacks supply clean past observations and actions. Repeated calls to $D_\theta$ turn an initial noisy sample into the clean observation at the top. The separate policy $\pi_\phi$ then chooses an action, which joins that image in the next column’s conditioning. Horizontal time $t$ advances only when the imagined environment advances. This direction agrees with the caption and Equation (5); the diagram separates the inner sampling loop from the outer interaction loop. e03e05e06

What it supports. The figure explains how diffusion becomes a simulator: each generated image affects later conditioning and policy decisions. It also locates the cost of imagination inside each environment step. Three denoising evaluations per Atari frame are an implementation choice documented in Table 3, rather than three actions executed by the agent.

Where the evidence stops. This schematic omits reward, termination and value prediction. Those functions require separate networks in Section 3.2 and Appendix D. It therefore cannot support a claim that one diffusion network jointly emits future images and actions.

2. Motivation

2.1 The problem and the proposed response

Source description

Discrete visual bottlenecks may erase small objects or changes that determine good actions. The problem is to preserve those details while generating enough stable imagined experience for sample-efficient reinforcement learning; expensive image synthesis alone would not solve it. e02e03

2.2 What this reading follows

A world model can generate a recognizable game screen and still confuse the few pixels that tell an agent what to pursue or avoid. DIAMOND studies that gap by predicting images directly with diffusion and training a recurrent policy inside the resulting simulator. The central design problem is cost: an imagined rollout needs many frames, so hundreds of denoising calls per frame would be impractical. Read the figures as a chain from economical generation to policy utility, then examine where that chain weakens. The Atari results, denoising ablation and later offline 3D demonstrations provide different kinds of evidence and should be assessed separately. e02e03e04e08e09e13e18e19e20

3. Research context

We place the paper in the collection through its world–action interface. The catalog labels and the reading’s assessment are shown separately.

Catalog dimensionRecorded classification
Major categoryFoundational work
ArchitectureNot applicable
Prediction paradigmNot applicable
QuadrantNot applicable

3.1 Evidence-based assessment

Supports the recorded classification

Reader analysis

The recorded foundational/classical model-based RL category fits. Separate dynamics, reward/termination and actor/value networks are explicit. Actions condition future-image prediction and come from an RL actor; there is neither joint future/action generation nor inverse-dynamics action extraction. The catalog’s Not applicable quadrant is supported without reclassification. e03e05e06

This is the collection’s architectural analysis, not a new related-work survey. Benchmark comparisons and their protocols appear in Section 6.

4. Problem formulation

4.1 Inputs and outputs

InputsOutputs
  • Atari: four previous 64×64 RGB observations and corresponding discrete actions, plus the noisy next frame and diffusion time (e05, e06).
  • Diffusion model: next image; separate reward/termination model: reward and episode-end predictions; separate actor: an action (e03, e05, e06, e07).

4.2 Equations and their role

Dθ(xt+1τ,ytτ)=cskipτxt+1τ+coutτFθ(cinτxt+1τ,ytτ)D_\theta(x_{t+1}^{\tau},y_t^{\tau})=c_{\mathrm{skip}}^{\tau}x_{t+1}^{\tau}+c_{\mathrm{out}}^{\tau}F_\theta(c_{\mathrm{in}}^{\tau}x_{t+1}^{\tau},y_t^{\tau})
Equation (6): D is the denoiser and F the learned network with parameters θ. Environment time is t; diffusion time is τ. The x term is the noisy future image. Conditioning y includes transformed noise level, clean observation history and action history. The c coefficients scale input, output and skip connection. e04
Λt={rt+γ(1dt)[(1λ)Vϕ(xt+1)+λΛt+1],t<H,Vϕ(xH),t=H.\Lambda_t=\begin{cases}r_t+\gamma(1-d_t)\left[(1-\lambda)V_\phi(x_{t+1})+\lambda\Lambda_{t+1}\right],&t<H,\\V_\phi(x_H),&t=H.\end{cases}
Equation (14): Λ is the lambda-return over horizon H; r and d are predicted reward and boolean termination, γ is discount, λ mixes bootstrapping with longer returns, and V is the value network with parameters φ. Termination removes future value after an imagined episode ends. e07

5. Method in detail

5.1 Why the denoising target matters when only a few calls are affordable

Author claim

Start from a future frame so noisy that almost none of its original content remains. The history and action must then supply most of the information needed to reconstruct a useful successor. In EDM, the skip coefficient shrinks as noise grows, shifting responsibility toward the network’s clean-signal prediction. At low noise, the skip path preserves much of the image and the network predicts a residual. The authors argue that this adaptive target avoids the unhelpful high-noise behavior of their DDPM noise-prediction baseline. Figure 3 is the practical test: under the same architecture and expert data, EDM tolerates much shorter sampling schedules. This is an empirical design argument, not a theorem that all diffusion parameterizations will exhibit the same stability gap. e04e10

5.2 How imagined pixels turn into an executed policy

Source description

The denoiser alone cannot teach an agent which actions are good. DIAMOND first learns reward and termination prediction from real replay, separately from image generation. It then starts short imagined rollouts from replay histories, using conditioning frames to initialize recurrent states. For each step, the actor chooses an action, the world-model components produce the next observation and scalar feedback, and the value network estimates future return. Lambda-returns combine those predictions over the 15-step horizon; stopped-gradient targets train value regression and a REINFORCE policy update. Real interaction is the next source of corrective data for the models. Algorithm 1 alternates these updates rather than optimizing the agent directly on real transitions. Consequently, visual accuracy and reward accuracy are distinct possible bottlenecks for real policy performance. e05e06e07

5.3 Keep the two 3D experiments separate from Atari control

Reader analysis

Appendix M asks whether a diffusion model can generate plausible short videos from recorded actions. Its early CS:GO Clean and motorway-driving experiments use six conditioning frames, 64×64 images and video-quality metrics; frame-stack DIAMOND uses 20 denoising steps. Section 6 is a later addition using the much larger CS:GO Online dataset. That model generates dynamics at 56×30 and adds a diffusion upsampler for 280×150 output. It is an interactive game-engine demonstration with no RL agent, and its quantitative capability evaluation is deferred. Neither experiment establishes the Atari policy result in a continuous-control setting. The offline action examples also expose causal confusion and state instability. Reader interpretation: favorable video distances and playable appearance are useful evidence, but counterfactual control requires its own evaluation. e01e13e19e20e21

5.4 Training and inference

During training

Source description

Each epoch collects real experience, fits dynamics and reward/termination models on replay, then updates actor/value networks on imagined trajectories. Reward and termination use cross-entropy. REINFORCE uses a value baseline and entropy regularization; value regression uses stopped-gradient lambda-returns. e06e07

Source description

Atari uses 1000 epochs, 100 collected steps and 400 training steps per epoch, batch 32, AdamW at 1e−4, and imagination horizon 15. Preprocessing uses frame skip 4, maximum 30 no-ops, life-loss termination and rewards clipped to −1, 0 or 1. e06

During inference

Source description

For each imagined step, the actor samples an action, a separate model predicts reward/termination, and three deterministic Euler denoising steps generate the next image from random initial noise. The frame and action enter subsequent history. Real interaction executes the actor’s action and obtains actual feedback; the described policy does not run lookahead search. e03e05e06e16

5.5 Implementation flow

  1. Condition the denoiser

    A 2D U-Net concatenates previous frames with the noisy future image along channels. Actions and diffusion time enter adaptive group normalization. The environment state is unobserved; this short history supplies the dynamics model’s memory. e03e05

  2. Adapt the target to noise

    EDM blends the noisy input with the U-Net output. At high noise the network target emphasizes clean signal; at low noise it becomes a scaled residual. Log-normal noise sampling concentrates training away from extreme noise levels. e04

  3. Complete the imagined environment

    A distinct CNN–LSTM predicts reward and termination. Another shares its trunk between policy and value heads. Conditioning sequences initialize their recurrent states before imagined rollouts. These modules are not a joint image/action predictor. e05e06

6. Experiments & results

DIAMOND learns an action-conditioned diffusion simulator in pixels, then trains a separate recurrent policy on imagined experience. EDM preconditioning makes short denoising schedules practical. Atari returns support its usefulness, while offline 3D demonstrations expose memory and causal limitations (e03, e04, e05, e09, e14, e21).

6.1 Read the original evidence

Table 1. The mean-score improvement coexists with a near tie in IQM. Original paper, p. 6 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Begin with the bottom rows, where random and human references anchor the normalized aggregates at zero and one. DIAMOND is the rightmost column; STORM is immediately to its left. Then inspect individual game rows to see what the aggregate conceals. Breakout improves over IRIS, whereas BankHeist remains close to its random reference and far below DreamerV3. The game rows report raw returns, so their magnitudes cannot be compared directly across games. Section 4.1 specifies five DIAMOND seeds and a 100,000-action real-environment budget per game; the many imagined training transitions do not count as extra real experience. e08e09e16

What it supports. DIAMOND reports mean HNS 1.459 against STORM’s 1.266, while their IQMs are 0.641 and 0.636. It exceeds the human reference on 11 games. The supported conclusion is a strong mean result among these imagination-trained agents, accompanied by uneven game-level performance and little separation on IQM.

Where the evidence stops. These are point estimates; Figure 2 on page 5 supplies aggregate bootstrap intervals. The table alone does not establish statistical superiority on every metric. Appendix J reports stronger overall means for BBF and EfficientZero under different algorithmic approaches.

Figure 5. Small temporal inconsistencies can alter the information available for learning a policy. Original paper, p. 8 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read each game’s row as two separate four-frame trajectories: IRIS on the left and DIAMOND on the right. The top row is Asterix, where white boxes highlight changes between enemy and reward appearance. The middle row highlights inconsistent bricks and score in Breakout; the bottom row marks Road Runner reward dots. These are temporal-consistency examples, not paired reconstructions of an identical reference trajectory. Both models use the same static expert data and 64×64 output resolution. Inspect object identity and persistence rather than general screenshot attractiveness, because those properties are the authors’ proposed connection to action learning. e12e09

What it supports. The selected DIAMOND examples preserve the highlighted details more consistently. This gives a concrete reason why a small visual error can affect policy learning: pursuing a reward and avoiding an enemy require different behavior. The examples make the performance hypothesis plausible, but the raw Atari returns are separate evidence.

Where the evidence stops. These are selected generations from different model architectures. Matching data and resolution does not isolate tokenization or visual fidelity as the cause of the policy gap. Native image resolution also limits how much detail can be inferred from enlarged crops.

6.2 Results and evaluation conditions

Task & protocolReported resultComparison & interpretation
Atari 100k aggregate policy performance

26 games; 100,000 real actions per game; five DIAMOND training seeds; agent optimization entirely in imagination.

Mean 1.459; IQM 0.641; 11 games above the human reference.

Mean and interquartile mean human-normalized score (HNS); higher is better.

STORM: mean 1.266, IQM 0.636. IRIS: mean 1.046, IQM 0.501.

The mean leads the listed imagination-trained baselines; IQM is essentially tied with STORM. Figure 2 supplies bootstrap intervals without textual endpoints. This is not an overall Atari record: Appendix J gives higher means for BBF (2.247) and EfficientZero (1.943). e08e09e16

Atari 100k Breakout

Same online Atari protocol and five-seed DIAMOND evaluation as Table 1.

132.5

Game return; higher is better.

IRIS 83.7; human reference 30.5.

A concrete gain on a visually analyzed game, but the visual comparison uses a separate static expert dataset and is not itself this policy evaluation. e08e09e12

Denoising-step ablation on ten selected Atari games

Table 7’s selected ten-game subset; default three-step scores versus one-step scores from only one seed.

Three steps: mean HNS 3.052; Boxing 86.9.

Subset mean HNS; Boxing game return.

One step: mean HNS 1.962; Boxing 41.9.

Suggestive support for iterative denoising, weakened by unequal seeds and game selection. Asterix, Krull and Pong favor one step. These subset means are not comparable to the full benchmark’s 1.459. e18

Early offline 3D video modeling

CS:GO Clean: 150k/40k train/test frames; driving: 3.6/0.8 hours. Evaluation uses 1024 real/generated 16-frame clips, six conditioning frames and recorded actions. Frame-stack DIAMOND uses 20 denoising steps.

CS:GO 34.8; driving 80.3.

Fréchet Video Distance (FVD); lower is better.

IRIS with 64 tokens: 85.7 and 276.9; cross-attention DIAMOND: 81.4 and 299.9.

Visual-modeling evidence on offline splits, without RL. Architecture variants differ in size, so conditioning is not isolated at matched parameter count. e19e20

Later CS:GO interactive world-model generation

Section 6’s Online dataset: 5M training and 0.5M held-out frames; dynamics plus upsampler; no RL agent.

10 Hz on an RTX 3090.

Reported interactive generation rate.

No quantitative capability baseline or gameplay score for this larger model.

A qualitative neural game-engine demonstration, distinct from Appendix M’s smaller Clean-data experiment. Quantitative capability assessment is deferred. e13

6.3 Ablations and diagnostic examples

Read component removals and qualitative examples within their stated evaluation conditions.

Figure 3. Few-step EDM avoids the severe color drift visible in the tested DDPM variant. Original paper, p. 7 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Compare corresponding rows across the two panels before comparing rows vertically. Each row fixes the denoising count, labelled n, and each column advances the autoregressive trajectory from the common initial observation toward timestep 1000. Moving downward reduces the computation available for each generated frame. The DDPM panel increasingly loses the background and brick structure, while EDM keeps a recognizable game layout even in its one-step row. Section 5.1 controls the architecture and uses a shared static expert Breakout dataset of 100k frames. This is a model-stability comparison on that dataset, separate from the online Atari agent-training protocol. e04e10e17

What it supports. The displayed EDM trajectories remain visually stable under sampling budgets that cause pronounced DDPM color corruption. The authors connect this to EDM’s noise-dependent target: at high noise the network learns to predict signal instead of mainly reproducing the noisy input. The example supports their practical framework choice in this regime.

Where the evidence stops. Recognizable frames need not track the exact true trajectory. Figure 8 on page 24 checks pixel drift over 400 held-out reference trajectories and still finds EDM divergence. Neither figure establishes equally accurate dynamics for every game or diffusion configuration.

Figure 4. Three denoising steps retain an uncertain opponent more clearly than one. Original paper, p. 7 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Follow the top row from timestep zero to five, then repeat along the lower row. The labels n=1 and n=3 refer to denoising calls per next-image prediction. The white player is controlled by the policy, so its action is available to the world model. The black opponent’s movement is uncertain under the supplied history. In the one-step sequence its shape spreads and fades; in the three-step sequence a distinct opponent remains. Section 5.2 explains the difference through a squared-error predictor averaging possible outcomes, while iterative sampling can move a generation toward a particular mode of the conditional distribution. e11e18

What it supports. This diagnostic explains why a stable one-step Breakout rollout is insufficient justification for using one step everywhere. Uncertainty matters: a conditional mean can be a poor training observation even when the background remains intact. The example motivates the default three-step Atari sampler and the separate policy-return ablation.

Where the evidence stops. Sharpness does not measure whether sampled opponent positions have correct probabilities, and this selected sequence provides no calibration test. The policy-return evidence in Table 7 is suggestive but has unequal seed counts; it is not a definitive causal test of multimodality.

Table 7. The one-step shortcut reduces the reported subset mean, with important exceptions. Original paper, p. 24 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Compare the final two columns within each game, then examine the normalized mean at the bottom. The left DIAMOND column repeats the default three-step results; the right column tests one-step sampling. Boxing connects directly to the uncertainty example, while Breakout shows that visually stable generation does not automatically preserve agent return. Also read the exceptions: Asterix, Krull and Pong score higher in the one-step column. The paragraph above this crop is essential context: it identifies a selected ten-game subset and says the one-step arm has only one seed. Those restrictions apply to every apparent win and loss here. e18e08e09e11

What it supports. The subset mean HNS is 3.052 for three steps and 1.962 for one; Boxing returns are 86.9 and 41.9. These observations support investigating a performance cost for the faster sampler. They do not show a uniform benefit, because several listed games improve with one step.

Where the evidence stops. The default scores come from the five-seed benchmark, whereas one-step scores use one seed. Selection and seed imbalance prevent a clean effect estimate. The subset’s mean must not be compared directly with the 26-game benchmark mean.

7. Analysis & limitations

7.1 What the evidence leaves open

Source description

The authors identify discrete-control evaluation, short frame-stack memory and separate reward/termination modeling as limitations. CS:GO may forget state when visibility is lost and permit impossible successive jumps. e14

Reader analysis

Better mean performance is not uniform improvement: BankHeist scores 19.7 versus DreamerV3’s 649.0. Visual examples cannot establish that perceptual fidelity alone explains every policy gain. e09e12

Reader analysis

Offline driving exhibits causal confusion: slowing the controlled vehicle is associated with traffic ahead slowing too. Unfamiliar CS:GO actions destabilize generations. Plausible video therefore does not establish reliable counterfactual control. e21

7.2 Questions for discussion

  1. Does three-step denoising improve real policy return with matched seeds and imagined training volume?
  2. Would selectively preserving task-critical pixels close the IRIS–DIAMOND gap without changing the policy learner?
  3. How much offline action failure comes from short memory versus absent counterfactual coverage?

8. Reproducibility audit

8.1 Requirements and known gaps

Source description

Use Appendix D–G architecture, preprocessing and alternating updates, including recurrent burn-in and the EDM noise distribution. Budget approximately 2.9 RTX 4090 days and 12 GB VRAM per Atari seed/game. The runtime profile is representative, not guaranteed. e04e05e06e15

Reader analysis

Resolve precise diffusion loss weighting: Equation (7) uses a preconditioned target, while Algorithm 1 summarizes unweighted image reconstruction. Establish software/environment versions and evaluation episode count, unspecified in the PDF. Matched-seed denoising and controlled visual-detail interventions are useful proposed checks. e22e11e12e18

8.2 Proposed reproduction checks

The following checks are proposals motivated by the paper. They have not been run as part of this reading.

Check 1: Repeat the denoising ablation with matched seeds and training volume

Proposed, not run: preselect Boxing, Breakout and Asterix and train one-step and three-step variants with five matched seeds, the same 100k real-action budget, architecture, 15-step imagination horizon and number of imagined training transitions. Evaluate actual game returns with identical episode counts and report paired uncertainty plus wall-clock cost. Separately compare held-out successor images and long rollouts to distinguish blur from accumulated drift. A reproducible three-step return advantage in Boxing would strengthen the multimodality explanation; a gap that disappears across seeds would weaken the single-seed inference. Add an equal-wall-time comparison as a separate efficiency test, keeping its protocol distinct. e06e08e10e11e17e18

Check 2: Test whether task-critical pixels matter beyond generic image degradation

Proposed, not run: freeze a trained DIAMOND simulator and reward/termination model, then train otherwise identical actors with three observation conditions: intact images, masks over annotated task-critical Asterix object regions, and equal-area masks over background regions. Apply the same actor-input intervention during imagined training and real-game evaluation, while leaving the simulator’s conditioning and reward computation unchanged. Match seeds and training volume. A larger return loss under object masking than background masking would support the importance of those visual details. Similar losses would favor a broader degradation explanation. This tests the information hypothesis; it would not by itself establish that IRIS tokenization caused the original performance gap. e05e06e09e12

8.3 Reading coverage

Visual audit: Visually inspected the title/author page, all Figures 1–12 and Tables 1–8, and the supporting equations, architectures, training configuration, RL objectives, Algorithm 1, compute profile and offline diagnostics on the declared pages. All six final original crops were individually viewed at their generated resolution; axes, panel labels and table rows are retained. Figure 1’s upward denoising arrows were checked against its caption and Equation (5); the sampling/training account was checked against Appendices A and D–G. References and checklist pages were read in the complete text pass, not visually inspected. External videos, code and separate supplements were outside this pass.

PDF pages inspected for this edition: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28. Appendix coverage: reviewed.

Original figures and tables remain the work of the source’s authors. Extractions preserve their scientific content; any HTML wrapper layout is disclosed with each figure. The surrounding reading notes are our own.

Text reading scope & known omissions
  • Abstract; Sections 1–9, including all subsections
  • Acknowledgments and References
  • Appendices A–G: sampling, DDPM link, EDM preconditioning, architectures, hyperparameters, RL objectives and algorithm
  • Appendices H–L: comparisons, compute profile, drift and denoising ablation
  • Appendix M.1–M.4: offline environments, architectures, metrics, compute and diagnostics
  • NeurIPS checklist items 1–15

Outside the original text pass

  • Version scope: the supplied NeurIPS 2024 PDF identifies itself as the final version of Alonso et al. (2023). The title matches the catalog apart from a dagger footnote marker; all seven author names match. Section 6 was added after acceptance. No numbered revision or exact revision date is printed; earlier editions were not supplied or compared (e01).
  • Text extraction does not reconstruct figure images; inspect the retained PDF for figures and equation/table layout.
  • Separate supplemental material availability has not been fully verified.
  • Visual inspection addressed the text extraction limitation: all twelve figures, all eight tables and the supporting method/evaluation pages were inspected.
  • Linked code, datasets, interactive models and external videos were not inspected. No experiment was reproduced. All ten supplied text chunks were read. References and the checklist were read as text; visual inspection covered PDF pages 1–10 and 16–28.

The visual audit above records the subsequent illustrated pass.

8.4 Traceable evidence

e01PDF p. 1, title block and dagger footnote; p. 8, Section 6 footnote 4Inspect

All seven title-page authors match the catalog. The dagger refers to this as the final version of Alonso et al. (2023). Section 6 explicitly says it was added after NeurIPS acceptance.

Go to primary source ↓
e02PDF pp. 1–2, Abstract and Section 1Inspect

Discrete compression can discard small visual details relevant to decisions. DIAMOND trains an RL agent inside an image-space diffusion world model.

Go to primary source ↓
e03PDF p. 2, Figure 1 and caption; p. 4, Section 2.3, Equation (5)Inspect

Past observations and actions condition next-image denoising. Reverse diffusion runs from noise to a clean frame within each environment step; generated images and policy actions enter subsequent conditioning.

Go to primary source ↓
e04PDF p. 4, Section 3.1, Equations (6)–(7); p. 16, Appendix C, Equations (9)–(13)Inspect

EDM combines a scaled noisy-input skip connection with a preconditioned network prediction. Its target depends on noise level; sigma_data is 0.5 and log sigma has mean −0.4 and standard deviation 1.2.

Go to primary source ↓
e05PDF p. 5, Sections 3.1–3.2; p. 17, Appendix D and Table 2Inspect

The Atari U-Net stacks four prior frames and uses adaptive group normalization for actions and diffusion time. Separate CNN–LSTM reward/termination and actor/value networks initialize recurrent states with conditioning sequences.

Go to primary source ↓
e06PDF p. 18, Appendix E, Table 3; p. 20, Appendix G, Algorithm 1Inspect

The loop alternates real collection, diffusion and reward/termination training, and imagined actor/value training. Settings include 1000 epochs, 100 collected steps and 400 training steps per epoch, batch 32, horizon 15, AdamW at 1e−4, and three Euler steps.

Go to primary source ↓
e07PDF p. 19, Appendix F, Equations (14)–(16)Inspect

Predicted rewards and termination define lambda-returns with a terminal value bootstrap. Value regression and REINFORCE use stopped-gradient targets/advantages, with policy entropy regularization.

Go to primary source ↓
e08PDF p. 5, Sections 4.1–4.2 and Figure 2Inspect

Atari 100k comprises 26 games with 100,000 real actions per game. DIAMOND trains from scratch with five seeds. Aggregate uncertainty is shown using stratified bootstrap confidence intervals.

Go to primary source ↓
e09PDF p. 6, Table 1, aggregate rows and Breakout/BankHeist rowsInspect

DIAMOND: mean HNS 1.459, IQM 0.641, 11 superhuman games. STORM: 1.266 and 0.636; IRIS: 1.046 and 0.501. Breakout: DIAMOND 132.5, IRIS 83.7, human 30.5. BankHeist: DIAMOND 19.7, DreamerV3 649.0.

Go to primary source ↓
e10PDF p. 6, Section 5.1; p. 7, Figure 3 and explanatory paragraphsInspect

DDPM and EDM use the same architecture and static 100k-frame expert Breakout dataset. Trajectories compare 1, 3, 5 and 10 denoising steps through timestep 1000; DDPM shows severe color drift at low step counts.

Go to primary source ↓
e11PDF p. 7, Section 5.2 and Figure 4Inspect

Single-step Boxing generation blurs the unpredictable black opponent; three-step generation stays sharper. The controlled white player is predictable in both. The authors explain iterative sampling as resolving multimodal uncertainty.

Go to primary source ↓
e12PDF p. 8, Section 5.3 and Figure 5Inspect

Shared static expert datasets of 100k frames support the IRIS–DIAMOND visual comparison. White boxes mark IRIS inconsistencies in Asterix objects, Breakout bricks/score and Road Runner rewards. Both render 64×64 images; DIAMOND uses 3 NFE versus IRIS 16.

Go to primary source ↓
e13PDF pp. 8–9, Section 6 and Figure 6Inspect

The later CS:GO Online experiment uses 5M training and 0.5M held-out frames, without RL. A 56×30 dynamics model plus 280×150 upsampler totals 381M parameters, including 51M for upsampling. Training takes 12 RTX 4090 days; reported generation is 10 Hz on RTX 3090. Quantitative capability evaluation is deferred.

Go to primary source ↓
e14PDF p. 10, Section 8; p. 9, Section 6, final paragraphInspect

Limitations include discrete-control evaluation, short frame-stack memory, separate reward/termination modeling, drift in less visited CS:GO areas, state forgetting after lost visibility and erroneous successive jumps.

Go to primary source ↓
e15PDF p. 5, Section 4.1; p. 21, Table 4; p. 22, Table 5Inspect

An Atari run takes approximately 2.9 days on one RTX 4090 using about 12 GB VRAM. Table 4 lists 13M total parameters. The representative profile separates 2.5 training days from 0.4 days for collection, evaluation and checkpointing.

Go to primary source ↓
e16PDF p. 23, Appendix J, Table 6, mean and IQM rowsInspect

The broader comparison gives mean HNS 2.247 for BBF, 1.943 for EfficientZero, and 1.459 for DIAMOND. BBF uses different learning techniques and EfficientZero uses lookahead search.

Go to primary source ↓
e17PDF p. 24, Appendix K, Figure 8 and captionInspect

Pixel drift uses 400 held-out 1000-step expert Breakout reference trajectories with matched initial frames/actions. Curves show means with standard deviations; EDM trajectory divergence is distinguished from DDPM pathological color shift.

Go to primary source ↓
e18PDF p. 24, Appendix L, Table 7 and preceding paragraphInspect

On ten selected games, three-step mean HNS is 3.052 versus 1.962 for one step; Boxing returns are 86.9 versus 41.9. One-step results use one seed. Asterix, Krull and Pong favor one step.

Go to primary source ↓
e19PDF p. 25, Appendix M.1Inspect

Early offline experiments use CS:GO Clean data with 150k/40k train/test frames and daylight motorway driving with 3.6/0.8 train/test hours. Images are 64×64; neither experiment includes RL.

Go to primary source ↓
e20PDF pp. 25–26, Appendix M.2–M.4, Figure 9 and Table 8Inspect

The comparison uses six conditioning frames, real actions, 1024 real/generated 16-frame clips, and 120k updates at batch 64. Frame-stack DIAMOND uses 20 denoising steps, reporting FVD 34.8/80.3 for CS:GO/driving versus IRIS-64 at 85.7/276.9 and cross-attention at 81.4/299.9. Runs use up to four A6000 GPUs.

Go to primary source ↓
e21PDF pp. 27–28, Appendix M.4, Figures 11–12 and captionsInspect

Fixed-action tests expose causal confusion in motorway traffic speed and CS:GO degeneration under unfamiliar sustained look directions. The authors attribute these issues to offline data.

Go to primary source ↓
e22PDF p. 4, Equation (7); p. 20, Algorithm 1, update_diffusion_model; pp. 5 and 18, evaluation/configuration descriptionsInspect

Equation (7) uses a preconditioned target divided by c_out; Algorithm 1 summarizes clean-image reconstruction without an explicit noise-level weight. The evaluation/configuration descriptions do not state an evaluation episode count or software versions.

Go to primary source ↓

8.5 Primary sources

Scroll across the image to inspect details. Press Esc to close.