Diffusion for World Modeling:
Visual Details Matter in Atari^††thanks: To prevent confusion, this is the final version of (Alonso et al.,, 2023) and is not related to (Ding et al.,, 2024).

We model the environment as a standard Partially Observable Markov Decision Process (pomdp) (Sutton and Barto,, 2018), $(\mathcal{S},\mathcal{A},\mathcal{O},T,R,O,\gamma)$, where $\mathcal{S}$ is a set of states, $\mathcal{A}$ is a set of discrete actions, and $\mathcal{O}$ is a set of image observations. The transition function $T:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\to[0,1]$ describes the environment dynamics $p(\mathbf{s}_{t+1}\mid\mathbf{s}_{t},\mathbf{a}_{t})$, and the reward function $R:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\to\mathbb{R}$ maps transitions to scalar rewards. Agents cannot directly access states $s_{t}$ and only see the environment through image observations $x_{t}\in\mathcal{O}$, emitted according to observation probabilities $p(\mathbf{x}_{t}\mid\mathbf{s}_{t})$, described by the observation function $O:\mathcal{S}\times\mathcal{O}\to[0,1]$. The goal is to obtain a policy $\pi$ that maps observations to actions in order to maximize the expected discounted return $\mathbb{E}_{\pi}[\sum_{t\geq 0}\gamma^{t}r_{t}]$, where $\gamma\in[0,1]$ is a discount factor. World models (Ha and Schmidhuber,, 2018) are generative models of environments, i.e. models of $p(s_{t+1},r_{t}\mid s_{t},a_{t})$. These models can be used as simulated environments to train RL agents (Sutton,, 1991) in a sample-efficient manner (Wu et al.,, 2023). In this paradigm, the training procedure typically consists of cycling through the three following steps: collect data with the RL agent in the real environment; train the world model on all the collected data; train the RL agent in the world model environment (commonly referred to as "in imagination").

Diffusion models (Sohl-Dickstein et al.,, 2015) are a class of generative models inspired by non-equilibrium thermodynamics that generate samples by reversing a noising process.

We consider a diffusion process $\{\mathbf{x}^{\tau}\}_{\tau\in[0,\mathcal{T}]}$ indexed by a continuous time variable $\tau\in[0,\mathcal{T}]$, with corresponding marginals $\{p^{\tau}\}_{\tau\in[0,\mathcal{T}]}$, and boundary conditions $p^{0}=p^{data}$ and $p^{\mathcal{T}}=p^{prior}$, where $p^{prior}$ is a tractable unstructured prior distribution, such as a Gaussian. Note that we use $\tau$ and superscript for the diffusion process time, in order to keep $t$ and subscript for the environment time.

This diffusion process can be described as the solution to a standard stochastic differential equation (SDE) (Song et al.,, 2020),

where $\mathbf{w}$ is the Wiener process (Brownian motion), $\mathbf{f}$ a vector-valued function acting as a drift coefficient, and $g$ a scalar-valued function known as the diffusion coefficient of the process.

To obtain a generative model, which maps from noise to data, we must reverse this process. Remarkably, Anderson, (1982) shows that the reverse process is also a diffusion process, running backwards in time, and described by the following SDE,

where $\bar{\mathbf{w}}$ is the reverse-time Wiener process, and $\nabla_{\mathbf{x}}\log p^{\tau}(\mathbf{x})$ is the (Stein) score function, the gradient of the log-marginals with respect to the support. Therefore, to reverse the forward noising process, we are left to define the functions $f$ and $g$ (in Section 3.1), and to estimate the unknown score functions $\nabla_{\mathbf{x}}\log p^{\tau}(\mathbf{x})$, associated with marginals $\{p^{\tau}\}_{\tau\in[0,\mathcal{T}]}$ along the process. In practice, it is possible to use a single time-dependent score model $\mathbf{S}_{\theta}(\mathbf{x},\tau)$ to estimate these score functions (Song et al.,, 2020).

At any point in time, estimating the score function is not trivial since we do not have access to the true score function. Fortunately, Hyvärinen, (2005) introduces the score matching objective, which surprisingly enables training a score model from data samples without the knowledge of the underlying score function. To access samples from the marginal $p^{\tau}$, we need to simulate the forward process from time $0$ to time $\tau$, as we only have clean data samples. This is costly in general, but if $f$ is affine, we can analytically reach any time $\tau$ in the forward process in a single step, by applying a Gaussian perturbation kernel $p^{0\tau}$ to clean data samples (Song et al.,, 2020). Since the kernel is differentiable, score matching simplifies to a denoising score matching objective (Vincent,, 2011),

where the expectation is over diffusion time $\tau$, noised sample $\mathbf{x}^{\tau}\sim p^{0\tau}(\mathbf{x}^{\tau}\mid\mathbf{x}^{0})$, obtained by applying the $\tau$-level perturbation kernel to a clean sample $\mathbf{x}^{0}\sim p^{data}(\mathbf{x}^{0})$. Importantly, as the kernel $p^{0\tau}$ is a known Gaussian distribution, this objective becomes a simple $L_{2}$ reconstruction loss,

with reparameterization $\mathbf{D}_{\theta}(\mathbf{x}^{\tau},\tau)=\mathbf{S}_{\theta}(\mathbf{x}^{\tau},\tau)\sigma^{2}(\tau)+\mathbf{x}^{\tau}$, where $\sigma(\tau)$ is the variance of the $\tau$-level perturbation kernel.

The score-based diffusion model described in Section 2.2 provides an unconditional generative model of $p_{data}$. To serve as a world model, we need a conditional generative model of the environment dynamics, $p(\mathbf{x}_{t+1}\mid\mathbf{x}_{\leq t},a_{\leq t})$, where we consider the general case of a pomdp, in which the Markovian state $s_{t}$ is unknown and can be approximated from past observations and actions. We can condition a diffusion model on this history, to estimate and generate the next observation directly, as shown in Figure 1. This modifies Equation 4 as follows,

During training, we sample a trajectory segment $\mathbf{x}_{\leq t}^{0},a_{\leq t},\mathbf{x}_{t+1}^{0}$ from the agent’s replay dataset, and we obtain the noised next observation $\mathbf{x}_{t+1}^{\tau}\sim p^{0\tau}(\mathbf{x}_{t+1}^{\tau}\mid\mathbf{x}_{t+1}^{0})$ by applying the $\tau$-level perturbation kernel. In summary, this diffusion process for world modeling resembles the standard diffusion process described in Section 2.2, with a score model conditioned on past observations and actions.

To sample the next observation, we iteratively solve the reverse SDE in Equation 2, as illustrated in Figure 1. While we can in principle resort to any ODE or SDE solver, there is an inherent trade-off between sampling quality and Number of Function Evaluations (NFE), that directly determines the inference cost of the diffusion world model (see Appendix A for more details).

Building on the background provided in Section 2, we now introduce diamond as a practical realization of a diffusion-based world model. In particular, we now define the drift and diffusion coefficients $\mathbf{f}$ and $g$ introduced in Section 2.2, corresponding to a particular choice of diffusion paradigm. While ddpm (Ho et al.,, 2020) is an example of one such choice (as described in Appendix B) and would historically be the natural candidate, we instead build upon the edm formulation proposed in Karras et al., (2022). The practical implications of this choice are discussed in Section 5.1. In what follows, we describe how we adapt edm to build our diffusion-based world model.

We consider the perturbation kernel $p^{0\tau}(\mathbf{x}_{t+1}^{\tau}\mid\mathbf{x}_{t+1}^{0})=\mathcal{N}(\mathbf{x}_{t+1}^{\tau};\mathbf{x}_{t+1}^{0},\sigma^{2}(\tau)\mathbf{I})$, where $\sigma(\tau)$ is a real-valued function of diffusion time called the noise schedule. This corresponds to setting the drift and diffusion coefficients to $\mathbf{f}(\mathbf{x},\tau)=\mathbf{0}$ (affine) and $g(\tau)=\sqrt{2\dot{\sigma}(\tau)\sigma(\tau)}$.

We use the network preconditioning introduced by Karras et al., (2022) and so parameterize $\mathbf{D}_{\theta}$ in Equation 5 as the weighted sum of the noised observation and the prediction of a neural network $\mathbf{F}_{\theta}$,

where for brevity we define $y_{t}^{\tau}\coloneqq(c_{\text{noise}}^{\tau},\mathbf{x}^{0}_{\leq t},a_{\leq t})$ to include all conditioning variables.

The preconditioners $c_{\text{in}}^{\tau}$ and $c_{\text{out}}^{\tau}$ are selected to keep the network’s input and output at unit variance for any noise level $\sigma(\tau)$, $c_{\text{noise}}^{\tau}$ is an empirical transformation of the noise level, and $c_{\text{skip}}^{\tau}$ is given in terms of $\sigma(\tau)$ and the standard deviation of the data distribution $\sigma_{\text{data}}$, as $c_{skip}^{\tau}=\sigma_{data}^{2}/(\sigma_{data}^{2}+\sigma^{2}(\tau))$. These preconditioners are fully described in Appendix C.

Combining Equations 5 and 6 provides insight into the training objective of $\mathbf{F}_{\theta}$,

The network training target adaptively mixes signal and noise depending on the degradation level $\sigma(\tau)$. When $\sigma(\tau)\gg\sigma_{\text{data}}$, we have $c_{\text{skip}}^{\tau}\to 0$, and the training target for $\mathbf{F}_{\theta}$ is dominated by the clean signal $\mathbf{x}_{t+1}^{0}$. Conversely, when the noise level is low, $\sigma(\tau)\to 0$, we have $c_{\text{skip}}^{\tau}\to 1$, and the target becomes the difference between the clean and the perturbed signal, i.e. the added Gaussian noise. Intuitively, this prevents the training objective to become trivial in the low-noise regime. In practice, this objective is high variance at the extremes of the noise schedule, so Karras et al., (2022) sample the noise level $\sigma(\tau)$ from an empirically chosen log-normal distribution in order to concentrate the training around medium-noise regions, as described in Appendix C.

We use a standard U-Net 2D for the vector field $\mathbf{F}_{\theta}$ (Ronneberger et al.,, 2015), and we keep a buffer of $L$ past observations and actions that we use to condition the model. We concatenate these past observations to the next noisy observation channel-wise, and we input actions through adaptive group normalization layers (Zheng et al.,, 2020) in the residual blocks (He et al.,, 2015) of the U-Net.

As discussed in Section 2.3 and Appendix A, there are many possible sampling methods to generate the next observation from the trained diffusion model. While our codebase supports a variety of sampling schemes, we found Euler’s method to be effective without incurring the cost of additional NFE required by higher order samplers, or the unnecessary complexity of stochastic sampling.

Given the diffusion model from Section 3.1, we now complete our world model with a reward and termination model, required for training an RL agent in imagination. Since estimating the reward and termination are scalar prediction problems, we use a separate model $R_{\psi}$ consisting of standard cnn (LeCun et al.,, 1989; He et al.,, 2015) and lstm (Hochreiter and Schmidhuber,, 1997; Gers et al.,, 2000) layers to handle partial observability. The RL agent involves an actor-critic network parameterized by a shared cnn-lstm with policy and value heads. The policy $\pi_{\phi}$ is trained with reinforce with a value baseline, and we use a Bellman error with $\lambda$-returns to train the value network $V_{\phi}$, similar to Micheli et al., (2023). We train the agent entirely in imagination as described in Section 2.1. The agent only interacts with the real environment for data collection. After each collection stage, the current world model is updated by training on all data collected so far. Then, the agent is trained with RL in the updated world model environment, and these steps are repeated. This procedure is detailed in Algorithm 1, and is similar to Kaiser et al., (2019); Hafner et al., (2020); Micheli et al., (2023). We provide architecture details, hyperparameters, and RL objectives in Appendices D, E, F, respectively.

For comprehensive evaluation of diamond we use the established Atari 100k benchmark (Kaiser et al.,, 2019), consisting of 26 games that test a wide range of agent capabilities. For each game, an agent is only allowed to take 100k actions in the environment, which is roughly equivalent to 2 hours of human gameplay, to learn to play the game before evaluation. As a reference, unconstrained Atari agents are usually trained for 50 million steps, a 500 fold increase in experience. We trained diamond from scratch for 5 random seeds on each game. Each run utilized around 12GB of VRAM and took approximately 2.9 days on a single Nvidia RTX 4090 (1.03 GPU years in total).

We compare with other recent methods training an agent entirely within a world model in Table 1, including storm (Zhang et al.,, 2023), DreamerV3 (Hafner et al.,, 2023), iris (Micheli et al.,, 2023), twm (Robine et al.,, 2023), and SimPle (Kaiser et al.,, 2019). A broader comparison to model-free and search-based methods, including bbf (Schwarzer et al.,, 2023) and EfficientZero (Ye et al.,, 2021), the current best performing methods on this benchmark, is provided in Appendix J. bbf and EfficientZero use techniques that are orthogonal and not directly comparable to our approach, such as using periodic network resets in combination with hyperparameter scheduling for bbf, and computationally expensive lookahead Monte-Carlo tree search for EfficientZero. Combining these additional components with our world model would be an interesting direction for future work.

Table 1 provides scores for all games, and the mean and interquartile mean (IQM) of human-normalized scores (HNS) (Wang et al.,, 2016). Following the recommendations of Agarwal et al., (2021) on the limitations of point estimates, we provide stratified bootstrap confidence intervals for the mean and IQM in Figure 2, as well as performance profiles and additional metrics in Appendix H.

Our results demonstrate that diamond performs strongly across the benchmark, outperforming human players on 11 games, and achieving a superhuman mean HNS of 1.46, a new best among agents trained entirely within a world model. diamond also achieves an IQM on par with storm, and greater than all other baselines. We find that diamond performs particularly well on environments where capturing small details is important, such as Asterix, Breakout and Road Runner. We provide further qualitative analysis of the visual quality of the world model in Section 5.3.

As explained in Section 2, we could in principle use any diffusion model variant in our world model. While diamond utilizes edm (Karras et al.,, 2022) as described in Section 3, ddpm (Ho et al.,, 2020) would also be a natural candidate, having been used in many image generation applications (Rombach et al.,, 2022; Nichol and Dhariwal,, 2021). We justify this design decision in this section.

To provide a fair comparison of ddpm with our edm implementation, we train both variants with the same network architecture, on a shared static dataset of 100k frames collected with an expert policy on the game Breakout. As discussed in Section 2.3, the number of denoising steps is directly related to the inference cost of the world model, and so fewer steps will reduce the cost of training an agent on imagined trajectories. Ho et al., (2020) use a thousand denoising steps, and Rombach et al., (2022) employ hundreds steps for Stable Diffusion. However, for our world model to be computationally comparable with other world model baselines (such as iris which requires 16 NFE for each timestep), we need at most tens of denoising steps, and preferably fewer. Unfortunately, if the number of denoising steps is set too low, the visual quality will degrade, leading to compounding error.

To investigate the stability of the diffusion variants, we display imagined trajectories generated autoregressively up to $t=1000$ timesteps in Figure 3, for different numbers of denoising steps $n\leq 10$. We see that using ddpm (Figure 3(a)) in this regime leads to severe compounding error, causing the world model to quickly drift out of distribution. In contrast, the edm-based diffusion world model (Figure 3(b)) appears much more stable over long time horizons, even for a single denoising step. A quantitative analysis of this compounding error is provided in Appendix K.

This surprising result is a consequence of the improved training objective described in Equation 7, compared to the simpler noise prediction objective employed by ddpm. While predicting the noise works well for intermediate noise levels, this objective causes the model to learn the identity function when the noise is dominant ($\sigma_{noise}\gg\sigma_{data}\implies\xi_{\theta}(\mathbf{x}^{\tau}_{t+1},y_{t}^{\tau})\to\mathbf{x}^{\tau}_{t+1}$), where $\xi_{\theta}$ is the noise prediction network of ddpm. This gives a poor estimate of the score function at the beginning of the sampling procedure, which degrades the generation quality and leads to compounding error.

In contrast, the adaptive mixing of signal and noise employed by edm, described in Section 3.1, means that the model is trained to predict the clean image when the noise is dominant ($\sigma_{noise}\gg\sigma_{data}\implies\mathbf{F_{\theta}}(\mathbf{x}^{\tau}_{t+1},y_{t}^{\tau})\to\mathbf{x}^{0}_{t+1}$). This gives a better estimate of the score function in the absence of signal, so the model is able to produce higher quality generations with fewer denoising steps, as illustrated in Figure 3(b).

While we found that our edm-based world model was very stable with just a single denoising step, as shown for Breakout in the last row of Figure 3(b), we discuss here how this choice would limit the visual quality of the model in some cases. We provide more a quantitative analysis in Appendix L.

As discussed in Section 2.2, our score model is equivalent to a denoising autoencoder (Vincent et al.,, 2008) trained with an $L_{2}$ reconstruction loss. The optimal single-step prediction is thus the expectation over possible reconstructions for a given noisy input, which can be out of distribution if this posterior distribution is multimodal. While some games like Breakout have deterministic transitions that can be accurately modeled with a single denoising step (see Figure 3(b)), in some other games partial observability gives rise to multimodal observation distributions. In this case, an iterative solver is necessary to drive the sampling procedure towards a particular mode, as illustrated in the game Boxing in Figure 4. As a result, we therefore set $n=3$ in all of our experiments.

We now compare to iris (Micheli et al.,, 2023), a well-established world model that uses a discrete autoencoder (Van Den Oord et al.,, 2017) to convert images to discrete tokens, and composes these tokens over time with an autoregressive transformer (Radford et al.,, 2019). For fair comparison, we train both world models on the same static datasets of 100k frames collected with expert policies. This comparison is displayed in Figure 2 below.

We see in Figure 2 that the trajectories imagined by diamond are generally of higher visual quality and more faithful to the true environment compared to the trajectories imagined by iris. In particular, the trajectories generated by iris contain visual inconsistencies between frames (highlighted by white boxes), such as enemies being displayed as rewards and vice-versa. These inconsistencies may only represent a few pixels in the generated images, but can have significant consequences for reinforcement learning. For example, since an agent should generally target rewards and avoid enemies, these small visual discrepancies can make it more challenging to learn an optimal policy.

These improvements in the consistency of visual details are generally reflected by greater agent performance on these games, as shown in Table 1. Since the agent component of these methods is similar, this improvement can likely be attributed to the world model.

Finally, we note that this improvement is not simply the result of increased computation. Both world models are rendering frames at the same resolution ($64\times 64$), and diamond requires only 3 NFE per frame compared to 16 NFE per frame for iris. This is further reflected by the fact that diamond has significantly fewer parameters and takes less time to train than iris, as provided in Appendix H.