Diffusion for World Modeling: Visual Details Matter in Atari

Authors: eloi  alonso
adam  jelley
vincent  micheli
anssi  kanervisto
amos  storkey
tim  pearce
françois  fleuret
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 16:32:42

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.

Main paper: 
 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 The environment time t flows along the horizontal axis, while the vertical axis represents the denoising time τ flowing backward from T to 0. Concretely, given (clean) past observations x 0 <t , actions a <t , and starting from an initial noisy sample x T t , we simulate a reverse noising process {x τ t } τ =0 τ =T by repeatedly calling D θ , and obtain the (clean) next observation x 0 t . The imagination procedure is autoregressive in that the predicted observation x 0 t and the action a t taken by the policy become part of the conditioning for the next time step. Animated visualizations of this procedure can be found at https://diamond-wm.github.io.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. 
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").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) ,EQUATION)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,EQUATION)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) ,EQUATION)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,EQUATIONwith reparameterization D θ (x τ , τ ) = S θ (x τ , τ )σ 2 (τ ) + x τ , where σ(τ ) is the variance of the τ -level perturbation kernel.The score-based diffusion model described in Section 2.2 provides an unconditional generative model of p data . To serve as a world model, we need a conditional generative model of the environment dynamics, p(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,EQUATION)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). 
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 θ ,EQUATIONwhere 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 θ ,EQUATION)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.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. 
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. 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. 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. 
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. The initial observation at t = 0 is common, and each row corresponds to a decreasing number of denoising steps n. We observe that DDPM-based generation suffers from compounding error, and that the smaller the number of denoising steps, the faster the error accumulates. In contrast, our EDM-based world model appears much more stable, even for n = 1.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 .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.Figure 4 : Single-step (top row) versus multi-step (bottom row) sampling in Boxing. Movements of the black player are unpredictable, so that single-step denoising interpolates between possible outcomes and results in blurry predictions. In contrast, multi-step sampling produces a crisp image by driving the generation towards a particular mode. Interestingly, the policy controls the white player, so his actions are known to the world model. This information removes any ambiguity, and so we observe that both single-step and multi-step sampling correctly predict the white player's position.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. In Asterix (top row), an enemy (orange) becomes a reward (red) in the second frame, before reverting to an enemy in the third, and again to a reward in the fourth. In Breakout (middle row), the bricks and score are inconsistent between frames. In Road Runner (bottom row), the rewards (small blue dots on the road) are inconsistently rendered between frames. None of these inconsistencies occur with DIAMOND. In Breakout, the score is even reliably updated by +7 when a red brick is broken3 .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 Offensive4 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.  
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. 
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. 
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. 
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. • 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. 
 
 Appendix: 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.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, p(x i |x i-1 ) = N (x i ; 1 -β i x i-1 , β i I).(8)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) . EQUATIONwhere σ data = 0.5.The noise parameter σ(τ ) is sampled to maximize the effectiveness of training as follows:EQUATIONwhere P mean = -0.4, P std = 1.2. Refer to Karras et al. (2022) for an in-depth analysis.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. 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,EQUATION)The value network V ϕ is trained to minimize L V (ϕ), the expected squared difference with λ-returns over imagined trajectories,EQUATION)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 RE-INFORCE 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,EQUATION)We provide performance profiles (Agarwal et al., 2021) for DIAMOND and baselines below. 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. 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. Figure 8 provides a quantitative measure of the compounding error demonstrated qualitatively in Figure 3 for DDPM and EDM based world models.Figure 8 : Average pixel drift between an imagined trajectory and the corresponding reference trajectory collected with an expert in Breakout. The trajectories are each 1000 timesteps, starting from the same frame and following the same sequence of actions. Each line displays the average and shaded standard deviation of 400 reference trajectories held out from training data. DDPM becomes more stable with increasing number of denoising steps, but is less stable than 1-step EDM, even with 10 denoising steps. The drift we observe for EDM corresponds to differences in the imagined trajectory rather than a pathological color shift as we see in Figure 3a .L Quantitative ablation on reducing the number of denoising stepsTable 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. 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) 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.We consider two potential diffusion model architectures, summarized in Figure 9 . 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.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.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.----CS:GO --------Driving ---- 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. 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. 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).https://en.wikipedia.org/wiki/Breakout_(video_game)#GameplayThis section was added after NeurIPS acceptance, following community interest in later CS:GO experiments.https://github.com/commaai/research 


Decision: This paper proposes a new RL agent trained with a diffusion-based world model. Authors argue that this approach preserves better visual details than discrete latent models, and the results show SOTA performance on the Atari 100k benchmark .
Authors show strong empirical results on the Atari 100k benchmark, outperforming other world model-based RL agents. This is a novel application of diffusion models to world modeling in RL. The paper is also well written and the code is open sourced.

I would have liked to see more quantitative experiments for the design choices. The method could also be applied to domains beyond Atari. However, reviewers generally agree on the paper's strengths. Three out of four reviewers strongly recommend acceptance, while one reviewer suggests a borderline reject. The authors' responses and additional analyses appear to have satisfied most reviewers. 

I recommend accepting this paper. The novel application of diffusion models to world modeling in RL represents an interesting contribution and has potential for several follow up studies.

Reviews: review: paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: One 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. 

Experimental 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. 

As 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. 

Empirically, 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?
questions: questions: Can we see some more experimental results other than Atari to demonstrate that the idea can work? 
More 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.
Have 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. 
There 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
limitations: limitations: Lack of enough experimental evidence other than Atari benchmarks
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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: limitations: The authors addressed all limitations adequately.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: One 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.
questions: questions: - The performance on a few tasks (BankHeist, Frostbite, UpNDown) are considerably worse than the baselines. Have you investigated why this is the case?
- 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?
- Have you investigated how the diffusion world model evolves as the agent used for collecting the training data improves?
limitations: limitations: The authors have included a dedicated section that discusses the limitations of the work.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: See weaknesses.
limitations: limitations: Discussion in paper.

Comments: title: Follow up on our rebuttal
Follow up on our rebuttal
Thank you for your response
Thank you for your response
Thank you for your response
Follow up on our rebuttal
Please could you respond to our rebuttal?
comment: Thank you for your response, and thanks again for your positive and constructive review!
Thank you for the response, I maintain that this is good work that deserves to be published and continue to recommend acceptance!
Dear Reviewer hqzV,

Thank you for your review and for the positive evaluation of our paper. We believe we've addressed your concerns and hope you appreciate our responses and additions.

As the discussion period is coming to a close, please let us know if there are any further points we can address to improve your rating.

Kind regards,

The Authors
Dear Reviewer 6rRD,

Thank you again for your thorough review of our paper. With less than 24 hours remaining in the discussion period, we would be grateful for your feedback on our response. 

Please let us know if there is anything else we can address, as we do feel that your current rating does not fairly reflect the value of our contribution.

Kind regards,

The Authors
Thank you for your response and updating your score, and thanks again for your constructive review.
Thank you for your response and updating your score, and thanks again for your constructive review.
Thank you for the thorough response! I have updated my score accordingly.
Thank you for the detailed responses to my questions and comments. 
I have updated my score to reflect this.
Dear Reviewers,

Thanks again for your time and effort reviewing our work. With less than 48 hours left in the discussion period, we wanted to kindly follow up on our rebuttal. We would greatly appreciate your feedback and are happy to clarify any points if needed.

Thank you for your time and consideration.

Best regards,

The Authors
Dear Reviewers, 

Thank you again for your time and effort reviewing our paper. We are now more than halfway through the discussion period. We would be grateful if you could reply to our responses soon so that we have time to address anything remaining.

Please could you let us know if we have addressed your concerns, or if there is anything else we can address for you to consider increasing your ratings of our paper?

Kind regards,

The Authors