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 2024 spotlight (2024)
Keywords: World models
diffusion models
reinforcement learning
generative models
Atari
Primary area: reinforcement_learning
TL;DR: We introduce DIAMOND, a diffusion world model, to train sample-efficient RL agents on Atari 100k.

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
Eloi Alonso"
Adam Jelley*
Vincent Micheli University of Geneva
Anssi Kanervisto Microsoft Research
University of Geneva
University of Edinburgh
Amos Storkey
Tim Pearce
François Fleuret University of Geneva
University of Edinburgh
Microsoft Research
Abstract
World models constitute 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 compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become dominant approach for image generation, challenging establissheed methods modeling discrete latents. Motivated by this paradigm shift, introduce DIAMOND (DIffusion As Model Of eNviromeent 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 mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within 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 htttps / //wi//ddddmdddd///////// amond wm. www.giithubbb..
Introduction
Generative models of environments, or "world models" (Ha and Schmidhuber) 2018). are becoming increasingly important as component for generalist agents 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., 20222, but is well-known be sample inefficient, which limits real-world applications. World models have shown promise for training reinforcement learning agents across diverse environments (Hainer et al., 2023 Schrittwieser et al.. 2020), with greatly improved sample-efficiency (Ye al. experience in the real world (Wu et al. 2003).
2021) which enable learning from
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 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 loss of generality and reconstruction quality. This may be problematic for more real-world scenarios where the information
To prevent confusion, this the final version Alooso et al. 2023 and not related to Ding et al. 2024 Equal contribution. 'Equal supervision. Contact: eloi al www.ouningeineooooomo ch and adam. aala.jeellll.l jelleyded
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Enviromment time (t)
Figure 1: Unrolling imagination of DIAMOND over time. The top row depicts a policy a taking a sequence of actions in the imagina- tion of our learned diffusion world model Dg. The environment time / flows along the hori- zontal axis, while the vertical axis represents the denoising time т flowing backward from T to O. Concretely, given (clean) past obser- vations x actions agt, and starting from an initial noisy sample x! we simulate reverse noising process (xī )=)) by repeatedly calling Dg, and obtain the (clean) next observation x!. The imagination procedure is autoregressive in that the predicted observation x and the ac- tion a taken by the policy become part of the conditioning for the next time step. Animated visualizations of this procedure can be found at https: hopp:://dimmodd -m gi thub.
x
x
0c+1
x
D
D
Denoising time
Da
Do
Do
X
x
Conditioning
X
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 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. (2220) have become dominant paradigm for high- nigh-resolution image generation RRombach et al. |0022] Podell
eet al. 2023). This class of methods, in which the model learns to reverse a noising process, cha lenges well established approaches modeling discrete tokens (Esser et al., 202T Ramesh al. [021 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 world model 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 eNviromment Dreams), a reinforcement learning agent trained in diffusion world model. Careful design choices are necessary to ensure our diffusion world model is efficient and stable over long-time horizons, and provide qualitative analysis to illustrate their importance. DIAMOND achieves 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 world model. Additionally, operating in image space has the benefit of enabling our diffusion world model to be 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 Sointer-Strike Global Offensive (CSGO) gameplay (Pearce and Zhu (022) to produce an interactive neural game engine for the popular in- in-game map, Dust II. We release our code, agents and playable world models at wtthhttt:/ccee http:::///dammoddeeee wm. www.iiiu...i......
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,7), where S is a set of states, A is a set of discrete
actions, and a set of image observations. The transition function T S x A X + [0, 1] describes the environment dynamics p(s st, az), and the reward function R : S x A x S - R maps transitions to scalar rewards. Agents cannot directly access states 8; and only see the environment through image observations 2, E O, emitted according to observation probabilities p(x; st).
2

---

The goal is to obtain policy that maps observations to actions in order maximize the expected discounted return E Ex r'r where 7 E [o, 1] is a discount factor. World models (Ha and Schmidhuber 2018) are generative models of environments, i.e. models of p( Tr 5cc Ge). These models can be used as simulated environments to train RL agents (Sutton 1991) in a sample-efficien 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").
described by the observation function O S 0 [0.
2.2 Score-based diffusion models
Diffusion models (Sohl-Dickstein et al. 2015) are a class of generative models inspired by non- cquilibrium thermodynamics that generate samples by reversing noising process.
We consider a diffusion process (x" hre(o,T] indexed by a continuous time variable T E 77, with corresponding marginals {p" [e( and boundary conditions p" = =paata and pT = parior, where yerior is tractable unstructured prior distribution, such
Gaussian. Note that use and superscript for the diffusion process time, in order to keep and subscript for the environment time.
This diffusion process can be described as the solution to standard stochastic differential equation (SDE) (Song et al. 2020).
dx f(x, + 9())dw,
(1)
where is the Wiener process (Brownian motion), f vector-valued function acting drift coefficient, and g a 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. Remark- ably, Anderson (1982) shows that the reverse process also a diffusion process, running backwards in time, and described by the following SDE,
dx = [f(x,T) )(T)² log (x)]drx - g(T)dw,
(2)
where is the reverse-time Wiener process, and V log (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 and g (in Section B.I). and estimate the unknown score functions Vx log (x)), associated with marginals {p" ] along the process. In practice, it is possible to use a single time-dependent score model (x 7) to estimate these score functions (Song et al., 2020)..
At any point time, estimating the score function 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", 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, analytically reach any time in the forward process in single step, by applying a Gaussian perturbation kernel prr to clean data samples (Song et al., 2020 Since the kernel is differentiable, score matching simplifies to denoising score matching objective (Vincent 2011).
(()) E [x Vx- Vx- log "x 1x)]11],
where the expectation is over diffusion time т, noised sample x* ~ px x9), obtained by applying the r-level perturbation kernel to a clean sample x" (dota(x)), Importantly, as the kernel Dr is known Gaussian distribution, this objective becomes simple L2 reconstruction loss,
C(0) =E
E [|| .T) -x"11]]
(4)
with reparameterizattio Do(x", = , (7) +x", where (T) is the variance of the --level perturbation kernel.
3

---

2.3 Diffusion for world modeling
The score-based diffusion model described Section provides an unconditional generative model of Pdata- To serve as world model, we need conditional generative model of the environment dynamics, p(x(+1 X<t: age), where we consider the general case of POMDP, in which the Markovian state &, 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 ass follows,
(0) E [I T, X age) X++1
During training, we sample a trajectory segment x agt, xl41 from the agent's replay dataset, and we obtain the noised next observation x[+1 Td (xī_ x(+1) by applying the r-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 While we can in principle resort to any ODE or SDE solver, there is 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
now introduce DIAMOND as 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
Building the background provided in Section
particular choice of diffusion paradigm. While DDPM (Ho et al. 2020) is an example of one such choice (as descrihed 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 Section 5.1 In what follows, we describe how we adapt EDM to build our diffusion-based world model.
We consider the perturbation kernel ""T(XT
x(+1) N [x x o² (T).), where (T)) is a real-valued function of diffusion time called the noise schedule. This corresponds to setting the drift and diffusion coefficients to f(x, r) 0 (affine) and g(T) 22
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ø.
Do (x v.) Cakip X
C Fø x(41. v(),
(6)
where for brevity we define yf = (ci x xb,, a(g) to include all conditioning variables.
The preconditioners < and ciu are selected keep the network's 's input and output at unit variance for any noise level (T), Groise is an empirical transformation of the noise level, and Caip is given in terms
of ar (T) and the standard deviation of the data distribution ddatas Cakip = oistal (a + o² (r)).
These preconditioners are fully described in Appendix C
Combining Equations 5 and 6 provides insight into the training objective of Fg.
((0)= B[1 Fo Network prodiction
(x) 1 Caip x [12]
Court
Network training target
The network training target adaptively mixes signal and noise depending on the degradation level (T)). When (TT) >> Odata, we have Skip - O, and the training target for Fo is dominated by the clean signal x Conversely, when the noise level is low, o(T) + O, we have Glip + 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 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 (T) from an empirically chosen log-normal distribution order to concentrate the training around medium-nois regions, described in Appendix C

---

We use a standard U-Net 2D for the vector field Fg (Ronneberger et al. 2015). and we keep buffer of L past observations and actions that we use 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 reward and termination model, required for training an RL agent in imagination. Since estimating the reward and termination are scalar prediction problems, we use separate model Ry consisting of standard CNN (LeCun et al., 1989 He et al. 2013) and LSTM (Hochreiter and Schmidhuber 1997, Gers et al. 2000) layers to handle partial observability. The RL agent involves actor-critic network parameterized by a shared CNN-LSTM with policy and value heads. The policy Tg is trained with REINFORCE with value baseline, and we use Bellman error with X-returns to train the value network Vz, 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 detailed in Algorithm[] I and is similar to Kaiser et al. (2019); Hafner et (2020); (00h)) Michhll et al. (2023). We provide architecture details, hyperparameters, and RL objectives Appendices
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 hours of human gameplay, to learn to play the game before evaluation. As 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 I including STORM (Zhang et al., 2023), DreamerV3 (Hafner et al. (002)), IRIS (Micheli et al. 2023 TWM (Robine et 2023), and SimPle (Kaiser et al. 2019). A broader comparison to model- free
3, and search-based methods, including BBF (Schwarzer et al.] 2023) and EfficientZero (Ye et al. 2021), the current best performing methods this benchmark. provided in Appendix BBF and EfficientZero use techniques that are orthogonal and not directly comparable to our approach, such using periodic network resets in combination with hyperparameter scheduling for BBF, and computationally expensive lookahead Monte-Carde tree search for EfficientZero. Combining these additional components with our world model would be interesting direction for future work.
4.2 Results on the Atari 100k benchmark
Insoquariia
STORM
Table provides scores for all games, and the mean and interquartile mean (IQM) of human- normalized scores (HNS) (Wang et al.] 2016).
TWM
Following the recommendations of |Agawaal et al. (222l) on the limitations of point estimates, we provide stratified bootstrap confidence inter- vals for the mean and IQM in Figure 2 as well as performance profiles and additional metrics
Sautt
Muman Mormal lord Score
Figure 2: Mean and interquartile mean human nor- malized scores. DIAMOND, in blue, obtains a mean HNS of 1.46 and an IQM of 0.64.
in Appendix H
5

---

Table I: Returns on the 26 games of the Atari 100k benchmark after 2 hours of real-time experience, and human-normalized aggregate metrics. Bold numbers indicate the best performing methods.
DIAMOND notably outperforms other world model baselines in terms of mean score over 5 seeds.
Game
Random
Human
SimPLe
TWM
IRIS
DreamerV3
STORM
DIAMOND (ours)
Alien
227.8
7127.7
616.9
674.6
420.0
959.0
983.6
744.1
Amidar
5.8
1719.5
74.3
121.8
143.0
139.0
204.8
225.8
Assault
222.4
742.0
527.2
682.6
1524.4
706.0
801.0
1526.4
Asterix
210.0
8503.3
1128.3
1116.6
853.6
932.0
1028.0
3698,5
BankHeist
14.2
753.1
34.2
466.7
53.1
649.0
641.2
19.7
BattleZone
2360.0
37187.5
4031.2
5068.0
13074.0
12250.0
13540.0
4702.0
Boxing
0.1
12.1
7.8
77.5
70.1
78.0
79.7
86.9
Breakout
1.7
30.5
16.4
83.7
31,0
15.9
132.5
ChopperCommann
811.0
7387.8
979.4
1697.4
1565.0
420,0
1888.0
1369,8
Crazy Climber
10780.5
35829.4
62583.6
71820.4
59324.2
97190.0
66776.0
99167.8
DemonAttack
152.1
1971.
208.1
350.2
2034.4
303.0
164.6
288.1
Freeway
0.0
29.6
16.7
31.1
33.5
33,3
Frostbite
65.2
4334.7
236.9
1475.6
259.1
909.0
1316.0
274.1
Gopher
257.6
2412.5
596.8
1674.8
2236.1
3730,0
8239.6
5897.9
Hero
1027.0
30826.4
2656.6
7254.0
7037.4
11161.0
11044.3
5621.8
Jameshond
29.0
302.8
100.5
362.4
462.7
445.0
509.0
427.4
Kangaroo
52.0
3035.0
5122
1240.0
838.2
4098.0
4208.0
S382,2
Krull
1598.0
2665.5
2204.8
6349.2
6616.4
7782.0
8412.6
8610.1
KungteMaster
258.5
22736.3
14862.5
24554.6
21759.8
21420.0
26182.0
18713.6
MsPacman
307.3
6951.6 66
1480.0
1588.4
999.1
1327.0
2673.5
1958.2
Pong
-20.7
14.6
12.8
18.8
14.6
18.0
11.3
20.4
PrivateFye
24.9
69571.3
35.0
86.6
100.0
882.0
7781.9
114.3
Qbert
163.9
13455.0
1288.8
3330.8
745.7
3405.0
4522.5
4499.3
RoudRunner
11.5
7845.0
5640.6
9109.0
9614.6
15565.0
17564.0
20673.2
Seaquest
68.4
42054.7
683.3
774.4
661.3
618.0
525.2
551.2
UpNDown
533.4
11693.2
3350.3
15981.7
3546.2
9234.0
7985.0
3856.3
#Superhuman (1)
N/A
10
10
11
Mcan (?)
0,000
1,000
0.332
0.956
1.046
1.097
1.266
1.459
IQM (1)
0.000
1,000
0.130
0,459
0.501
0.497
0.636
0.641
Our results demonstrate that DIAMOND performs strongly across the benchmark, outperforming human players on 11 games, and achieving superhuman mean HNS of 1.46, a new best among agents trained entirely within 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
Analysis
5.1
Choice of diffusion framework
As explained in Section 2 could in principle use any diffusion model variant in our world model. While DIAMOND utilizes EDM (Karras et al.] 2022) as described Section DDPM (Ho et al. (020) would also be a natural candidate, having been used in many image generation applications (Rombach
ct al. 2022 Nichol and Dhariwal 2021 We justify this design decision in this section.
To provide 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 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 IRIS which requires 16 NFE for each timestep),
we need at most tens of denoising steps, and preferably fewer. Unfortunately, the number of denoising steps is set too low, the visual quality will degrade, leading 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 Ba) 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 single denoising step. A quantitative analysis of this compounding error is provided in Appendix K

---

1+0
E+5
t+10
t-50
1-100
t-500
E-1 0800
E+0
1+5
C-10
1-50
1-100
1-500
[-1 800
5+5
5
(a) DDPM-base world model trajectories.
(b) EDM-based world model trajectories.
Figure 3: Imagined trajectories with diffusion world models based on DDPM (left) and EDM (right). The initial observation at t = 0 is common, and each row corresponds to 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
This surprising result is 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 (Onoise > data = prediction network of DDPM. This gives poor estimate of the sampling procedure, which degrades the generation quality and leads to compounding error.
Eo(x++1 yf + x 1), where &o is the noise function at the beginning of the
In contrast, the adaptive mixing of signal and noise employed by EDM, described in Section 3.1 means that the model trained to predict the clean image when the noise dominant (Gnoise
ddta = Fx MFE + x This gives 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 the number of denoising steps
While found that our EDM- based world model was very stable with just single denoising step, 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 quantitative analysis in Appendix L
As discussed in Section 2.2 our score model is equivalent to denoising autoencoder (Vincent) et al. 2008) trained with an Lz reconstruction loss. The optimal single-step prediction is thus the expectation over possible reconstructions for given noisy input, which can 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 As result, therefore set
in the game Boxing in Figure
3 in all of our experiments.
t=0
t=1
t=2
t=3
t=4
t=5
[
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-ster sampling produces crisp image by driving the generation towards 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.
7

---

5.3 Qualitative visual comparison with IRIS
We now compare to IRIS (iicheli al. 2023), well-established world model that uses discrete autoencoder (Van Den Oord et al. (017) to convert images to discrete tokens, and composes these tokens over time with an autoregressive transformer (Radford 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 below.
(a) IRIS
(b) DIAMOND
Figure 5: Consecutive frames imagined with IRIS (left) and DIAMOND (right). The white boxes highlight inconsistencies between frames, which see only arise in trajectories generated with IRIS. In Asterix (top row), an enemy (orange) becomes a reward (red) in the second frame, before reverting to enemy in the third, and again to 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 broker
We see in Figure 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 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 make it more challenging learn optimal policy.
enemies, these small visual discrepancies
These improvements in the consistency of visual details are generally reflected by greater agent performance these games, as shown in Table 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 the same resolution (64 x 64), and DIAMOND requires only 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 train than IRIS, as provided in Appendix
Scaling the diffusion world model to Counter-Strike. Global Offensive
To investigate the ability of DIAMOND's ' diffusion world model to learn to model more complex 3D environments, train the world model in isolation 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 16Hz from the map Dust I/ by Pearce and Zhu (002)) We randomly hold out 0.5M frames (corresponding to 500 episodes, or hours) for testing, and use the remaining 5M frames (87 hours) for training. There is reinforcement learning agent online data collection involved in these experiments.
hattpa://en wikipedia. .oo...rrg/......ooooooeo... org/viki/Breakout, (video_ gammee #Gameplay
This section added after NeurlPS acceptance, following community interest later CS: experiments.

---

To reduce the computational cost, we reduce the resolution from (280 x 150) to (56 x 30) for world modeling. We then introduce second, smaller diffusion model as an upsampler to improve the generated images at the original resolution (Saharia et al. (2022)). 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 5 51M for the upsampler). The combined model was trained for 12 days 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 denoising steps). This enables 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 Sigure6obelooe
shoct reload
mUyS le Flach
reco
navigat tng turnel
arable
J'anp onto box
-ight
throuch
Figure 6: Images captured from people playing with keyboard and mouse inside DIAMOND'S diffusion world model. This model was trained on hours of static Counter-Strike. Global Offensive (CS:GO) gameplay (Pearce and Zhu 2022) to produce interactive neural game engine for the popular in-game map, Dust II. Best viewed as videos at wtthttts:///// /ttp::://damondd diamonddmdddd/////////.do//ddddd wm github.
We find the model is able to generate stable trajectories over hundreds of timesteps, although is more likely to drift out-of-disttibution in less frequently visited areas of the map. Due the limited memory of the model, approaching walls or losing visibility may cause the model to forget the current state and instead generate new weapon or area map. Interestingly, we find the model wrongly enables successive jumps by generalizing the effect of a jump 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 address many of these limitations, with the exception of the memory 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 RI. from the latent space of recurrent state space model (RSSM). DreamerV2 (Hafner et al. 202I) demonstrated that using discrete latents could help to reduce compounding error, and DreamerV3 (Hafner et al. 2023) able to achieve human-level performance on a wide range of domains with fixed hyperparameters. TWM (Robine et al. 2023 adapts DreamerV2' RSSM to
transformer architecture, while STORM (Zhang et al. 2023
adapts DreamerV3 V3 in a similar way but with a different tokenization approach. Alternatively, IRIS (Micheli al.) 2023) builds language of image tokens with 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 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 RRamesh
al. ) convert images to discrete tokens with discrete autoencoders (Van Den Oord 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 (aan et al.) (02T) provides a minimal video generation architecture by combining a discrete autoen- coder with an autoregressive transformer. Godiva (Wu al.) (021) enables text conditioning with promising generalization. Phenaki (Villegas et al. (2023) allows arbitrary length video generation with sequential prompt conditioning. TECO (Yan et al., (023) 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 31) U-Nets to provide high quality but short-duration video (Singer et al.) [023 Bar-Tall 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 al., 2022; Pearce et al.) (023), as planners (Janner
et al. 2022 Liang et al. 2023). 3. as reward models (Nuti et al. 2023), and trajectory modeling for
data augmentation in offliee 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 al. 2014) while BBafford and Lucas (2020) use a Neural GPU (Kaiser and Sutskever 2015). Concurrent work includes Genie (Bruce et al. (024). 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 larger scale.
8 Limitations
We identify three main limitations of our work for future research. First, our main evaluation is focused discrete control environments, and applying DIAMOND 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 (202)), 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 diffusion model is not trivial (Luo et al. 2023 Xu al. 2023) and would make our world model unnecessarily complex.
Conclusion and Broader Impact
We have introduced DIAMOND, reinforcement learning agent trained in 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 low number of denoising steps. DIAMOND achieves a mean human normalized score of 1.46 on the well-established Atari 100k benchmark; 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 's diffusion world model can successfully model 3D environments and serve as real-time neural game engine by training on static Counter-Strike: Global Offensive gameplay.
World models constitute 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.
10

---

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 Valcareel 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). conditional generative modeling all you need for decision decision-making! making? International Conference on Learning Representations.
Alonso, E., Jelley, A., Kanervisto, A., and Pearce, (2023). Diffusion world models. OpenReview.
Anderson, B. D. (1982). Reverse-time diffusion equation models. Stochastic Processes and their Applications, 133)3313-326.
Ascher, U. M. and Petzold, L. R. (1998). Computer methods for ordinary differential equations and aefferentiul-algebbee equations. Society for Industrial and Applied Mathematics.
Bamford, 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- J., Shi, Y., Hughes, E.,, Lai, M., Mavalankar, A., Steigerwald, R., Apps, C., et al. (2024). Genie: Generative interactive environments. Forty-first International Conference 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 1111-113221
Ç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-A -Asisted Intervention-MICCAl 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, (2018). Bert: Pre-training of deep bidirectional transformers for language understanding, arXiv preprint 171.0:1810.048855

---

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- 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//VE 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)24411-241.
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, (2018). Recurrent world models facilitate policy evolution. Advances in Neural Information Processing Systems, 31:2451-2466
Hafner, D., Lillicrap, T., Ba, 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, Ba, J., and Lillicrap. T. (2023). Mastering diverse domains through world models. arXiv preprint arXiv:2301.04144.
He, K., Zhang, X., Ren, S., and Sun, J. (2015). Deep residual learning for image recognition. arXiv preprint 1rXi::15122033555
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. (2017). Gans trained by a two time-scale update rule converge to 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, 68511.
Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, J. (2022). Video diffusion models. URL https://arxiv. org/abs/2204.03458.
Hochreiter, S. and Schmidhuber, J. (1997). Long short-term memory. Neural Computation, 987: 1780.
Hu, A., Russell, L.,, Yeo, H., Murez, Z.,, Fedoseev, G., Kendall, A., Shotton, J., and Corrado, (2023). Gaia-l: 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 240.66660
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 1031v:9903.00374.
Kaiser, L. and Sutskever, I. (2015). Neural gpus learn algorithms. arXiv preprint arXiv:1511.0822
12

---

Kapturowski, S., Ostrovski, G., Quan, J., Munos, R., and Dabney, W. (2018). Recurrent experience replay in distributed reinforcement learning. International Conference on Learning Representa- lions.
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). 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, 144):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 Machine Learning.
Lu, C., Ball, P. J., and Parker-Holder, (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., 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):599-533.
Nichol, A. Q. and Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. Interna- tional Conference on Machine Learning.
Nuti, F., Franzmeyer, T., and Henriques, F. (2023). Extracting reward functions from diffusion models. arXiv preprint arXhv:2306.0/804
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. 7., Momennejad, I., Hofmann, K., and Devlin, (2023). Imitating human behaviour with diffusion models. The Eleventh International Conference Learning Representations.
Pearce, T. and Zhu, (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 IEFE/CVF International Conference on Computer Vision (ICCV), pages 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 bigh-resolution image synthesis. arXiv preprint 207.02230..
13

---

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 22t::22044006225,5 1(2:3.
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C.,, Radford, A., Chen, M., and Sutskever, I. (2021). seoo text-to-image generation. International Conference on Machine Learning.
Robine, J., Höftmann, M., Uelwer, T., and Harmeling, S. (2023). Transformer-basee 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 InterventinnMMICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer.
Saharia, C., Chan, W. Saxena, S., L., Whang, J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., al. (2022a). Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479- 36494.
Saharia, C., Ho, J., Chan, W., Salimans, T., Fleet, D. J., and Norouzi, M. (2022b). Image super- resolution via iterative refinement. IEEE transactions on pattern analysis and machine intelligence.
Santana, E. and Hotz, G. (2016). Learning a driving simulator. arXiv preprint
0/230.
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L.., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel. T., al. (2020). Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839)6604-00.
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., Lanetot, M., 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 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 IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 3626-3636.
fohh-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. (2015). Deep unsupervised learning using noncquilibrium thermodynamics. International Conference on Machine Learning.
Song. Y.. Dock-eickst,in, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole. B. (2020). Score- based 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, 160)-160-
Sutton, R. S. and Barto, A. G. (2018). Reinforcement learning: An introduction. MIT press.
14

---

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:/812.0171
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, L., and Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems, 30.
Villegas, R., Babaeizadeh, M., Kindermans, P.-.,, 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, 2367):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 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. 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 104x:2210441155
Ye, W., Liu, S., Kurutach, T., Abbeel, 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. 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 trans- former based world models for reinforcement learning. In Thirty-eventth Conference on Neural Information Processing Systems.
Zheng, H., Fu, J., Zeng, Y., Luo, J., and Zha, 7.-J. (2020). Learning semantic-aware normalization for generative ad versariial networks. In Advances in Neural Information Processing Systems.
15

---

A Sampling observations in DIAMOND
We describe here how sample observation x from our diffusion world model. We initialize the procedure with a noisy observation x ~ pprtoor, and iteratively solve the reverse SDE in Equation 2 from T observations x2, and actions act. This procedure is illustrated in Figure
using the learned score model Sq (x] T; x= act)) conditioned on past
In fact, there are many possible sampling methods for a given learned score model Sn (Karras et al. 2022). Notably, [Song et al. (2020) introduce corresponding "probability flow" ordinary differential equation (ODE), with marginals equivalent to the stochastic process described in Section 2.2 that case, the solving procedure is deterministic, and the only randomness comes from sampling the initial condition. In practice, this means that for given score model, we can resort to any ODE or SDE solver, from simple first order methods like Euler (deterministic) and Suler-Maruyana (stochastic) schemes, to higher-order methods like Heun's '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 sample. This local error generally scales superlinearly with respect to the step size (for instance Euler's method is O(h²) for step size ht), so increasing the number denoising steps improves the visual quality of the generated next frame. Therefore, there is 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 dis- crete 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 discrete noise schedule ,,,....,,....................... ....... < 1, and variance-preservingg Gaussian transition kernel,
pxx11x =N(xx 3
3,''=1,600 ,,I)).
(8)
In the continuous time limit N oo, the Markov chain becomes a diffusion process, and the discrete noise schedule becomes time-dependent function ((T). This diffusion process can be described by
an SDE with drift coefficient f(x, 7) - xx and diffusion coefficient g(T) v (B( (Song
ct al. 2020).
C EDM network preconditioners and training
Karras et al. (022) use the following preconditioners for normalization and rescaling purposes (as mentioned in Section 3.1 1 to improve network training:
on
(9)
Vo((T² o Jaata o(T)Odata
Cau Va(T)² + o o
(10)
Croise log(o(T))
(11)
0
C,iip
(12)
σ data + oz
where ddata 0.5.
The noise parameter a (T) is sampled to maximize the effectiveness of training as follows:
log(a(T)) (P p22),
(13)
where Pmean -0.4, Pard 1.2. Refer to Karras et al. 2022) for an in-depth analysis.
16

---

D Model architectures
The diffusion model Dg is a standard U-Net 2D (Ronneberger al. 2013). conditioned on the last 4 frames and actions, 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 Ry layers 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.; 2013) followed by an LSTTM cell (Mnih 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 initialize the hidden and cell states of the LSTM.
The weights of the policy Ts and value network Vo are shared except for the last layer. In the following, refer to (x, V) the "actor-critic" network,
though V is technically state- value network, not a critic. This network takes as input a frame, and forwards it through convolutional trunk followed by 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, SiLU activation (Elfwing et al. 2018)) and 3x3 convolution with stride 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 beloo for hyperparameter values, and to Algorithm for a detailed summary of the training procedure.
Table 2: Architecture details for DIAMOND.
Hyperparameter
Value
Diffusion Model (D₂)
Observation conditioning mechanism
Frame stacking
Action conditioning mechanism
Adaptive Group Normalization Adaptive Group Normalization
Diffusion time conditioning mechanism
Residual blocks layers
[2, 2, 2, [64, 64, 64, 64] 256
Residual blocks channels
Residual blocks conditioning dimension
Reward/Termination Model (R.))
Action conditioning mechanisms
Adaptive Group Normalization [2, 2, 2, 2] [32, 32, 32, 32] 128
Residual blocks layers
Residual blocks channels
Residual blocks conditioning dimension
LSTM dimension
512
Actor-Critte Model (To and V()
Residual blocks layers
[1, 1, 1, 1] [32, 32, 64, 64] 512
Residual blocks channels
LSTM dimension
17

---

E
Training hyperparameters
Table 3: Hyperparameters for DIAMOND.
Hyperparameter
Value
Training loop
Number of epochs
1000
Training steps per epoch Batch size
400
32
Environment steps per epoch Epsilon (greedy) for collection
100
0.01
RL hyperparameters
Imagination horizon (If)
15
Discount factor (r)
0.985
Entropy weight (1)
0.001
X-returns coefficient (A)
0.95
Sequence construction during training For Da, number of conditioning observations and actions (T.)
4
For Rg. burn-in length (BR), set to L in practice
4
For Ry. training sequence length (BR + H)
19
For Tę and Vg. burn-in length (Bx, v), set to L in practice
4
Optimization
Optimizer
AdamW
Learning rate
le-4
Epsilon
1e-8
Weight decay (Da)
1e-2
Weight decay (Rg.)
1e-2
Weight decay (Tg and Vi)
0
Diffusion Sampling
Method
Euler 3
Number of steps
Environment
Image observation dimensions
x4xx6x3
Action space
Discrete (up 18 actions)
Frameskip
4
Max noop
30
Termination on life loss
True
Reward clipping
(-1, 1}
18

---

Reinforcement learning objectives
In what follows, we note X,, r; and de the observations, rewards, and boolean episode terminations predicted by our world model. We note H the imagination horizon, Vo the value network, T the policy network, and ae the actions taken by the policy within the world model.
We use X-returns to balance bias and variance as the regression target for the value network. Given imagined trajectory of length H, we can define the A-return recursively as follows,
re 7 de) di) V)Vo + MA++1
t <
A,
(14)
(x(xH)
t= H.
The value network V& trained to minimize Lv (6), the expected squared difference with A-returns over imagined trajectories.
Lv( = ET Σ
- -g(xx))s(AA)))
(15)
where sg(-) denotes the gradient stopping operation, meaning that the target is a constant in the gradient-based optimization, classically established in the literature (Mnih et al. 2015 Hatner et al. 2021 Micheli et al. 2023).
As we can generate large amounts of on-policy trajectories in imagination, we use simple RE- INFORCE objective to train the policy, with the value (x(x,) as a baseline reduce variance of the gradients (Sutton and Barto, 2018). The policy is trained to minimize the following objec- tive, combining REINFORCE and weighted entropy maximization objective to maintain sufficient exploration,
Σ log (T& (a, xxx)) sg (A, Vo (xx)) + 17 (mo (a, (xsi)]
L (Ø) -
(16)
19

---

G DIAMOND algorithm
We summarize the overall training procedure of DIAMOND in AlgorithmmTbelow We denote as 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.
Algorithm 1: DIAMOND
Procedure training.loop(/ for epochs do
gellect_experiente (steps_colled)
for steps_diflusion, ddiflssionmoodel mooel
update_dis ltl.:ffffsoooooeeeeeceo ()
for steps steps_reward @wward_end_at model do
update_rewar leee_reewrdd_edddeeeeeeeeeeeeeed serrdddendddddeee ()
for steps_ace acttraccttiittt _rrttic do
update_ ceattte_atorcrrttiio
Procedure collect_experienet (n):
x - env. reset() do Sample aa(at x
for 1
x++1 rt, dc +- env step (0:)
D DU ax )
d = 1 then
x env reset)
Procedure update_difitteeeeeeeeeee www.usiionmooeeeeeeee
Sample sequence (xl (xi-L+1: 0-111 xi: ot, x(+1)
Sample log(o) (N Patd) //
sigma distribution from EDM 11 default identity schedule from EDM
Define T 0
Sample X,41 ~
// Add independent Gaussian noise
Compute 8 Do(x T, 1, 14--44:00-++11 xf...)))....))
Compute reconstruction loss C(0) [80 -x--1l1ll x2+1l?
Update D3
Procedure update_rres penttte_rwardd_er end fondd_mode[[ ():
Sample indexes I = (t....,t t+ L H - 1} // burn-in imagination horizon Sample sequence (x. at Ti, di)iez D 1/ LSTM hidden and cell states
Initialize = c= 0
for i I do
L Compute A, d, h, c= (x, a, h,c)
Compute L(v) = EiEz CE(A, sign(r:) (r:)) + CE(d,,) d,) // CE: cross-entropy loss
Update Rs
Procedure update_actor, actor_ critic():
Sample initial buffer (x) (xt-L+1: at L+...... x2) D
Burn-in buffer with R.. no and V to initialize LSTM states
for i to t
do
Sample a; (T xx)
Sample reward and termination d; with R.
Sample next observation xf.1 by simulating reverse diffusion process with D
Compute (x;) for i = t.....tt tt H
Compute RL losses Lv (Ø) and (Ø)
Update T and Vi
20

---

H Additional performance comparisons
We provide performance profiles Aggarall al. 202
for DIAMOND and baselines below.
SimPLe
IRIS
STORM
1.00
TWM
Dreamerv3
DIAMOND (ours)
Fraction of runs with score
0.75
0.50
0.25
0.00
0
Human Normalized Score (T)
Figure 7: Performance profiles, i.e. fraction of runs above given human normalized score.
As additional angles of comparison, we also provide parameter counts and approximate training times for TRIS, DreamerV3 and DIA MOAD in Table beblew 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 4: Number of parameters, training time, and mean human-normalln score (HNS).
IRIS
DreamerV3
DIAMOND (ours)
#parameters (L)
30M
18M
13M
Training days (1)
4.1
<1
2.9
Mcan HNS (1)
1.046
1.097
1.459
A full training time profile for DIAMOND is provided in Appendix !
21

---

I Training time profile
Table provides full training time profile for DIAMOND.
Table 5: Detailed breakdown of training time. Profiling performed using a Nvidia RTX 4090 with the default hyperparameters specified in Appendices and T profiling measures representative, since exact durations will depend on the machine, the environment, and the training stage.
Single update Total
Time (ms)
Detail (ms)
543
88 + 115 + 340
Diffusion model update
88
Reward/Termination model update
115
Actor-Critic model update
340
x 20.4 + 34
Imagination step (x 15)
20.4
12.7 +.0 + 0.7
Next observation prediction
12.7
3x22
Denoising step (x 3)
4.2
Reward/Terme Dewaad//ermination prediction
7.0
Action prediction
0.7
Loss computation and backward
34
Epoch
Time (s)
Detail (s)
Total
217
35 + 46 + 136
Diffusion model
35
400 x x 10
Reward/Termination model
46
400 115 10-3
Actor-Critic model
136
400 x 340 x 10-3
Run
Time (days)
Detail (days)
Total
2.9
2.5 0.4
Training time
2.5
1000 x 217/(24 x 3600)
Other (collection, evaluation, checkpointing)
0.4
22

---

Broader comparison to model-free and search-based methods
Table provides scores for model-free and search-based methods, including the current best per- forming 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., 202T) and BBF (Schwarzer al., (023) respectively) can still provide greater performance overall, DIAMOND promisingly still outperforms these methods on some games.
Table 6: Raw scores and human-normalicn normalized metrics for scarch-based and model-free methods.
Seurch-based
Model-free
Game
Human
MuZero
EfficientZero
CURL
SR-SPR
DIAMOND (ours)
Alien
7127.7
530.0
808.5
7110
841.9
1107.8
1173.2
744.1
Amidar
1719.5
113.7
179.7
203.4
244.6
225.8
Assault
742.0
500.1
1263.1
500.9
565.6
1088.9
2098.5
1526.4
Asterix
8503.3
1734.0
25557.8
567.2
962.5
903.1
3946.1
3698.5
BankHeist
753.1
192.5
351.0
65.3
345.4
531.7
732.9
19.7
BattleZone
37187.5
7687.5
13871.2
8997.8
14834.1
17671.0
24459.8
4702.0
Boxing
12.1
52.7
45.8
86.9
Breakout
30.5
48.0
414.1
2.6
19.6
25.5
370.6
132.5
ChopperCommann
7387.8
1350.0
1117.3
783.5
946.3
2362.
7549.3
1369.8
CrazyClimber
35829.4
56937.0
83940.2
9154.4
36700.5
45544.1
58431.8
99167.8
DemonAttack
1971.0
3527.0
13003.9
646.5
517.6
2814.4
13341.4
288.1
Freeway
29.6
21.8
21.8
28.3
19.3
25.4
25.5
33.3
Frostbite
4334.7
255.0
296.3
1226.5
1170.7
2584.8
2384.8
274.1
Gopher
2412.5
1256.0
3260.3
400.9
660.6
712.4
1331.2
5897.9
Hero
30826.4
3095.0
9315.9
4987.7
5858.6
8524.0
7818.6
5621.8
Jamesbond
302.8
87.5
$17.0
3310
366.5
389.1
1129.6
427.4
Kangaroo
3035.0
62.5
724.1
740.2
3617.4
3631.7
6614.7
5382.2
Krull
2665.5
4890.8
5663.3
3049.2
3681.6
5911.8
8223.4
8610.1
KungFuMaster
22736.3
18813.0
30944.8
8155.6
14783.2
18649.4
18991.7
18713.6
MsPacman
6951.6
1265.6
1281.2
1064.0
1574.1
1958.2
Pong
14.6
-6.7
20.1
-18.5
-5.4
2.9
16.7
20.4
PrivateFye
69571.3
56.3
96.7
81.9
86.0
97.9
40,5
114.3
Qbert
13455.0
3952.0
13781.9
727.0
866.3
4044.1
4447.1
4499.3
RoadRunner
7845.0
2500.0
17751.3
5006.1
12213.1
13463.4
33426.8
20673.2
Scaquest
42054.7
208.0
1100.2
315.2
558.1
819.0
1232.5
551.2
UpNDown
11693.2
2896.9
17264.2
2646.4
10859,2
112450.3
12101.7
3856.3
#Superhuman
N/A
5
14
11
Mean (1)
1.000
0.562
1.943
0.261
0.616
1.271
2.247
1.459
IQM (T)
1.000
0,288
1.047
0.113
0.337
0,700
1.139
0.641
23

---

K Quantitative analysis of autoregressive model drift
Figure g provides quantitative measure of the compounding error demonstrated qualitatively in Figure
for DDPM and EDM based world models.
Drift from reference trajectory for WM with denoising steps
DDPM n=1
1.4
DDPM n=3
Average pixel distance
DDPM n=5
1.2
DDPH n=10
n=1
1.0
EDM
n=3
EDM
n=5
0.8
EDM
n=10
0.6
0.4
0.2
0.0
10
50
100
500
1000
Timestep
Figure 8: Average pixel drift between 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 in Figure 3a.
trajectory rather than a pathological color shift
L Quantitative ablation on reducing the number of denoising steps
Table | provides quantitative ablation of the effect of reducing the number of denoising steps used for our EDM diffusion world model from 3 (used for Table to 1, for DIAMOND's 10 highest performing games. Note that the 1-step results correspond to 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 g This additional evidence supports our qualitative analysis in Section 5.2
Table 7: Quantitative ablation reducing the number of denoising steps from 3 (default) to 1.
Game
Random
Human DIAMOND (n 3)
DIAMOND (n 1)
Amidar
1719.5
225.8
191.8
Assault
222.4
742.0
1526.4
782.5
Asterix
210.0
8503.3
3698.5
6687.0
Boxing
0.1
12.1
86.9
41.9
Breakout
30.5
132.5
50.8
CrazyClimher
10780.5
35829.4
99167.8
87233.0
Kangaroo
52.0
3035.0
5382,2
1710.0
Krull
1598.0
2665.5
8610.1
9105.1
Pong
-20.7
14.6
20.4
20.9
RoadRunner
11.5
7845.0
20673.2
5084.0
Mean HNS (1)
0.000
1.000
3.052
1.962

---

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 world model on the well-established Atari 100k benchmark (Kaiser et al.) 2019). and demonstrated DIAMOND's 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 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 I/ 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 64x64 pixels, and use augmentation.
Motorway driving. We use the dataset from Santana and Hotz (2016 which contains camera and metadata captured from human drivers US motorways. We select only trajectories captured in daylight, and exclude the first and last minutes of each trajectory (typically traveling to/from 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 x4,, and for actions the (normalized) steering angle and acceleration. During training, 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 DIA DOAMON''s diffusion model on these static datasets, and that we do not perform reinforcement learning, since there no standard reinforcement learning protocol for these environments.
M.2 Diffusion Model Architectures
consider two potential diffusion model architectures, summarized in Figure
Diffusion video generation
Diflusion world model architectures
2) Cross-atiention
Achin
U-Net 3D
stacking
Coneaticiion
Figure 9: We tested two architectures for DIAMOND'S ' diffusion model which condition on previous image observations in different ways. To illustrate differences with typical video generation models, we also visualize a U-Net 3D (Çiçek et al. 2016) which diffuses a block of frames simultaneously.
Frame-stacking. The simplest way to condition previous observations is by concatenating the previous frames together with the next noised frame, concat [x; xx__ which is compatible with standard U-Net 2D (Ronneberger et al.) 2015). This architecture 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 (Ciçek et al. 2016). also displayed for comparison in Figure 9 is 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 direct input, but which cross-attends to the activations of separate history encoder network. This encoder a lightweight version of the U-Net 2D architecture. Parameters are shared
:/ htttp:://github.coo/c/mo/ http:///ggggttbo.ooc//////eeerrt
25

---

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
Metries. To evaluate the visual quality of generated trajectories, we use the standard Fréchet Video Distance (FVD) (Unterthiner et al.) 2018) as implemented by Skorookoodo et al. (2022). This is computed between I024 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 al.) 2017), which measures the visual quality of individual observations, ignoring the temporal dimension. For these same sets of videos, we also compute the L.PPS 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, sequence, by a single Nvidia RTX A6000 GPU, per second.
Baselines. We compare against two well-established world model methods; DreamerV3 (Hafne al. 2023) and IRIS (Micheli et al. 2023 adapting the original implementations to train on static dataset. We ensured baselines used a similar number of parameters to DIAMOND. Two variants of IRIS are reported; image observations are discretized into
16 tokens (as used in the original work), or into K = 64 tokens (achieved with one less down/up -assamlin layer in the autoencoder, see Appendix E Micheli et al. (2023). ) which provide the potential for modeling nigher-fidelity visuals.
Compute. All models (baselines and DIAMOND) trained for 120k updates with 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 observations and real sequence of actions.
CS:GO
Driving FVD
Sample rate
Parameters (#)
Method
FID
FVD
LPIPS
FID 1
LPIPS
(Hz)
DreamerV3 IRIS (K 16)
106.8
509.1
0.173
167.5
733.7
0.160
266.7
181M
24.5
110.1
0.129
51.4
368.7
0.188
4.2
123M
IRIS (K 64)
22.8
85.7
0.116
44.3
276.9
0.148
1.5
111M
DIAMOND frame-stack (ours)
9.6
34.8
0.107
16.7
80.3
0.058
7.4
122M
DIAMOND cross-attention (ours)
11.6
81.4
0.125
35.2
299.9
0.119
2.5
184M
Table&rreport 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 the wider video generation literature, where diffusion models currently lead, 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 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 found corresponds to 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.8x 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 fiip-side of this is the low visual quality of its trajectories.
26

---

Figure I0 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, car is plausibly imagined overtaking on the left.
Figure 10: Example trajectories sampled every 25 timesteps from DIAMOND (frame stack) for the modern 3D first-person shooter CS:GO (top row), and real-world motorway driving (bottom row).
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 TT shows the effect real frames, we generate
of the actions in motorway driving conditioned on the same trajectories conditioned on five different action sequences. In general the effects are as intended, e.g. steer straight/left/rihht moves the camera as expected. Interestingly. when 'slow down' 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 expected, though they are unstable beyond 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 symptom of training world models offline data, rather than being an inherent weakness of DIAMOND.
27

---

Continue siraight
Steer
Sleer right
Speed
up
1-25
t=
Figure 11: Effect of fixed actions on sampled trajectories in motorway driving. Conditioned on the same initial observations, we rollout the model applying differing actions. Interestingly, the model has learnt to associate 'Slow down' and 'Speed up" actions to the whole traffic slowing down and speeding up.
Look
Look
Lock
Look
Fire
left
right
down
up
8
Figure 12: Effect of fixed actions sampled trajectories in CS:GO. Conditioned the same initial observation, we rollout the model applying differing actions. Whilst in immediate frames these have the intended effect, for longer roll-outs the observations can degenerate. For instance, it would have been very unlikely for the human demonstrator to look directly into ground in this game state, so the world model unable to generate plausible trajectory here, and instead snaps another area of the map when looking down does make sense.
28

---

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 the title, abstract and introduction are limited to the paper's scope Atari 100k, and match our results in Sections 4 and We include motivational
example from the real world in the introduction but it 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 the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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.
is fine to include aspirational goals as motivation as long as is clear that these goals 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
Guidelines:
The answer NA means that the paper has 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant 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
29

---

Question: For each theoretical result, does the paper provide the full set of assumptions and 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 world modeling. Sections and 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 cross- referenced.
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 they appear in the supplemental material, the authors are encouraged provide 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 ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all the necessary details reproduce results. Our method is explained in Section 3 with the full training procedure detailed in Algorithm iin Appendix G We provide full details of architectures in Appendix [ training hyperparameters in Additionally our entire codebase will be open-source, and provide an anonymized version with this submission: attps://githhe
Appendix E and reinforcement learning objectives in Appendix [
www.m///llaalooooooooooooooooool dimmod
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 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 specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, 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 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 submis- sions 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) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
30

---

(c) If the contribution is new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or 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. the case of closed-source models, 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 have some path 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 instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: The environments consider open-source, and provide open access to our anonymized codebase here: htttps://iithut com/eloialonseeeeeooooooooooooao di diamond Our codebase contains detailed installation and reproducibility instructions. also includes feature to interactively play inside DIAMOND'S 's world model.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
Cooe//ooeeeeeeeeeeeeeeeeeee wwriiissiionooioooiooi 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. only 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Our full training procedure is detailed in Algorithm 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: /ttl::///tttiiee github...... writthub.co/eeeeee www.////walwo/....oooooooooooooa /odioommoo//ddddd
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 level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
31

---

7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We follow the recommended evaluation protocol of Agarwal et al. (2021) provide confidence intervals for our experimental results in Figure 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, confi- dence intervals, statistical significance tests, 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 the standard error of the mean.
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 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 range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We detail the type of GPU, memory and time of execution in Section 4.1 and provide performance comparison in Table
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 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 failed experiments that didn't make into the paper).
9. Code Of Fthics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: http:///eeriip..eee https///c/.pl.cttccccccccceeeee.
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 not foresee any direct harmful societal impact of our work, although we discuss the potential broader impact of our work in Section 9
32

---

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 special consid- eration 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 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 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how 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 models that have 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 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 best faith effort.
33

---

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 the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We' we 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 dataset.
The authors should state which version of the asset is used and, if possible, include 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 in the package should be provided. For popular datasets, paperswi www.de....c........ com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's 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/moddl 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS 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.
34

---

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 equivalent approval/review based the requirements of your country 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 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.
35


Reviews:
[{"review_id": "RjmRMD3Vug", "reviewer": "Reviewer_6rRD", "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.", "questions": "Can 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", "limitations": "Lack of enough experimental evidence other than Atari benchmarks", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "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. \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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We are pleased you believe our work does a good job in demonstrating that using diffusion for world modeling can work well in Atari and recognize our experimental gains on this competitive benchmark.\n\nYour main concern seems to be the focus of our evaluation on the Atari 100k benchmark. First, many recent works (*IRIS* [1], *TWM* [2], *STORM* [3], *BBF* [4]) were widely adopted having solely evaluated on these 26 Atari 100k games, indicating that they are generally considered to be diverse enough to comprehensively evaluate the advantages and drawbacks of various approaches.\n\nSecond, we do apply our diffusion world model to other more visually complex and realistic domains. In particular, we demonstrate that our world model has improved visual quality over *DreamerV3* [5] and *IRIS* on the popular video game *Counter Strike: Global Offensive*, and a real-world motorway driving dataset in Appendix J. We did not consider the less visually complex locomotion tasks to be relevant, but since our method makes no Atari-specific assumptions, our open-source codebase could be applied to other environments of interest by the community.\n\nRegarding comparisons with other non-diffusion model-based approaches you mentioned, we do indeed compare with many non-diffusion methods, including the state-of-the-art model-based approaches on this benchmark, and additional model-free methods in Appendix I. For flow matching approaches, we are not aware of world models that have been designed to leverage this new class of generative models. However, we believe that this would be an interesting direction to investigate in future work, given flow matching provides straighter integration paths more robust to few-step denoising, and enables exact likelihood computation, both of which may be valuable for world modeling.\n\nThank you again for your insightful review. We hope we have addressed your concern regarding our evaluation, and answered your questions on the place of our work in the broader literature.\n\n---\n\n**References**\n\n[1] Micheli et al., *Transformers are Sample-Efficient World Models*, ICLR 2023\n\n[2] Robine et al., *Transformer-based World Models Are Happy With 100k Interactions*, ICLR 2023\n\n[3] Zhang et al., *STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning*, NeurIPS 2023\n\n[4] Schwarzer et al., *Bigger, Better, Faster: Human-level Atari with human-level efficiency*, ICML 2023\n\n[5] Hafner et al., *Mastering diverse domains through world models*, arXiv 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hmSDhKUApC", "reviewer": "Reviewer_h6bf", "summary": "This paper introduces a new world model for learning behaviors in imagination using reinforcement learning.\nIn particular, a diffusion model is used to generate the next frame, conditioned on previous frames and actions.\nAt each environment step, multiple denoising steps are performed to convert a noise image into the next frame.\nBy training an RL agent on synthesized trajectories, the method achieves state-of-the-art performance on the Atari 100k benchmark.", "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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- S1: The world model achieves strong performance on the Atari 100k benchmark, compared with other world models.\n- 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.\n- 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)?\n\nTypo in L240: The single denoising step is shown in the \"last row\" (instead of \"first row\").", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your clear and concise review. We are pleased that you appreciate our idea and strong results.\n\nRegarding your concern with our analysis being mainly qualitative (W1), we agree that a more quantitative measure of the compounding error of different methods for long trajectories would be valuable. Following your suggestion, we extended our analysis to include a plot demonstrating the average drift in generated observations with respect to reference trajectories (provided in our additional one-page PDF). Specifically, we generated 1000-step trajectories with our world model and the real environment, starting from the same frame and following the same action sequence. This new plot confirms that DDPM suffers from accumulating error for small number of denoising steps, and that EDM is more stable even with low number of denoising steps, as illustrated in Figure 3 of the paper.\n\nWhile this plot clearly demonstrates the benefits of EDM over DDPM, it does not show any significant difference between the different numbers of denoising steps for EDM. Therefore, we decided to investigate if EDM with 1-step denoising would affect the down-the-line performance of the model-based RL agent, compared to our 3-step default, for our 10 highest performing games. Our results are displayed in Table 1 in our additional one-page PDF. Even though there is some variance due to the fact we only had time to run a single seed for this ablation, we see there is some signal that the agents trained on the 1-step EDM model perform worse, as the mean HNS on these games dropped from 3.1 to 2.0. The drop is particularly evident on the game *Boxing*, which again confirms our qualitative analysis in Figure 4.\n\nThese additional quantitative results strengthen the justification for our design choices, and would be included (with additional seeds) in a camera-ready version of our paper.\n\nRegarding the breakdown of training times (W2), following your suggestion we have now included a table in our additional one-page PDF (Table 2) demonstrating the breakdown of our training time into individual model steps. We see that a world model step requires around twice as much time to generate the next frame (12.7ms) compared to the reward/termination prediction (7.0ms) using our default procedure with 3 denoising steps. This breakdown also confirms that integrating the reward prediction into the diffusion model (suggested in our limitation section) would be a promising future direction, as it would likely speed up the world model’s imagination. We hope this provides additional insight into our training time and would include this table in a camera-ready version of our paper.\n\nThanks for noticing the typo in L240; we have now fixed it.\n\nRegarding your question around the effect of increasing the imagination horizon (Q1), during development we ran some experiments on this question and did not find much signal that increasing the horizon led to better scores on the games we investigated. In any case, increasing the horizon comes at an increased computational cost, since a longer trajectory must be generated for a single agent update. Additionally, this update may be higher variance earlier in training when the world model is less reliable, so we decided to stick with the default value of 15.\n\nWe hope we have addressed both of your concerns, and believe your suggestions have improved our paper, so thank you again for your review and constructive feedback!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3s4TS4meGq", "reviewer": "Reviewer_JScL", "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.", "questions": "- 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?", "limitations": "The authors have included a dedicated section that discusses the limitations of the work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This 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.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful review. We are pleased that you appreciate the potential of our diffusion-based approach, our thorough analysis of design choices and the clarity of our paper.\n\nWe can understand your minor concern with a lack of more quantitative analysis of our design choices. To address this, we have now added a plot displaying the average drift from reference trajectories for DDPM and EDM-based world models for different numbers of denoising steps in our additional one-page PDF. This plot confirms and quantifies the insights provided by our qualitative analysis in Figure 3, and would be included in a camera-ready version of our paper. \n\nWhile this plot clearly demonstrates the benefits of EDM over DDPM, it does not show any significant difference between the different numbers of denoising steps for EDM. Therefore, we decided to investigate if EDM with 1-step denoising would affect the down-the-line performance of the model-based RL agent, compared to our 3-step default, for our 10 highest performing games. Our results are displayed in Table 1 in our additional one-page PDF. Even though there is some variance due to the fact we only had time to run a single seed for this ablation, we see there is some signal that the agents trained on the 1-step EDM model perform worse, as the mean HNS on these games dropped from 3.1 to 2.0. The drop is particularly evident on the game *Boxing*, which again confirms our qualitative analysis in Figure 4.\n\nThese additional quantitative results strengthen the justification for our design choices, and would be included (with additional seeds) in a camera-ready version of our paper.\n\nTo address your questions:\n\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    \nWhile the performance on these games is indeed worse than some baselines, it is generally in-line with the performance of *IRIS*, for which we have a very similar reinforcement learning pipeline, as described in Section 3.2. This suggests that the performance difference is likely due to differences in reinforcement learning (such as hyper parameter choices not being as well suited to these environments) rather than due to differences in the performance of the world model.\n    \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    \nThese are indeed important details. The current diffusion model is updated in each epoch, not trained from scratch. The training data used is the union of data collected in all epochs collected so far. These details are mentioned in Section 3.2 and Algorithm 1.\n \n> Have you investigated how the diffusion world model evolves as the agent used for collecting the training data improves?\n    \nYes, we found the world-model is generally quite restricted to the policy of the current agent collecting the training data. In *Breakout* for example, we unsurprisingly found that the world model was not able to predict a brick being broken from a higher row (different color) before this happened in the data collected by the real agent, although it was able to generalize to a brick being broken with a ball coming from a different location.\n    \nIn any case, we agree that this is useful for the community to be able to investigate, so have added an option `--pick-checkpoint` to select a particular checkpoint to play with to our codebase. In the following code snippet, the training command is modified to save every checkpoint, and the play command includes the new option to pick a checkpoint.\n    \n```bash\npython src/main.py checkpointing.save_agent_every=1 checkpointing.num_to_keep=null\ncd outputs/<DATE>/<TIME>\npython src/play.py --pick-checkpoint\n```\n\nWe hope that our responses have addressed all your questions and believe that our additions following your review have improved our paper, so thank you again for your constructive feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4BlrTVhXAC", "reviewer": "Reviewer_hqzV", "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.", "questions": "See weaknesses.", "limitations": "Discussion in paper.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Clear and well-written paper\n- Strong empirical results beating prior MBRL baselines on Atari 100K\n- Impressive scaled-up experiments for visual diffusion world modeling, with consistent visual quality over long horizons even with a low number of diffusion steps\n- Good set of ablations over key design choices of the algorithm \n- 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?\n- 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\n\nMinor:\n- 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\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your clear and concise review. We’re pleased you appreciate our carefully written paper, strong results, ablations and open-source code.\n\nTo address your concerns:\n\n> 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?\n    \nWe agree the training time is an interesting perspective, so had already included a training time comparison with latent world model baselines in Appendix H. \n\nIn addition to this comparison with baselines, we have also now included a full training time profile analysis in our additional one-page PDF, which breaks down our total training time in Appendix H into individual steps at the model level. We hope this provides additional insight into the training time of our approach, and we would include this analysis in a camera-ready version of our paper.\n    \n> 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\n    \nWe also agree that it is helpful to compare to model-free algorithms to understand the relative utility of MBRL, so had already included this comparison to SOTA algorithms such as BBF [1] in Appendix I of our paper. We found it interesting to see that there are some games for which model-free algorithms are dominant and some for which *DIAMOND* performs better.\n    \n> 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\n    \nWhile we agree that our background on diffusion is more comprehensive than many diffusion papers and that some readers may already have this background, we initially found the diffusion literature to be hard to parse coming from a RL background. Since our target audience is primarily the RL community, we tried our best to provide a concise, pragmatic and clear introduction to diffusion so that our paper can stand alone and be useful to our colleagues in RL. All appendices in the paper, such as the results on other environments and discussions on the architectures used, are referenced from the main text at the relevant points of the paper, so we hope would still be read by the interested reader.\n    \nWe hope we have addressed all your concerns and questions, and thank you again for your clear and constructive review.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all the reviewers for taking the time to review our paper, and for their positive and constructive feedback.

We are pleased to see a general consensus regarding the motivation of our work, the clarity of our paper, and the strong results achieved by our method.

The main suggestion to improve our paper appeared to be to provide **more quantitative analysis of our ablations** to justify our design decisions. In our paper, we demonstrated visually that an EDM-based diffusion world model is more stable than a DDPM-based world model (Figure 3) and that the use of EDM enabled the use of only 3 denoising steps to provide a fast and reliable world model (Figure 4). **To support this qualitative analysis, we now measure the average drift from reference trajectories** for DDPM and EDM-based world models for different numbers of denoising steps, and provide these results in our additional one-page PDF. **These results confirm and quantify the insights** provided in the qualitative analysis in Figure 3 of the paper.

While this new plot clearly demonstrates the benefits of EDM over DDPM, it does not show any significant difference between the different numbers of denoising steps for EDM. Therefore, we decided to **investigate if EDM with 1-step denoising would affect the downstream performance** of the model-based RL agent, compared to our 3-step default, for our 10 highest performing games. Our results are displayed in Table 1 in our additional one-page PDF. Even though there is some variance due to the fact we only had time to run a single seed for this ablation, we see there is already some signal that the **agents trained on the 1-step EDM model perform worse**, as the mean HNS on these games dropped from 3.1 to 2.0. The drop is particularly evident on the game *Boxing*, which again confirms our qualitative analysis in Figure 4.

**These additional quantitative results strengthen the justification for our design choices**, and would be included (with additional seeds) in a camera-ready version of our paper.

Another point of interest mentioned by multiple reviewers was the **training time of our method**. While we had already included overall training times for comparison with the primary baselines in our paper’s appendix, **we have now included a full profiling analysis** demonstrating the breakdown of our training time into individual model calls in our additional one-page PDF. This breakdown confirms that integrating the reward prediction into the diffusion model (as suggested in our limitation section) would be a promising future direction, as it would likely speed up the world model’s imagination. We would also include this table in a camera-ready version of our paper to **provide additional insight into the training time** and potential improvements to our method.

We hope that we have addressed all of the suggestions raised, and thank all of the reviewers again for their helpful and constructive feedback, which we believe has improved our paper. We look forward to the coming discussion period!

Meta-review: 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.
Decision: Accept (spotlight)