DIAMOND

A world model has to answer $p(x_{t+1} \mid x_{\leq t}, a_{\leq t})$. DIAMOND answers it with a conditional diffusion model over the image, dropping the discrete latent bottleneck its predecessors relied on.

The conditioning is the architecture. The last $L$ clean frames are concatenated to the noised next frame along the channel axis --- plain frame stacking, which keeps the network a standard U-Net 2D --- while the actions and the diffusion time enter through adaptive group normalisation inside the residual blocks. Nothing about the image path knows it is a world model.

EDM, not DDPM, and the difference is not cosmetic. The network $F_\theta$ is wrapped in Karras et al.'s preconditioners:

so the training target adaptively mixes signal and noise with the degradation level. When $\sigma \gg \sigma_{\text{data}}$ the skip vanishes and the network is asked for the clean frame; when $\sigma \to 0$ it is asked for the added noise. DDPM asks for the noise at every level, which makes the high-noise regime an identity map and leaves the score badly estimated exactly where sampling begins. Autoregressed over a thousand imagined steps that error compounds --- the paper shows DDPM drifting out of distribution while the EDM model stays stable even at a single denoising step.

Three steps, and the reason it is not one. A denoiser trained under an $L_2$ loss predicts the expectation over possible reconstructions, so when the next frame is genuinely multi-modal --- an opponent whose move the agent cannot predict --- one step returns a blur between the outcomes. Iterating drives the sample onto a single mode. Hence $n = 3$, everywhere in the paper.

What diffusion does not model. Rewards and episode ends are scalar predictions, not images, so they get their own CNN-LSTM; the actor-critic gets another. The agent is trained entirely inside the imagined environment, reaching a mean human-normalised score of 1.46 on Atari 100k with 13M parameters --- fewer than IRIS's 30M or DreamerV3's 18M.