Title: Diffusion for World Modeling: Visual Details Matter in Atari
Venue: NeurIPS 2024
Accepted: yes
Keywords: World models

diffusion models

reinforcement learning

generative models

Atari
Abstract: World models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static *Counter-Strike: Global Offensive* gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io.

Diffusion for World Modeling: Visual Details Matter in Atari †

Abstract

World models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static Counter-Strike: Global Offensive gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io .

1 Introduction

Generative models of environments, or 'world models" (Ha and Schmidhuber, 2018), are becoming increasingly important as a component for generalist agents to plan and reason about their environment (LeCun, 2022). Reinforcement Learning (RL) has demonstrated a wide variety of successes in recent years (Silver et al., 2016; Degrave et al., 2022; Ouyang et al., 2022), but is well-known to be sample inefficient, which limits real-world applications. World models have shown promise for training reinforcement learning agents across diverse environments (Hafner et al., 2023; Schrittwieser et al., 2020), with greatly improved sample-efficiency (Ye et al., 2021), which can enable learning from experience in the real world (Wu et al., 2023).

Recent world modeling methods (Hafner et al., 2021; Micheli et al., 2023; Robine et al., 2023; Hafner et al., 2023; Zhang et al., 2023) often model environment dynamics as a sequence of discrete latent variables. Discretization of the latent space helps to avoid compounding error over multi-step time horizons. However, this encoding may lose information, resulting in a loss of generality and reconstruction quality. This may be problematic for more real-world scenarios where the information

A A

A A

a

t

-

1

A A

A A

a

t

A A

A A

a

T

-

1

x

0

t

-

1

.

.

.

.

.

.

a

t

x

T

t

x

τ

t

x

0

t

x

0

t

.

.

.

.

.

.

a

t

+1

x

T

t

+1

x

τ

t

+1

x

0

t

+1

x

0

T

-

1

.

.

.

.

.

.

x

T

T

x

τ

T

x

0

T

π

ϕ

π

ϕ

D

θ

D

θ

D

θ

D

θ

D

θ

D

θ

. . .

Environment time (

t

)

Denoising time (

τ

)

Conditioning

required for the task is less well-defined, such as training autonomous vehicles (Hu et al., 2023). In this case, small details in the visual input, such as a traffic light or a pedestrian in the distance, may change the policy of an agent. Increasing the number of discrete latents can mitigate this lossy compression, but comes with an increased computational cost (Micheli et al., 2023).

In the meantime, diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020) have become a dominant paradigm for high-resolution image generation (Rombach et al., 2022; Podell et al., 2023). This class of methods, in which the model learns to reverse a noising process, challenges well-established approaches modeling discrete tokens (Esser et al., 2021; Ramesh et al., 2021; Chang et al., 2023), and thereby offers a promising alternative to alleviate the need for discretization in world modeling. Additionally, diffusion models are known to be easily conditionable and to flexibly model complex multi-modal distributions without mode collapse. These properties are instrumental to world modeling, since adherence to conditioning should allow a world model to reflect an agent's actions more closely, resulting in more reliable credit assignment, and modeling multi-modal distributions should provide greater diversity of training scenarios for an agent.

Motivated by these characteristics, we propose DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. Careful design choices are necessary to ensure our diffusion world model is efficient and stable over long-time horizons, and we provide a qualitative analysis to illustrate their importance. DIAMOND achieves a mean human normalized score of 1.46 on the well-established Atari 100k benchmark; a new state of the art for agents trained entirely within a world model. Additionally, operating in image space has the benefit of enabling our diffusion world model to be a drop-in substitute for the environment, which provides greater insight into world model and agent behaviors. In particular, we find the improved performance in some games follows from better modeling of critical visual details. To further demonstrate the effectiveness of our world model in isolation, we train DIAMOND's diffusion world model on 87 hours of static Counter-Strike: Global Offensive (CSGO) gameplay (Pearce and Zhu, 2022) to produce an interactive neural game engine for the popular in-game map, Dust II . We release our code, agents and playable world models at https://diamond-wm.github.io .

2 Preliminaries

2.1 Reinforcement learning and world models

We model the environment as a standard Partially Observable Markov Decision Process (POMDP) (Sutton and Barto, 2018), ( S , A , O , T, R, O, γ ) , where S is a set of states, A is a set of discrete actions, and O is a set of image observations. The transition function T : S×A×S → [0 , 1] describes the environment dynamics p ( s t +1 | s t , a t ) , and the reward function R : S × A × S → R maps transitions to scalar rewards. Agents cannot directly access states s t and only see the environment through image observations x t ∈ O , emitted according to observation probabilities p ( x t | s t ) ,

described by the observation function O : S × O → [0 , 1] . The goal is to obtain a policy π that maps observations to actions in order to maximize the expected discounted return E π [ ∑ t ≥ 0 γ t r t ] , where γ ∈ [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 | 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").

2.2 Score-based diffusion models

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

We consider a diffusion process { x τ } τ ∈ [0 , T ] indexed by a continuous time variable τ ∈ [0 , T ] , with corresponding marginals { p τ } τ ∈ [0 , T ] , and boundary conditions p 0 = p data and p T = p prior , where p prior is a tractable unstructured prior distribution, such as a Gaussian. Note that we use τ 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 w is the Wiener process (Brownian motion), 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 ¯ w is the reverse-time Wiener process, and ∇ x log p τ ( 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 ∇ x log p τ ( x ) , associated with marginals { p τ } τ ∈ [0 , T ] along the process. In practice, it is possible to use a single time-dependent score model S θ ( x , τ ) 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 τ , we need to simulate the forward process from time 0 to time τ , as we only have clean data samples. This is costly in general, but if f is affine, we can analytically reach any time τ in the forward process in a single step, by applying a Gaussian perturbation kernel p 0 τ 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 τ , noised sample x τ ∼ p 0 τ ( x τ | x 0 ) , obtained by applying the τ -level perturbation kernel to a clean sample x 0 ∼ p data ( x 0 ) . Importantly, as the kernel p 0 τ is a known Gaussian distribution, this objective becomes a simple L 2 reconstruction loss,

with reparameterization D θ ( x τ , τ ) = S θ ( x τ , τ ) σ 2 ( τ )+ x τ , where σ ( τ ) is the variance of the τ -level perturbation kernel.

2.3 Diffusion for world modeling

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 ( x t +1 | x ≤ t , a ≤ 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 x 0 ≤ t , a ≤ t , x 0 t +1 from the agent's replay dataset, and we obtain the noised next observation x τ t +1 ∼ p 0 τ ( x τ t +1 | x 0 t +1 ) by applying the τ -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).

3 Method

3.1 Practical choice of diffusion paradigm

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 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 τ ( x τ t +1 | x 0 t +1 ) = N ( x τ t +1 ; x 0 t +1 , σ 2 ( τ ) I ) , where σ ( τ ) is a real-valued function of diffusion time called the noise schedule. This corresponds to setting the drift and diffusion coefficients to f ( x , τ ) = 0 (affine) and g ( τ ) = √ 2˙ σ ( τ ) σ ( τ ) .

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

where for brevity we define y τ t := ( c τ noise , x 0 ≤ t , a ≤ t ) to include all conditioning variables.

The preconditioners c τ in and c τ out are selected to keep the network's input and output at unit variance for any noise level σ ( τ ) , c τ noise is an empirical transformation of the noise level, and c τ skip is given in terms of σ ( τ ) and the standard deviation of the data distribution σ data , as c τ skip = σ 2 data / ( σ 2 data + σ 2 ( τ )) . These preconditioners are fully described in Appendix C.

Combining Equations 5 and 6 provides insight into the training objective of F θ ,

The network training target adaptively mixes signal and noise depending on the degradation level σ ( τ ) . When σ ( τ ) ≫ σ data, we have c τ skip → 0 , and the training target for F θ is dominated by the clean signal x 0 t +1 . Conversely, when the noise level is low, σ ( τ ) → 0 , we have c τ skip → 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 σ ( τ ) 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 F θ (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.

3.2 Reinforcement learning in imagination

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 ψ 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 π ϕ is trained with REINFORCE with a value baseline, and we use a Bellman error with λ -returns to train the value network V ϕ , 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.

4 Experiments

4.1 Atari 100k benchmark

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.

4.2 Results on the Atari 100k benchmark

Table 1 provides scores for all games, and the mean and interquartile mean (IQM) of humannormalized 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.

0

.

4

0

.

8

1

.

2

1

.

6

SimPLe

TWM

IRIS

STORM

DIAMOND

Mean

0

.

2

0

.

4

0

.

6

Interquartile Mean

Human Normalized Score

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.

5 Analysis

5.1 Choice of diffusion framework

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 ≤ 10 . We see that using DDPM (Figure 3a) 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 3b) 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.

t=0

t=5

t=10

t=50

t=500

t=1000

t=0

t=5

t=10

t=50

t=500

t=1000

n=10

n=10

n=5

n=5

n=3

n=3

n= 1

n=1

(a) DDPM-based world model trajectories.

(b) EDM-based world model trajectories.

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 ( σ noise ≫ σ data = ⇒ ξ θ ( x τ t +1 , y τ t ) → x τ t +1 ), where ξ θ 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 ( σ noise ≫ σ data = ⇒ F θ ( x τ t +1 , y τ t ) → 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 3b.

5.2 Choice of the number of denoising steps

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 3b, 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 3b), 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.

t=0

t=1

t=2

t=3

t=4

t=5

n=1

27

n=3

5.3 Qualitative visual comparison with IRIS

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 5 below.

(a) IRIS

(b) DIAMOND

We see in Figure 5 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 × 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.

6 Scaling the diffusion world model to Counter-Strike: Global Offensive 4

To investigate the ability of DIAMOND's diffusion world model to learn to model more complex 3D environments, we train the world model in isolation on static data from the popular video game Counter-Strike: Global Offensive (CS:GO). We use the Online dataset of 5.5M frames (95 hours) of online human gameplay captured at 16Hz from the map Dust II by Pearce and Zhu (2022). We randomly hold out 0.5M frames (corresponding to 500 episodes, or 8 hours) for testing, and use the remaining 5M frames (87 hours) for training. There is no reinforcement learning agent or online data collection involved in these experiments.

To reduce the computational cost, we reduce the resolution from (280 × 150) to (56 × 30) for world modeling. We then introduce a second, smaller diffusion model as an upsampler to improve the generated images at the original resolution (Saharia et al., 2022b). We scale the channels of the U-Net to increase the number of parameters from 4M for our Atari models to 381M for our CS:GO model (including 51M for the upsampler). The combined model was trained for 12 days on an RTX 4090.

Finally, we introduce stochastic sampling and increase the number of denoising steps for the upsampler to 10, which we found to improve the resulting visual quality of the generations, while keeping the dynamics model the same (in particular, still using only 3 denoising steps). This enables a reasonable tradeoff between visual quality and inference cost, with the model running at 10Hz on an RTX 3090. Typical generations of the model are provided in Figure 6 below.

shoot&reload

muzzle flash & recoil

navigating tunnel

kill enemy & jump

sign in arabic

jump onto box

shadow&light

goingthroughdoors

We find the model is able to generate stable trajectories over hundreds of timesteps, although is more likely to drift out-of-distribution in less frequently visited areas of the map. Due to the limited memory of the model, approaching walls or losing visibility may cause the model to forget the current state and instead generate a new weapon or area of map. Interestingly, we find the model wrongly enables successive jumps by generalizing the effect of a jump on the geometry of the scene, since multiple jumps do not appear often enough in the training gameplay for the model to learn that mid-air jumps should be ignored. We expect scaling the model and data to address many of these limitations, with the exception of the memory of the model. Quantitative measurements of the capabilities of the CS:GO world model and attempts to address these limitations are left to future work.

7 Related work

World models. The idea of reinforcement learning (RL) in the imagination of a neural network world model was introduced by Ha and Schmidhuber (2018). SimPLe (Kaiser et al., 2019) applied world models to Atari, and introduced the Atari 100k benchmark to focus on sample efficiency. Dreamer (Hafner et al., 2020) introduced RL from the latent space of a recurrent state space model (RSSM). DreamerV2 (Hafner et al., 2021) demonstrated that using discrete latents could help to reduce compounding error, and DreamerV3 (Hafner et al., 2023) was able to achieve human-level performance on a wide range of domains with fixed hyperparameters. TWM (Robine et al., 2023) adapts DreamerV2's RSSM to use a transformer architecture, while STORM (Zhang et al., 2023) adapts DreamerV3 in a similar way but with a different tokenization approach. Alternatively, IRIS (Micheli et al., 2023) builds a language of image tokens with a discrete autoencoder, and composes these tokens over time with an autoregressive transformer.

Generative vision models. There are parallels between these world models and image generation models which suggests that developments in generative vision models could provide benefits to world modeling. Following the rise of transformers in natural language processing (Vaswani et al.,

2017; Devlin et al., 2018; Radford et al., 2019), VQGAN (Esser et al., 2021) and DALL·E (Ramesh et al., 2021) convert images to discrete tokens with discrete autoencoders (Van Den Oord et al., 2017), and leverage the sequence modeling abilities of autoregressive transformers to build powerful text-to-image generative models. Concurrently, diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020) gained traction (Dhariwal and Nichol, 2021; Rombach et al., 2022), and have become a dominant paradigm for high-resolution image generation (Saharia et al., 2022a; Ramesh et al., 2022; Podell et al., 2023).

The same trends have taken place in the recent developments of video generation methods. VideoGPT (Yan et al., 2021) provides a minimal video generation architecture by combining a discrete autoencoder with an autoregressive transformer. Godiva (Wu et al., 2021) enables text conditioning with promising generalization. Phenaki (Villegas et al., 2023) allows arbitrary length video generation with sequential prompt conditioning. TECO (Yan et al., 2023) improves upon autoregressive modeling by using MaskGit (Chang et al., 2022), and enables longer temporal dependencies by compressing input sequence embeddings. Diffusion models have also seen a resurgence in video generation using 3D U-Nets to provide high quality but short-duration video (Singer et al., 2023; Bar-Tal et al., 2024). Recently, transformer-based diffusion models such as DiT (Peebles and Xie, 2023) and Sora (Brooks et al., 2024) have shown improved scalability for both image and video generation, respectively.

Diffusion for reinforcement learning. There has also been much interest in combining diffusion models with reinforcement learning. This includes taking advantage of the flexibility of diffusion models as a policy (Wang et al., 2022; Ajay et al., 2022; Pearce et al., 2023), as planners (Janner et al., 2022; Liang et al., 2023), as reward models (Nuti et al., 2023), and trajectory modeling for data augmentation in offline RL (Lu et al., 2023; Ding et al., 2024; Jackson et al., 2024). DIAMOND represents the first use of diffusion models as world models for learning online in imagination.

Generative game engines. Playable games running entirely on neural networks have recently been growing in scope. GameGAN (Kim et al., 2020) learns generative models of games using a GAN (Goodfellow et al., 2014) while Bamford and Lucas (2020) use a Neural GPU (Kaiser and Sutskever, 2015). Concurrent work includes Genie (Bruce et al., 2024), which generates playable platformer environments from image prompts, and GameNGen (Valevski et al., 2024), which similarly leverages a diffusion model to obtain a high resolution simulator of the game DOOM, but at a larger scale.

8 Limitations

We identify three main limitations of our work for future research. First, our main evaluation is focused on discrete control environments, and applying DIAMOND to the continuous domain may provide additional insights. Second, the use of frame stacking for conditioning is a minimal mechanism to provide a memory of past observations. Integrating an autoregressive transformer over environment time, using an approach such as Peebles and Xie (2023), would enable longer-term memory and better scalability. We include an initial investigation into a potential cross-attention architecture in Appendix M, but found frame-stacking more effective in our early experiments. Third, we leave potential integration of the reward/termination prediction into the diffusion model for future work, since combining these objectives and extracting representations from a diffusion model is not trivial (Luo et al., 2023; Xu et al., 2023) and would make our world model unnecessarily complex.

9 Conclusion and Broader Impact

We have introduced DIAMOND, a reinforcement learning agent trained in a diffusion world model. We explained the key design choices we made to adapt diffusion for world modeling and to make our world model stable over long time horizons with a low number of denoising steps. DIAMOND achieves a mean human normalized score of 1 . 46 on the well-established Atari 100k benchmark; a new best among agents trained entirely within a world model. We analyzed our improved performance in some games and found that it likely follows from better modeling of critical visual details. We further demonstrated DIAMOND's diffusion world model can successfully model 3D environments and serve as a real-time neural game engine by training on static Counter-Strike: Global Offensive gameplay.

World models constitute a promising direction to address sample efficiency and safety concerns associated with training agents in the real world. However, imperfections in the world model may lead to suboptimal or unexpected agent behaviors. We hope that the development of more faithful and interactive world models will contribute to broader efforts to further reduce these risks.

Acknowledgments and Disclosure of Funding

We would like to thank Andrew Foong, Bálint Máté, Clément Vignac, Maxim Peter, Pedro Sanchez, Rich Turner, Stéphane Nguyen, Tom Lee, Trevor McInroe and Weipu Zhang for insightful discussions and comments. Adam and Eloi met during an internship at Microsoft Research Cambridge, and would like to thank the Game Intelligence team, including Anssi Kanervisto, Dave Bignell, Gunshi Gupta, Katja Hofmann, Lukas Schäfer, Raluca Georgescu, Sam Devlin, Sergio Valcarcel Macua, Shanzheng Tan, Tabish Rashid, Tarun Gupta, Tim Pearce, and Yuhan Cao, for their support in the early stages of this project, and a great summer.

References

Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C., and Bellemare, M. (2021). Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems , 34:29304-29320.

Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T., and Agrawal, P. (2022). Is conditional generative modeling all you need for decision-making? International Conference on Learning Representations .

Alonso, E., Jelley, A., Kanervisto, A., and Pearce, T. (2023). Diffusion world models. OpenReview .

Anderson, B. D. (1982). Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326.

Ascher, U. M. and Petzold, L. R. (1998). Computer methods for ordinary differential equations and differential-algebraic equations . Society for Industrial and Applied Mathematics.

Bamford, C. and Lucas, S. M. (2020). Neural game engine: Accurate learning of generalizable forward models from pixels. In 2020 IEEE Conference on Games (CoG) , pages 81-88. IEEE.

Bar-Tal, O., Chefer, H., Tov, O., Herrmann, C., Paiss, R., Zada, S., Ephrat, A., Hur, J., Li, Y ., Michaeli, T., et al. (2024). Lumiere: A space-time diffusion model for video generation. arXiv preprint arXiv:2401.12945 .

Brooks, T., Peebles, B., Holmes, C., DePue, W., Guo, Y., Jing, L., Schnurr, D., Taylor, J., Luhman, T., Luhman, E., Ng, C., Wang, R., and Ramesh, A. (2024). Video generation models as world simulators.

Bruce, J., Dennis, M. D., Edwards, A., Parker-Holder, J., Shi, Y ., Hughes, E., Lai, M., Mavalankar, A., Steigerwald, R., Apps, C., et al. (2024). Genie: Generative interactive environments. In Forty-first International Conference on Machine Learning .

Chang, H., Zhang, H., Barber, J., Maschinot, A., Lezama, J., Jiang, L., Yang, M.-H., Murphy, K. P., Freeman, W. T., Rubinstein, M., Li, Y., and Krishnan, D. (2023). Muse: Text-to-image generation via masked generative transformers. In International Conference on Machine Learning , pages 4055-4075. PMLR.

Chang, H., Zhang, H., Jiang, L., Liu, C., and Freeman, W. T. (2022). Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11315-11325.

Çiçek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., and Ronneberger, O. (2016). 3d u-net: learning dense volumetric segmentation from sparse annotation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19 , pages 424-432. Springer.

Degrave, J., Felici, F., Buchli, J., Neunert, M., Tracey, B., Carpanese, F., Ewalds, T., Hafner, R., Abdolmaleki, A., de Las Casas, D., et al. (2022). Magnetic control of tokamak plasmas through deep reinforcement learning. Nature , 602:414-419.

Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 .

Dhariwal, P. and Nichol, A. (2021). Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794.

Ding, Z., Zhang, A., Tian, Y., and Zheng, Q. (2024). Diffusion world model. arXiv preprint arXiv:2402.03570 .

Elfwing, S., Uchibe, E., and Doya, K. (2018). Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks , 107:3-11.

Esser, P., Rombach, R., and Ommer, B. (2021). Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 12873-12883.

Gers, F. A., Schmidhuber, J., and Cummins, F. (2000). Learning to forget: Continual prediction with LSTM. Neural Computation , 12(10):2451-2471.

Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. (2014). Generative adversarial nets. Advances in Neural Information Processing Systems , 27.

Ha, D. and Schmidhuber, J. (2018). Recurrent world models facilitate policy evolution. Advances in Neural Information Processing Systems , 31:2451-2463.

Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. (2020). Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations .

Hafner, D., Lillicrap, T. P., Norouzi, M., and Ba, J. (2021). Mastering atari with discrete world models. In International Conference on Learning Representations .

Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T. (2023). Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 .

He, K., Zhang, X., Ren, S., and Sun, J. (2015). Deep residual learning for image recognition. arXiv preprint arXiv:1512.03385 .

Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. (2017). Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in Neural Information Processing Systems , 30.

Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851.

Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, D. J. (2022). Video diffusion models. URL https://arxiv. org/abs/2204.03458 .

Hochreiter, S. and Schmidhuber, J. (1997). Long short-term memory. Neural Computation , 9(8):17351780.

Hu, A., Russell, L., Yeo, H., Murez, Z., Fedoseev, G., Kendall, A., Shotton, J., and Corrado, G. (2023). Gaia-1: A generative world model for autonomous driving. arXiv preprint arXiv:2309.17080 .

Hyvärinen, A. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6:695-709.

Jackson, M. T., Matthews, M. T., Lu, C., Ellis, B., Whiteson, S., and Foerster, J. (2024). Policy-guided diffusion. arXiv preprint arXiv:2404.06356 .

Janner, M., Du, Y., Tenenbaum, J., and Levine, S. (2022). Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning , pages 9902-9915. PMLR.

Kaiser, L., Babaeizadeh, M., Milos, P., Osinski, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., et al. (2019). Model-based reinforcement learning for atari. arXiv preprint arXiv:1903.00374 .

Kaiser, Ł. and Sutskever, I. (2015). Neural gpus learn algorithms. arXiv preprint arXiv:1511.08228 .

Kapturowski, S., Ostrovski, G., Quan, J., Munos, R., and Dabney, W. (2018). Recurrent experience replay in distributed reinforcement learning. International Conference on Learning Representations .

Karras, T., Aittala, M., Aila, T., and Laine, S. (2022). Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems , 35:26565-26577.

Kim, S. W., Zhou, Y., Philion, J., Torralba, A., and Fidler, S. (2020). Learning to simulate dynamic environments with gamegan. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1231-1240.

LeCun, Y. (2022). A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. OpenReview .

LeCun, Y., Boser, B., Denker, J. S., Henderson, D., Howard, R. E., Hubbard, W., and Jackel, L. D. (1989). Backpropagation applied to handwritten zip code recognition. Neural computation , 1(4):541-551.

Liang, Z., Mu, Y., Ding, M., Ni, F., Tomizuka, M., and Luo, P. (2023). Adaptdiffuser: Diffusion models as adaptive self-evolving planners. International Conference on Machine Learning .

Lu, C., Ball, P. J., and Parker-Holder, J. (2023). Synthetic experience replay. arXiv preprint arXiv:2303.06614 .

Luo, G., Dunlap, L., Park, D. H., Holynski, A., and Darrell, T. (2023). Diffusion hyperfeatures: Searching through time and space for semantic correspondence. In Advances in Neural Information Processing Systems .

Micheli, V., Alonso, E., and Fleuret, F. (2023). Transformers are sample-efficient world models. International Conference on Learning Representations .

Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D., and Kavukcuoglu, K. (2016). Asynchronous methods for deep reinforcement learning. In Balcan, M. F. and Weinberger, K. Q., editors, Proceedings of The 33rd International Conference on Machine Learning , volume 48 of Proceedings of Machine Learning Research , pages 1928-1937, New York, New York, USA.

Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015). Human-level control through deep reinforcement learning. Nature , 518(7540):529-533.

Nichol, A. Q. and Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. International Conference on Machine Learning .

Nuti, F., Franzmeyer, T., and Henriques, J. F. (2023). Extracting reward functions from diffusion models. arXiv preprint arXiv:2306.01804 .

Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. (2022). Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems , 35:27730-27744.

Pearce, T., Rashid, T., Kanervisto, A., Bignell, D., Sun, M., Georgescu, R., Macua, S. V., Tan, S. Z., Momennejad, I., Hofmann, K., and Devlin, S. (2023). Imitating human behaviour with diffusion models. The Eleventh International Conference on Learning Representations .

Pearce, T. and Zhu, J. (2022). Counter-strike deathmatch with large-scale behavioural cloning. In 2022 IEEE Conference on Games (CoG) , pages 104-111. IEEE.

Peebles, W. and Xie, S. (2023). Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , pages 4195-4205.

Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Müller, J., Penna, J., and Rombach, R. (2023). Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952 .

Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. (2019). Language models are unsupervised multitask learners.

Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. (2022). Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 1(2):3.

Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., and Sutskever, I. (2021). Zero-shot text-to-image generation. International Conference on Machine Learning .

Robine, J., Höftmann, M., Uelwer, T., and Harmeling, S. (2023). Transformer-based world models are happy with 100k interactions. International Conference on Learning Representations .

Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. (2022). High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 10684-10695.

Ronneberger, O., Fischer, P., and Brox, T. (2015). U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer.

Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. (2022a). Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems , 35:3647936494.

Saharia, C., Ho, J., Chan, W., Salimans, T., Fleet, D. J., and Norouzi, M. (2022b). Image superresolution via iterative refinement. IEEE transactions on pattern analysis and machine intelligence .

Santana, E. and Hotz, G. (2016). Learning a driving simulator. arXiv preprint arXiv:1608.01230 .

Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. (2020). Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588(7839):604-609.

Schwarzer, M., Ceron, J. S. O., Courville, A., Bellemare, M. G., Agarwal, R., and Castro, P. S. (2023). Bigger, better, faster: Human-level atari with human-level efficiency. International Conference on Machine Learning .

Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. (2016). Mastering the game of go with deep neural networks and tree search. Nature , 529:484-489.

Singer, U., Polyak, A., Hayes, T., Yin, X., An, J., Zhang, S., Hu, Q., Yang, H., Ashual, O., Gafni, O., et al. (2023). Make-a-video: Text-to-video generation without text-video data. International Conference on Learning Representations .

Skorokhodov, I., Tulyakov, S., and Elhoseiny, M. (2022). Stylegan-v: A continuous video generator with the price, image quality and perks of stylegan2. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 3626-3636.

Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. (2015). Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning .

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2020). Scorebased generative modeling through stochastic differential equations. International Conference on Learning Representations .

Sutton, R. S. (1991). Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin , 2(4):160-163.

Sutton, R. S. and Barto, A. G. (2018). Reinforcement learning: An introduction . MIT press.

Unterthiner, T., Van Steenkiste, S., Kurach, K., Marinier, R., Michalski, M., and Gelly, S. (2018). Towards accurate generative models of video: A new metric & challenges. arXiv preprint arXiv:1812.01717 .

Valevski, D., Leviathan, Y., Arar, M., and Fruchter, S. (2024). Diffusion models are real-time game engines.

Van Den Oord, A., Vinyals, O., et al. (2017). Neural discrete representation learning. Advances in Neural Information Processing Systems , 30.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems , 30.

Villegas, R., Babaeizadeh, M., Kindermans, P.-J., Moraldo, H., Zhang, H., Saffar, M. T., Castro, S., Kunze, J., and Erhan, D. (2023). Phenaki: Variable length video generation from open domain textual descriptions. International Conference on Learning Representations .

Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674.

Vincent, P., Larochelle, H., Bengio, Y., and Manzagol, P.-A. (2008). Extracting and composing robust features with denoising autoencoders. In International Conference on Machine learning .

Wang, Z., Hunt, J. J., and Zhou, M. (2022). Diffusion policies as an expressive policy class for offline reinforcement learning. International Conference on Learning Representations .

Wang, Z., Schaul, T., Hessel, M., Hasselt, H., Lanctot, M., and Freitas, N. (2016). Dueling network architectures for deep reinforcement learning. International Conference on Machine Learning .

Wu, C., Huang, L., Zhang, Q., Li, B., Ji, L., Yang, F., Sapiro, G., and Duan, N. (2021). Godiva: Generating open-domain videos from natural descriptions. arXiv preprint arXiv:2104.14806 .

Wu, P., Escontrela, A., Hafner, D., Abbeel, P., and Goldberg, K. (2023). Daydreamer: World models for physical robot learning. In Conference on Robot Learning , pages 2226-2240. PMLR.

Wu, Y. and He, K. (2018). Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV) .

Xu, J., Liu, S., Vahdat, A., Byeon, W., Wang, X., and De Mello, S. (2023). Open-vocabulary panoptic segmentation with text-to-image diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 2955-2966.

Yan, W., Hafner, D., James, S., and Abbeel, P. (2023). Temporally consistent transformers for video generation. International Conference on Machine Learning .

Yan, W., Zhang, Y., Abbeel, P., and Srinivas, A. (2021). Videogpt: Video generation using vq-vae and transformers. arXiv preprint arXiv:2104.10157 .

Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y. (2021). Mastering atari games with limited data. Advances in Neural Information Processing Systems , 34.

Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. (2018). The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 586-595.

Zhang, W., Wang, G., Sun, J., Yuan, Y., and Huang, G. (2023). Storm: Efficient stochastic transformer based world models for reinforcement learning. In Thirty-seventh Conference on Neural Information Processing Systems .

Zheng, H., Fu, J., Zeng, Y., Luo, J., and Zha, Z.-J. (2020). Learning semantic-aware normalization for generative adversarial networks. In Advances in Neural Information Processing Systems .

A Sampling observations in DIAMOND

We describe here how we sample an observation x 0 t from our diffusion world model. We initialize the procedure with a noisy observation x T t ∼ p prior , and iteratively solve the reverse SDE in Equation 2 from τ = T to τ = 0 , using the learned score model S θ ( x τ t , τ, x 0 <t , a <t ) conditioned on past observations x 0 <t and actions a <t . This procedure is illustrated in Figure 1.

In fact, there are many possible sampling methods for a given learned score model S θ (Karras et al., 2022). Notably, Song et al. (2020) introduce a corresponding 'probability flow" ordinary differential equation (ODE), with marginals equivalent to the stochastic process described in Section 2.2. In that case, the solving procedure is deterministic, and the only randomness comes from sampling the initial condition. In practice, this means that for a given score model, we can resort to any ODE or SDE solver, from simple first order methods like Euler (deterministic) and Euler-Maruyama (stochastic) schemes, to higher-order methods like Heun's method (Ascher and Petzold, 1998).

Regardless of the choice of solver, each step introduces truncation errors, resulting from the local score approximation and the discretization of the continuous process. Higher order samplers may reduce this truncation error, but come at the cost of additional Number of Function Evaluations (NFE) - how many forward passes of the network are required to generate a sample. This local error generally scales superlinearly with respect to the step size (for instance Euler's method is O ( h 2 ) for step size h ), so increasing the number of denoising steps improves the visual quality of the generated next frame. Therefore, there is a trade-off between visual quality and NFE that directly determines the inference cost of the diffusion world model.

B Link between DDPM and continuous-time score-based diffusion models

Denoising Diffusion Probabilistic Models (DDPM, Ho et al. (2020)) can be described as a discrete version of the diffusion process introduced in Section 2.2, as described in Song et al. (2020). The discrete forward process is a Markov chain characterized by a discrete noise schedule 0 < β 1 , . . . , β i , . . . β N < 1 , and a variance-preserving Gaussian transition kernel,

In the continuous time limit N →∞ , the Markov chain becomes a diffusion process, and the discrete noise schedule becomes a time-dependent function β ( τ ) . This diffusion process can be described by an SDE with drift coefficient f ( x , τ ) = -1 2 β ( τ ) x and diffusion coefficient g ( τ ) = √ β ( τ ) (Song et al., 2020).

C EDM network preconditioners and training

Karras et al. (2022) use the following preconditioners for normalization and rescaling purposes (as mentioned in Section 3.1) to improve network training:

where σ data = 0 . 5 .

The noise parameter σ ( τ ) is sampled to maximize the effectiveness of training as follows:

where P mean = -0 . 4 , P std = 1 . 2 . Refer to Karras et al. (2022) for an in-depth analysis.

D Model architectures

The diffusion model D θ is a standard U-Net 2D (Ronneberger et al., 2015), conditioned on the last 4 frames and actions, as well as the diffusion time τ . We use frame stacking for observation conditioning, and adaptive group normalization (Zheng et al., 2020) for action and diffusion time conditioning.

The reward/termination model R ψ layers are shared except for the final prediction heads. The model takes as input a sequence of frames and actions, and forwards it through convolutional residual blocks (He et al., 2015) followed by an LSTM cell (Mnih et al., 2016; Hochreiter and Schmidhuber, 1997; Gers et al., 2000). Before starting the imagination procedure, we burn-in (Kapturowski et al., 2018) the conditioning frames and actions to initialize the hidden and cell states of the LSTM.

The weights of the policy π ϕ and value network V ϕ are shared except for the last layer. In the following, we refer to ( π, V ) ϕ as the "actor-critic" network, even though V is technically a statevalue network, not a critic. This network takes as input a frame, and forwards it through convolutional trunk followed by an LSTM cell. The convolutional trunk consists of four residual blocks and 2x2 max-pooling with stride 2. The main path of the residual blocks consists of a group normalization (Wu and He, 2018) layer, a SiLU activation (Elfwing et al., 2018), and a 3x3 convolution with stride 1 and padding 1. Before starting the imagination procedure, we burn-in the conditioning frames to initialize the hidden and cell states of the LSTM.

Please refer to Table 2 below for hyperparameter values, and to Algorithm 1 for a detailed summary of the training procedure.

E Training hyperparameters

F Reinforcement learning objectives

In what follows, we note x t , r t and d t the observations, rewards, and boolean episode terminations predicted by our world model. We note H the imagination horizon, V ϕ the value network, π ϕ the policy network, and a t the actions taken by the policy within the world model.

We use λ -returns to balance bias and variance as the regression target for the value network. Given an imagined trajectory of length H , we can define the λ -return recursively as follows,

The value network V ϕ is trained to minimize L V ( ϕ ) , the expected squared difference with λ -returns over imagined trajectories,

where sg( · ) denotes the gradient stopping operation, meaning that the target is a constant in the gradient-based optimization, as classically established in the literature (Mnih et al., 2015; Hafner et al., 2021; Micheli et al., 2023).

As we can generate large amounts of on-policy trajectories in imagination, we use a simple REINFORCE objective to train the policy, with the value V ϕ ( x t ) as a baseline to reduce the variance of the gradients (Sutton and Barto, 2018). The policy is trained to minimize the following objective, combining REINFORCE and a weighted entropy maximization objective to maintain sufficient exploration,

G DIAMOND algorithm

We summarize the overall training procedure of DIAMOND in Algorithm 1 below. We denote as D the replay dataset where the agent stores data collected from the real environment, and other notations are introduced in previous sections or are self-explanatory.

H Additional performance comparisons

We provide performance profiles (Agarwal et al., 2021) for DIAMOND and baselines below.

0

1

2

3

4

5

6

7

8

Human Normalized Score (

τ

)

0.00

0.25

0.50

0.75

1.00

Fraction of runs with score  >

τ

SimPLe

TWM

IRIS

STORM

DIAMOND (ours)

As additional angles of comparison, we also provide parameter counts and approximate training times for IRIS, DreamerV3 and DIAMOND in Table 4 below. We see that DIAMOND has the highest mean HNS, with fewer parameters than both IRIS and DreamerV3. DIAMOND also trains faster than IRIS, although is slower than DreamerV3.

A full training time profile for DIAMOND is provided in Appendix I.

I Training time profile

Table 5 provides a full training time profile for DIAMOND.

J Broader comparison to model-free and search-based methods

Table 6 provides scores for model-free and search-based methods, including the current best performing methods on the Atari 100k benchmark, EfficientZero (Ye et al., 2021) and BBF (Schwarzer et al., 2023). Both of these methods use approaches that are out of scope of our approach, such as computationally expensive lookahead Monte-Carlo tree search for EfficientZero, and using periodic network resets in combination with hyperparameter scheduling for BBF. We see that while the use of lookahead search and more advanced reinforcement learning techniques (for EfficientZero (Ye et al., 2021) and BBF (Schwarzer et al., 2023) respectively) can still provide greater performance overall, DIAMOND promisingly still outperforms these methods on some games.

K Quantitative analysis of autoregressive model drift

Figure 8 provides a quantitative measure of the compounding error demonstrated qualitatively in Figure 3 for DDPM and EDM based world models.

DDPMn=1

1.4

DDPMn=3

==

distance

DDPMn=5

1.2

DDPMn=10

EDM

n=1

1.0

一

EDM

n=3

pixel

EDM

n=5

0.8

EDM

n=10

Average

0.6

0.4

0.2

0.0

1

5

10

50

100

500

1000

Timestep

L Quantitative ablation on reducing the number of denoising steps

Table 7 provides a quantitative ablation of the effect of reducing the number of denoising steps used for our EDM diffusion world model from 3 (used for Table 1) to 1, for DIAMOND's 10 highest performing games. Note that the 1-step results correspond to a single seed only so will have higher variance. Nonetheless, these results provide some signal that agents trained with 1 denoising step perform worse than our default choice of 3, particularly for the game Boxing , despite the apparent similarity in Figure 8. This additional evidence supports our qualitative analysis in Section 5.2.

M Early investigations on visual quality in more complex environments

In the main body of the paper, we evaluated the utility of DIAMOND for the purpose of training RL agents in a world model on the well-established Atari 100k benchmark (Kaiser et al., 2019), and demonstrated DIAMOND's diffusion world model could be applied to model a more complex 3D environment from the game Counter-Strike: Global Offensive . In this section, we provide early experiments investigating the effectiveness of DIAMOND's diffusion world model by directly evaluating the visual quality of the trajectories they generate. The two environments we consider are presented in Section M.1 below.

M.1 Environments

CS:GO. We use the Counter-Strike: Global Offensive dataset introduced by Pearce and Zhu (2022). Here we use the Clean dataset containing 190k frames (3.3 hours) of high-skill human gameplay, captured on the Dust II map. This contains observations and actions (mouse and keyboard) captured at 16Hz. We use 150k frames (2.6 hours) for training and 40k frames (0.7 hours) for evaluation. We resize observations to 64 × 64 pixels, and use no augmentation.

Motorway driving. We use the dataset from Santana and Hotz (2016) 5 , which contains camera and metadata captured from human drivers on US motorways. We select only trajectories captured in daylight, and exclude the first and last 5 minutes of each trajectory (typically traveling to/from a motorway), leaving 4.4 hours of data. We use five trajectories for training (3.6 hours) and two for testing (0.8 hours). We downsample the dataset to 10Hz, resize observations to 64 × 64, and for actions use the (normalized) steering angle and acceleration. During training, we apply data augmentation of shift & scale, contrast, brightness, and saturation, and mirroring.

We note that the purpose of our investigation is to train and evaluate DIAMOND's diffusion model on these static datasets, and that we do not perform reinforcement learning, since there is no standard reinforcement learning protocol for these environments.

M.2 Diffusion Model Architectures

We consider two potential diffusion model architectures, summarized in Figure 9.

Action

1) Frame-stacking

Cross-attention

Action

…

2) Cross-attention

Diffusion world model architectures

Self-attention

…

U-Net 3D

Diffusion video generation

Frame-stacking. The simplest way to condition on previous observations is by concatenating the previous L frames together with the next noised frame, concat[ x τ t , x 0 t -1 , . . . , x 0 t -L ] , which is compatible with a standard U-Net 2D (Ronneberger et al., 2015). This architecture is particularly attractive due to its lightweight construction, requiring minimal additional parameters and compute compared to typical image diffusion. This is the architecture we used for the main body of the paper.

Cross-attention. The U-Net 3D (Çiçek et al., 2016), also displayed for comparison in Figure 9, is a leading architecture in video diffusion (Ho et al., 2022). We adapted this design to have an autoregressive cross-attention architecture, formed of a core U-Net 2D, that only receives a single noised frame as direct input, but which cross-attends to the activations of a separate history encoder network. This encoder is a lightweight version of the U-Net 2D architecture. Parameters are shared

for all L encoders, and each receives the relative environment timestep embedding as input. The final design differs from the U-Net 3D which diffuses all frames jointly, shares parameters across networks, and uses self-, rather than cross-, attention.

M.3 Metrics, Baselines and Compute

Metrics. To evaluate the visual quality of generated trajectories, we use the standard Fréchet Video Distance ( FVD ) (Unterthiner et al., 2018) as implemented by Skorokhodov et al. (2022). This is computed between 1024 real videos (taken from the test set), and 1024 generated videos, each 16 frames long (1-2 seconds). Models condition on L = 6 previous real frames, and the real action sequence. On this same data, we also report the Fréchet Inception Distance ( FID ) (Heusel et al., 2017), which measures the visual quality of individual observations, ignoring the temporal dimension. For these same sets of videos, we also compute the LPIPS loss (Zhang et al., 2018) between each pair of real/generated observations (Yan et al., 2023). Sampling rate describes the number of observations that can be generated, in sequence, by a single Nvidia RTX A6000 GPU, per second.

Baselines. We compare against two well-established world model methods; DreamerV3 (Hafner et al., 2023) and IRIS (Micheli et al., 2023), adapting the original implementations to train on a static dataset. We ensured baselines used a similar number of parameters to DIAMOND. Two variants of IRIS are reported; image observations are discretized into K = 16 tokens (as used in the original work), or into K = 64 tokens (achieved with one less down/up-sampling layer in the autoencoder, see Appendix E of Micheli et al. (2023)), which provide the potential for modeling higher-fidelity visuals.

Compute. All models (baselines and DIAMOND) were trained for 120k updates with a batch size of 64, on up to 4 × A6000 GPUs. Each training run took between 1-2 days.

M.4 Analysis

Table 8: Results for 3D environments. These metrics compare observations from real trajectories and generated trajectories. The generated trajectories are conditioned on an initial set of L = 6 observations and a real sequence of actions.

Table 8 reports metrics on the visual quality of generated trajectories, along with sampling rates and number of parameters, for the frame-stack and cross-attention DIAMOND architectures, compared to baseline methods. DIAMOND outperforms the baselines across all visual quality metrics. This validates the results seen in the wider video generation literature, where diffusion models currently lead, as discussed in Section 7. The simpler frame-stacking architecture performs better than cross-attention, something surprising given the prevalence of cross-attention in the video generation literature. We believe the inductive bias provided by directly feeding in the input, frame-wise, may be well suited to autoregressive generation. Overall, these results indicate DIAMOND frame-stack > DIAMOND cross-attention ≈ IRIS 64 > IRIS 16 > DreamerV3, which we found corresponds to our intuition from visual inspection.

In terms of sampling rate, DIAMOND frame-stack (with 20 denoising steps) is faster than IRIS ( K = 16 ). IRIS suffers from a further 2.8 × slow down for the K = 64 version, verifying its sample time is bottlenecked by the number of tokens K . On the other hand, DreamerV3 is an order of magnitude faster - this derives from its independent, rather than joint, sampling procedure, and the flip-side of this is the low visual quality of its trajectories.

Figure 10 below shows selected examples of the trajectories produced by DIAMOND in CS:GO and motorway driving. The trajectories are plausible, often even at time horizons of reasonable length. In CS:GO, the model accurately generates the correct geometry of the level as it passes through the doorway into a new area of the map. In motorway driving, a car is plausibly imagined overtaking on the left.

While the above experiments use real sequences of actions from the dataset, we also investigated how robust DIAMOND (frame stack) was to novel, user-input actions. Figure 11 shows the effect of the actions in motorway driving - conditioned on the same L = 6 real frames, we generate trajectories conditioned on five different action sequences. In general the effects are as intended, e.g. steer straight/left/right moves the camera as expected. Interestingly, when 'slow down' is input, the distance to the car in front decreases since the model predicts that the traffic ahead has come to a standstill. Figure 12 shows similar sequences for CS:GO. For the common actions (mouse movements and fire), the effects are as expected, though they are unstable beyond a few frames, since such a sequence of actions is unlikely to have been seen in the demonstration dataset. We note that these issues - the causal confusion and instabilities - are a symptom of training world models on offline data, rather than being an inherent weakness of DIAMOND.

Continue

Steer

Steer

Slow

Speed

straight

left

right

down

up

t=0

t=25

t=50

Look

Look

Look

Look

Fire

left

right

down

up

t=3

t=8

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: All claims made in the title, abstract and introduction are limited to the paper's scope on Atari 100k, and match our results in Sections 4 and 5. We include a motivational example from the real world in the introduction but it is clear that this scenario is motivational and not addressed by the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of our work in a separate "Limitations" section (see Section 8), and analyze the trade-offs in computational efficiency in Section 5.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: We do not have novel theoretical results, but instead apply the existing theoretical framework of score-based diffusion models to world modeling. Sections 2 and 3 provide the necessary theoretical background to diffusion for RL practitioners in a clearly stated and referenced way.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide all the necessary details to reproduce our results. Our method is explained in Section 3, with the full training procedure detailed in Algorithm 1 in Appendix G. We provide full details of architectures in Appendix D, training hyperparameters in Appendix E and reinforcement learning objectives in Appendix F. Additionally our entire codebase will be open-source, and we provide an anonymized version with this submission: https://github.com/eloialonso/diamond .

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The environments we consider are open-source, and we provide open access to our anonymized codebase here: https://github.com/eloialonso/diamond . Our codebase contains detailed installation and reproducibility instructions. It also includes a feature to interactively play inside DIAMOND's world model.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Our full training procedure is detailed in Algorithm 1 in Appendix G. We provide all hyperparameters, including environment configurations, number of epochs, batch sizes, data collection schedules, optimizer settings, architectures and reinforcement learning details in Appendices D, E and F. All these details can also be found in our (anonymized) open-source codebase: https://github.com/eloialonso/diamond .

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We follow the recommended evaluation protocol of Agarwal et al. (2021) to provide confidence intervals for our experimental results in Figure 2 and performance profiles in Appendix H.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We detail the type of GPU, memory and time of execution in Section 4.1, and provide a performance comparison in Table 4.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We do not foresee any harms caused by the research process (there are no human subjects involved in our work, and the environments we use are open-source). We also do not foresee any direct harmful societal impact of our work, although we discuss the potential broader impact of our work in Section 9.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Our introduction discuss the general motivation for our research, and position it into a broader picture. We consider both positive and negative societal impacts in Section 9.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We only release pretrained models on Atari games, and do not believe these to pose such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We've made our best effort to include all relevant citations and licenses in our paper and codebase.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide detailed documentation alongside our codebase.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
The paper proposes a visual diffusion-based world model for reinforcement learning. The authors argue that current world models using discrete latent variables may lose important visual details, which are critical for reinforcement learning tasks. DIAMOND addresses this by generating images in the original pixel space. The paper evaluates DIAMOND on the Atari 100k benchmark and achieves strong results relative to other choices for the world model.

STRENGTHS:
- Clear and well-written paper
- Strong empirical results beating prior MBRL baselines on Atari 100K
- Impressive scaled-up experiments for visual diffusion world modeling, with consistent visual quality over long horizons even with a low number of diffusion steps
- Good set of ablations over key design choices of the algorithm 
- Publicly available code for the experiments

WEAKNESSES:
- It would be useful to include a comparison of training time for DIAMOND compared to prior baselines, e.g. what is the change in speed by using a diffusion world model v.s. Latent world model?
- How does the algorithm compare to the current SOTA model-free algorithms on Atari 100K, e.g. [1]. This would help the author understand the relative utility of MBRL for this setting

Minor:
- Theoretical diffusion background is lengthy and the paper would be better served by including results on other environments that are currently in the appendix, or discussion on architecture used

[1] Bigger, Better, Faster: Human-level Atari with human-level efficiency. Max Schwarzer, Johan Obando-Ceron, Aaron Courville, Marc Bellemare, Rishabh Agarwal, Pablo Samuel Castro.

QUESTIONS:
See weaknesses.

LIMITATIONS:
Discussion in paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces DIAMOND (DIffusion As a Model Of eNvironment Dreams), a novel RL agent trained within a diffusion-based world model. DIAMOND's world model is a diffusion model that, conditioned on past observations and actions, generates an observation at the next time step. This approach diverges from previous methods that rely on discrete latent variables, and offers a promising alternative that leverages the strengths of diffusion models, such as the ability to model multi-modal distributions. The training process is iterative, involving three key steps: data collection in the real environment, world model training on this collected data, and RL agent training with rollouts using the learned world model. The authors analyze design choices for the diffusion world model, and demonstrate that the EDM formulation is superior to DDPM for their use case through visualization of generated trajectories. They also conduct qualitative experiments to determine the optimal number of denoising steps for sampling from the world model. Finally, DIAMOND's evaluation on the Atari 100k benchmark reveals superior performance compared to other world model-based RL agents, achieving a new state-of-the-art mean human-normalized score of 1.46.\n\nSTRENGTHS:\nThis paper introduces a new new diffusion-based approach to world modeling in reinforcement learning, and show that it is a promising alternative to prior approaches based on discrete latent variable methods. The proposed method achieves state-of-the-art results on the Atari 100k benchmark and surpasses other world model-based RL agents. The paper also provides a thorough qualitative analysis of the design choices involved in creating an effective diffusion world model, including the choice of diffusion framework (EDM over DDPM) and the number of denoising steps. Overall, the paper is well-written and explains the concepts in a clear manner.\n\nWEAKNESSES:\nOne minor concern is the lack of quantitative analysis for the design choices. While qualitative experiments and the results on the Atari-100k benchmark effectively motivate the design choices, a more thorough quantitative analysis across the different tasks would provide stronger evidence for these choices. Such a quantitative analysis would be potentially important for applications in real-world environments.\n\nQUESTIONS:\n- The performance on a few tasks (BankHeist, Frostbite, UpNDown) are considerably worse than the baselines. Have you investigated why this is the case?\n- Is the diffusion model retrained from scratch in each epoch of the training loop? Also, is the training data in a given epoch just what is collected in the current epoch or the union of data collected in all epochs so far?\n- Have you investigated how the diffusion world model evolves as the agent used for collecting the training data improves?\n\nLIMITATIONS:\nThe authors have included a dedicated section that discusses the limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
This paper introduces a new world model for learning behaviors in imagination using reinforcement learning.
In particular, a diffusion model is used to generate the next frame, conditioned on previous frames and actions.
At each environment step, multiple denoising steps are performed to convert a noise image into the next frame.
By training an RL agent on synthesized trajectories, the method achieves state-of-the-art performance on the Atari 100k benchmark.

STRENGTHS:
- S1: The world model achieves strong performance on the Atari 100k benchmark, compared with other world models.
- S2: Implementing world models using diffusion models is a logical step, given the success of diffusion models in image generation.

WEAKNESSES:
- W1: The analysis of the world model is mainly qualitative and not quantitative. The model seems to be really good at generating long trajectories without compounding errors (e.g., Figure 3(b)), but it would be nice to have a more objective measurement of this. One simple idea would be to generate long trajectories and compare the generated frames to the nearest neighbors in the replay buffer. This could also be compared with IRIS and DreamerV3.
- W2: As the training is rather slow, it would be interesting to see a breakdown of the training times of the individual components. For instance, how much time is spent to generate frames compared with the reward/termination model (which is a CNN + LSTM)?

Typo in L240: The single denoising step is shown in the "last row" (instead of "first row").

QUESTIONS:
- Q1: The imagination horizon is set to the usual value of 15 steps. I am wondering whether longer horizons would lead to better scores, or whether this is not required for Atari?

LIMITATIONS:
The authors addressed all limitations adequately.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an approach for learning world models with diffusion based approaches, compared to the recently proposed ones using transformers or the ones dependent on discrete latent variables in general. The core idea is that given success of image generation using diffusion models, the visual details of game playing RL tasks can be improved if a diffusion based world model can be learnt; which can in turn lead to improved performance specifically in pixel based tasks like Atari.\n\nSTRENGTHS:\nThe idea of using generative models for learning the world model in RL is perhaps nothing new; past works have tried to do this with transformers or other discrete variable models; however success has not been achieved much. In contrast, while the idea of using diffusion for world models perhaps may sound oblivious given current context, this paper does a good job in demonstrating that this can work empirically in Atari. Experimental comparisons with past works such as IRIS, that are dependent on using transformers, shows that the learnt difufison based approach can achieve empirical gains.\n\nWEAKNESSES:\nOne primary bottleneck of this approach would be the dependency of the diffusion driven word model for long horizon tasks. Past works often model a per-next time step state prediction dynamics model, or uses models that can handle longer sequences. In contrast, the diffusion based approach may signifcaintly suffer from the long horizon. \n\nExperimental approach and algorithmic novelty is perhaps nothing significantly new; this paper basically does a simple plug in of diffusion model with careful fine-tuning in the context of model based RL. \n\nAs stated by the authors as well, the practical choice of the diffusion approach would matter a lot here; and carefully analysis is required to ensure that the difufuson model cxan learn a good world model. \n\nEmpirically, other than Atari benchmarks, can this work demonstrate that the idea can work in other domains and compare to more extensive experimental analysis with Dreamer based approaches? For example, Dreamer line of work often shows good results in locomotion or humanoid driven tasks - I think this approach can fail in those cases since it is harder to learn a world model that can generate visual details in those tasks? Does the authors have any comments around that?\n\nQUESTIONS:\nCan we see some more experimental results other than Atari to demonstrate that the idea can work? \nMore experimental comparisons with other model based approaches are required. MBRL is a huge literature in the field and lots of works have compared to different algorithms in different benchmarks. In contrast, this paper lacks empirical evidence other than Atari domains which seems concerning.\nHave the authors considered other non-diffusion based approaches or more recent works based on flow matching for example? Or can we see how the learnt world model differs based on the type of generative model we use? Example, if we can use different variants of diffusion or FM based models, how does the eprfomance differ? I’d like to understand in general, the significance of using generative models for learning world models. \nThere used to be prior works on imagination augmented rollouts, or works that would learn a multi-step forward dynamics model for example; other than IRIS that are dependent on the use of transformers, can we see some more experimental comparisons with those prior approaches\n\nLIMITATIONS:\nLack of enough experimental evidence other than Atari benchmarks", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}