World Models via Policy-Guided Trajectory Diffusion
Authors: Marc Rigter
Jun Yamada
Ingmar Posner
Venue: Accepted by TMLR ()

Abstract: World models are a powerful tool for developing intelligent agents. By predicting the outcome of a sequence of actions, world models enable policies to be optimised via on-policy reinforcement learning (RL) using synthetic data, i.e. in “in imagination”. Existing world models are autoregressive in that they interleave predicting the next state with sampling the next action from the policy. Prediction error inevitably compounds as the trajectory length grows. In this work, we propose a novel world modelling approach that is not autoregressive and generates entire on-policy trajectories in a single pass through a diffusion model. Our approach, Policy-Guided Trajectory Diffusion (PolyGRAD), leverages a denoising model in addition to the gradient of the action distribution of the policy to diffuse a trajectory of initially random states and actions into an on-policy synthetic trajectory. We analyse the connections between PolyGRAD, score-based generative models, and classifier-guided diffusion models. Our results demonstrate that PolyGRAD outperforms state-of-the-art baselines in terms of trajectory prediction error for short trajectories, with the exception of autoregressive diffusion. For short trajectories, PolyGRAD obtains similar errors to autoregressive diffusion, but with lower computational requirements. For long trajectories, PolyGRAD obtains comparable performance to baselines. Our experiments demonstrate that PolyGRAD enables performant policies to be trained via on-policy RL in imagination for MuJoCo continuous control domains. Thus, PolyGRAD introduces a new paradigm for accurate on-policy world modelling without autoregressive sampling.

Published in Transactions on Machine Learning Research (02/2024)
World Models via Policy-Guided Trajectory Diffusion
Marc Rigter, Jun Yamada, and Ingmar Posner
Reviewed OpenReview: http:://ppaarrrsee..eee///mmorod
Abstract
World models are a powerful tool for developing intelligent agents. By predicting the outcome of a sequence of actions, world models enable policies to be optimised via on-policy reinforcement learning (RI.) using synthetic data, i.e. in "in imagination". Existing world models are autoregressive in that they interleave predicting the next state with sampling the next action from the policy. Prediction error inevitably compounds as the trajectory length grows, In this work, we propose a novel world modelling approach that is not autovegressive and generates entire on-policy trajectorics in a single pass through diffusion model. Our approach, Policy-Guided Trajectory Diffusion (PolyGRAD), leverages denoising model in addition to the gradient of the action distribution of the policy to diffuss trajectory of initially random states and actions into an on-policy synthetic trajectory. We analyse the connections between PolyGRAD, score-based generative models, and classifier-guided diffusion models. Our results demonstrate that PolyGRAD outperforms state-of of-tte-rrtt baselines in terms of trajectory prediction error for short trajectories, with the exception of autoregressive diffuusion. For short trajectories, PolyGRAD obtains similar errors autoregressive diffusion, but with lower computational requirements. For long trajectories, PolyGRAD obtains comparable performance to baselines. Our experiments demonstrate that PolyGRAD enables performant policies to be trained via on-policy RL in imagination for MuJoCo continuous control domains. Thus, PolyGRAD introduces new paradigm for accurate on-policy world modelling without autorogressive sampling.
1 Introduction
Model-based reinforcement lcarning (RL) trains a predictive model the environment dynamics, conditionod upon actions. Such models are often referred to as world models. Once a world model has been learnt, imagined (i.e. synthetic) data generated from the world model can be utilised for planning (Schrittwieser et al., 2020) on-policy RL (Hafner et al., 2021). By distilling the knowledge obtained about the environment into the world model, this approach facilitates improved sample efficiency relative to model-free approaches (Micheli et al., 2023) and zero-shot transfer to new tasks (Sekar et al., 2020).
A core challenge of this approach is learning a sufficiently accurate world model: If the world model is inaccurate, the imagined data will not be representative the real environment, and the actions cannot be optimised effectively (Janner et al., 2019). Previous approaches have mitigated modelling errors by (a) using ensembles of models (Chua et al., 2018); (b) predicting the dynamics in latent space with recurrent model (Ha & Schmidhuber, 2018; Hafner et al., 2021); (c) using powerful generative models such transformers (Micheli et al., 2023; Robine et al., 2023) diffusion models (Anonymous, 2023) to generate accurate predictions. Common to all of these existing approaches is that they Iearn single-step transition model, which conditions upon an action sampled from the current policy and predicts the next state. The model is unrolled autoregressively to generate imagined on-policy trajectories one step at a time. Modelling error accumulates at each stop, mcaning that the error in the trajectory inevitably grows as the trajectory length increases (Lambert et al., 2022). In this work, we propose an approach to world modelling that not autoregressive, and generates entire on-policy trajectories in single pass of diffusion.

---

Published in Transactions on Machine Learning Research (02/2024)
Policy-Guided Trajectory Diftusion
π
π
Policy,
Standard Diffusion
Illustration of Policy- Guided Trojectory Diffusion (PolyGRAD). PolyGRAD starts with trajectory of random states and actions and diffuses it into an on-policy trajectory using lcarnt denoising model, e, and the policy, т. Bottom: Training a standard diffusson model on trajectories can be used to generate synthetic trajectories, but these are not on-policy.
Figure 1: Top:
The core challenge for creating world model that is not autoregressive ensuring that the trajectories are generated using actions sampled according to the current policy. In other words: How do we sample actions from the policy at future states if these states have not yet been predicted? This prohibits the standard approach of directly sampling actions from the policy. Instead, Policy-Guided Trajectory Diffusion (PolyGRAD) starts with trajectory of random states and actions and gradually diffuses it into an on-policy synthetic Irajectory (Figure 1). To accomplish this, PolyGRAD utilises a learnt denoising model, in addition to the gradient of the action distribution of the current policy. The denoising model responsible for gencrating accurate prodictions of the environment dynamics, while the policy is used to guide the diffusion process to on-policy trajectories. We analyse how our work can be viewed either as using a score-based generative model to generate on-policy actions,
an instance of massifier-guided diifusson.
The core contributions of this work are: (a) proposing PolyGRAD, the first approach to world modelling that enables on-policy trajectory generation without autoregressive sampling, and (b) analysing the connection between PolyGRAD, score-based generative models, and classifier-guides diffusion models. Our results demonstrate that PolyGRAD outperforms state-of-thhear baselines terms of trajectory prediction error for short trajectories, with the exception of autoregressive diffusion. For short trajectories, PolyGRAD obtains similar errors to autoregressive diffusion, but with lower computational requirements. For long trajectories, PolyGRAD obtains comparable performance to baselines. We show that PolyGRAD can be used to optimise performant policies via ou-policy RL using only imagined data for MuJoCo continuous control domains. Thus, PolyGRAD introduces new paradigm for developing accurate world models without autoregressive sampling.
2 Related Work
World Models Model-based RT. methods (Sutton, 1991) learn model that predicts the transition dynamics and rewards in the environment, conditioned on previous observations and the currett action. Such a model is commonly referred to as "world model" (Ha & Schmidhuber, 2018; Kaiser et al., 2020). Earlier works on model-based RL used MLP ensembles (Chua et al., 2018; Yu et al., 2020) or Gaussian processes (Deisenroth & Rasmussen, 2011) to model the environment dynamics. More recently, the predominant approach has been to use VAE (Kingma & Welling, 2013) map observations to a compact latent space, and predict dynamics in the latent space using recurrent model (Ha & Schmidhuber, 2018; Hafner et al., 2019; Doerr et al., 2018; Hafner et al., 2021). Such models are able to learn from high- dighsdimensional observations in partially-observable environments. Transformers (Vaswani et al., 2017) have been used as more powerful sequence models for generaling accurate dynamics predictions (Micheli et al., 2023; Robine et al., 2023; Schubert et al., 2023; Chen et al., 2022). Likewise, concurrently to our work, diffusion models have also been utilised to generate accurate predictions of individual state transitions model has been learnt, can be utilised for planning (Schrittwieser et al., 2020; Ye et al., 2021) or policy optimisation (Janner et al., 2019). Crucially, by conditioning actions sampled from the current, policy,
a world model (Anonymous, 2023). Once a world
2

---

Published in Transactions on Machine Learning Research (02/2024)
world models enable the generation of unlimited synthetic on-policy data, thus enabling on-policy RL "in imagination" (Hafner et al., 2019). The generation of on-policy synthetic trajectories also makes it possible to visualise the behaviour of the policy before deploying it (Wu et al., 2023). Learning policies entirely from synthetic data has the potential to improve sample efficiency (Micheli et al., 2023) and to facilitate transfer to new lasks without any further environment interaction (Sekar et al., 2020; Rigter et al., 2023).
Common to all of the aforementioned approaches is that they learn a single-step transition model, and use this model autoregressively to generate imagined trajectories of states and actions, one step at time. In autoregressive models, error accumulates as the length of the prediction grows (Janner et al., 2019; Lambert et al., 2022) and therefore the length of the imagined trajectorics is typically limited to a small number of steps (Janner et al., 2019; Hafner et al., 2021). Notably, Lambert et al. (2021) propose to learn a predictive model of environment dynamics that is not autoregressive. Instead, the model conditions upon the parameters of simple controller, such as a PID or LQR controller, and is trained to predict the future state at an arbitrary point in time. Lambert et al. (2021) show that this approach achieves lower error at long-horizon predictions than simple autoregressive baselines. However, limitation of Lambert el al. (2021) is that, the controller parameters used to collect each transition in the dataset must be used as training label. Therefore, this approach is not readily applicable to doep RL, where the mumber of policy parameters can be very large (i.e. millions). Like Lambert et al. (2021) we investigate an approach to trajectory generation that is not autoregressive. However, we aim to sample trajectories under a neural network policy defined by a large number of parameters. To enable this, we propose a novel world model that generates entire on-policy trajectories in a single pass of diffusion. To our knowledge, this is the first work to propose a method for non-auttoregressive on-policy world modelling that is compatible with complex policies.
Diffusion Models for Decision-making Diffusion models are a powerful class of generative model that formulates data generation as an iterative denoising procedure (Sohl-Dickstein et al., 2015; Ho et al., 2020). Learning the denoising model is equivalent to learning the gradients of the data distribution, and therefore diffusion models are an instance of score-based generative models (Song et al., 2021). One advantage of diffusion models is that the iterative sampling procedure enables flexible conditioning, which is commonly utilised to generate bespoke high-quality synthetic images (Dhariwal & Nichol, 2021: Ho & Salimans, 2021).
In the context of sequential docision-making, the use of diffusion models was first introduced by Janner et al. (2022). The authors proposed to train diffussin model to generale trajectories of states and actions, and use classifier musssffer-guidanee guidance to guide the trajectories towards goal states or high rewards thus enabling the diffusion model to be used as planner. This approach was subsequently extended to classifier-fres guidance on state-only trajectories, with an inverse kinematics model to infer the actions (Ajay et al., 2023). Guided diffusion models have also proven to be highly effective in the multi-task setting (He et al., 2023). In each of these works, the output of the diffusion model is used to determine the action to be taken at each time step, making action selection very slow. Furthermore, unlike world models, trajectory generation cannot be conditioned on specific policy actions.
In another line of work, diffusion models have been used extensively as policies capable of accurately representing multi-modal action distributions. Such policies are particularly useful for offline RL (Hansen- Estruch et al., 2023; Wang et al., 2023) and imitation learning (Pearce et al., 2023; Chi et al., 2023), where it is important for the policy to remain near the action distribution in the datasct. In our work, we learn a separate leedforward policy that can be queried quickly during real-world rollouts, and use the diffusiin model to generate on-policy synthetic data for policy training.
There are several previous works that use diffusion models to generate synthetic data for RL (Zhu et al., 2023). SynthER (Lu et al., 2023) trains a diffusion model to generate individual steps of off-policy synthetic data. This data is used to augment the replay buffer of an off-policy RL algorithm to improve sample efficiency. Ding et al. (2024) build on SynthER by generating off-policy trajectories rather than individual transitions. Similarly, MTDIFF-S (He et al., 2023) generates off-policy trajectories to augment offline RL datasets, but in the multi-task setting. Concurrent works on world modelling with diffusion (Anonymous, 2023; Zhang et al., 2023; Yang et al., 2023) train diffusion model to generate single step of transition data, conditioned on the current action and previous observations. Like other existing world models, this approach enables the generation of synthetic trajectories by sampling actions from the policy and autoregressively querying the
3

---

Published in Transactions on Machine Learning Research (02/2024)
model one step at time. In contrast, approach generates entire on-policy trajectories in single pass of diffusion. Also concurrently, and most related to our work, Jackson et. al. (2023) guide diffusion with a policy to increase the likelihood of generated trajectories under the policy, in a similar manner to PolyGRAD. Jackson et al. (2023) use this approach to generate synthetic datasets for offine, off-policy RL. In contrast, PolyGRAD generales on-policy trajectories for online, on-policy RL. in imagination. Unlike Jackson et al. (2023), we also analyse the connection between PolyGRAD, classífier-uuddane models. Furthermore, we show that the scale of the action guidance can be tuned automatically online to approximately generate the correct on-policy action distribution.
and score-based generative
3 Preliminaries
Throughout this paper, we use subscript indices to refer to steps in a diffusion process and superseript indices to refer to time steps in a trajectory through the environment. Thus, x; refers to a data sample at the 2th step in a diffusion process, while a" refers to an action at the t step in a trajectory. If symbol has no subscript, it refers to data sample that has had no noise added (i.e. x =xx)). Bold symbols refer to matrices.
Denoising Diffusion Probabilistic Models Diffusion models (Ho et al., 2020; Sohl-Dickstein et al., 2015) are a class of generative models. Consider sequence of positive noise scales, 0 < S2, ..... .N < 1. In the forward process, for each training data point Xo ~ Px)ta(X), Markov chain Xo, xi...., XN is constructed such that p(x; x = N(x; /1 3,, 8,1). Therefore, Pa: (x, | xu) = NNxx; va,Xo, (1 a; J),, where o= (1 - B,)- We denote the perturbed data distribution as Pa. (x:) )) := Pdata [ p) (x; x)dx. The
noise scales are chosen such that xNN is distributed according to 1), Define s(x,, i) to be the score function of the perturbed data distribution: s(x,, i) == Vx, log (x;), for all Samples can be generated from diffusion model by starting from XN (O, I) and following the recursion:
1
X
VI 3 (x, + Biso(xx, i)) + vB,z,
(1)
where is a learnt approximation to the true score function 8, and is sample from the standard normal distribution. If we reparameterize the sampling of the noisy data points according to: x; vaixo + /1 - a,e, where e ~~ I),, we observe that
Xx, log (x,l xo) =
E
V1-ai
Thus, estimating the score function is equivalent to estimating the noise added. Therefore, we can define the estimated score function in terms of a function Eg that predicts the noise e used to generate each sample
80(x;, i) == eg ex
(3)
VI a1-a
The noise prediction model eø is trained to optimise the objective
0" = arg min
Panta (x)
=olvapxo + VI - a,e, i)]?a.
(4)
Sequential Decision-Making under Uncertainty We consider the setting of fully-observable Markov decision processes (MDPs). An MDP defined by the tuple, M = (S, A, T. R, Ho-7). and A are the state and action spaces, and μa the initial state distribution. T S A ASS) is the transition function, where A(X) denotes the set of possible distributions over X, and y is the discount factor. R: S A + A(R) is the reward function. A policy, n, maps each state to a distribution over actions: π S A(A). We will write ah, + = 50, a. r ro....,s*, shaah, rh to refer to a trajectory of states, actions, and rewards in an MDP with horizon h. -" refers to the sequence of actions only, and -"r refers to the sequence of states and rewards only. The standard objective for MDPs is to find the policy which maximises the total expected discounted reward.
Model-Basod Reinforcement Learning and World Models Model-based approaches to reinforcement learning (Sutton, 1991) (RL) utilise a predictive model of the environment dynamics, commonly referred to as world model (Ha & Schmidhuber, 2018). During online rollouts, trajectories are collected from the
4

---

Published in Transactions on Machine Learning Research (02/2024)
real environment and stored in data buffer D. World models typically use D to learn single-step transition model, T, as an approximation to the dynamics of the enviromment, T. If observations are high-dimensional, this dynamics model is usually learnt in a compressed latent representation of the state (Hafner et al., 2019). Additionally, model of the reward function R is also learnt from the data.
Once the world model has been trained, autoregressive sampling can be used to generate synthetic on-policy trajectories. To generate an imagined trajectory, F, we first sample an initial state, 8 We then sample an action from the policy conditioned on this state, a ~ T--ls)). From the world model, we sample a reward TO ~ R(-1s-, aa) and a successor state, s' ~ (-115 aa)) This process is repeated step-by-step until a trajectory of the desired length has been generated: 7 == ,0,a" aa, ,r0 ...,,,,, a", rh. Because the world model is only an approximation to the environment, the error between imagined and real trajectories accumulates as the length of the trajectory grows. To mitigate this issue, small values of the horizon are typically used (Janner et al., 2019; Hafner et al., 2021). The imagined trajectories can then be utilised for online planning (Argenson & Dulac-Arnold, 2021; Schrittwieser al., 2020) or on-poliey RL (Hafner et al., 2019; 2021) to optimise decision-making, In this work, we will focus on the latter case of on-policy RL for policy optimisation.
4 World Models via Policy-Guided Trajectory Diffusion
In this work, we propose new approach to world modelling: Policy-Guided tRAjectory Diffusion (PolyGRAD). A core novelty of PolyGRAD is that it enables the generation of entire on-policy trajectories in a single pass of diffusion, rather than autoregressively chaining a sequence of one-step predictions. The main challenge for creating a non-autoorerressvve world model is ensuring that the trajectorics generated are sampled according to the current policy. To address this challenge, PolyGRAD utilises diffussin to gradually diffuse a trajectory of initially random states and actions into an on-policy synthetic trajectory.
PolyGRAD utilises two learned components: denoising model, co, and a policy, Te, defined by paramcters 0 and & respectively. In this section, we first describe the denoising model, eg. Second, we describe our main contribution, PolyGRAD, which uses the denoising model in conjunction with the policy to generate synthetic on-policy trajectories via diffusion. Finally, we provide an algorithm for imagined RL in PolyGRAD world models that automatically tunes the magnitude of the action updates in an outer loop.
4.1 Denoising Model Training (Algorithm 1)
Denoising model training in Algorithm 1 follows the same pro-
Algorithm 1 Denoising Model Training denoising model, eg; data buffer, D; number of diffusion steps, N
cess as a standard diffusion model, with the exception that add action-conditioning. To train the denoising model, we sample a sequence of states and rewards, -er = ,", yy.....hh 5h rh and sequence of act ions, T" = a".... ".....,a", ah, from the data busfer, D. The step in the diffusion process, i, is also sampled. Ran- dom noise vī- is added to the state and reward sequence, where e sampled according lo a standard normal distribution and a; is defined by the diffusion noise schedule. The denoising model is trained to predict the noise added to the original state and reward sequence, conditioned upon the actions and the diffusion step. In our implementation, the denoising model is also trained to predict whether a state is terminal, but we omit this from our notation to avoid clutter.
Require:
2: while training: :
3: (T"r, T²) ~ D
4: i in Uniform(1,N)
5: N~,,1
Take gradient descent step on
Volle-eale eelooso + +II age i, (.7")) )1?
4.2 Puidey Trajectory Diffusion (Algorithm 2)
We now present PolyGRAD, our algorithm for generating imagined on-policy trajectories for world modelling (Algorithm 2). Each of the key steps of the algorithm is illustrated in Figure 2. PolyGRAD begins with a trajectory of random states, rewards, and actions, î. = (F +g), sampled from a standard normal distribution (Lines 2 and 3). Note that the subscript N refers to the step in the diffusion process. We also sample an initial state s uniformly from D. The trajectory is then iteratively refined over many diffusion steps. To condition the trajectory on the initial state from the dataset s we perform inpainting (Tugmayr
5

---

Published in Transactions on Machine Learning Research (02/2024)
Algorithm Policy-Guided Trajectory Diffusion (PolyGRAD)
1: Require: policy, Re; denoising model, Ea; action update scale, data buffer D
2: TN I0,
FN I).
4: 5 ~
for i = N,N 1:...,1:
6:
set initial state in 7," to
P Condition trajectory on initial state from datasel
c 1,,
Predict noise
8:
i:
70"
Ti
vo
1-a,
Predict fully-denoised state sequence Update actions using policy score
10:
T--1 4 Ti
+
V
log Tol (
T6) + 332
11: else:
12:
T:-1 + :/
13:
7- V1-0:
2 VB,z
Update states and rewards using diffusion (Equation 1)
1-aa
14: return 7- (70", 70!
ct al., 2022) by replacing the initial state in the trajectory by s at cach step of the diffusion process (Line 6). Generating rollouls branched from states in the dataset is standard practice in model-based RL (Janner et al., 2019). At each diffusion step, the denoising model used to predict the noise added to the state and reward sequence, conditioned on the current action sequence (Line 7). The predicted noise is denoted by
In all steps except. for the final diffusion step, the action sequence is updated in Lines and 10. To perform the action update, first use the predicted noise é to compute prediction of the fully-denoised state sequence, P.. We compute this prediction because the policy is trained on data without noise added, we do not wish to condition the policy on a noisy state sequence. In Line 10, we condition the policy on the predicted denoised state sequence. We update the action sequence in the direction the score of the policy action distribution, Va log (als), in Line 10. Note that this gradient differs from the standard gradient used in policy-gradient RL algorithms (which V log (a(a)s)), and can be easily computed for common policy parameterisations, such as Gaussian policies. The magnitude of the update to the actions is controlled by the action update scale, 6, which In Line 13 the state and reward sequence is updated using the noise prediction and a standard diffusion update (Equation 1). After completing all the diffusion steps, the final synthetic trajectory is returned in Line 14.
shall see in Section 4.3, is tuned automatically online.
In summary, the action sequence is updated using the gradient of the policy action distribution V. log n (a(s)) to increase the likelihood of the actions under the policy. The state and reward sequence is updated according to the learnt denoising model. We must ensure that the state and action sequences are consistent (i.e. the actions are sampled according to the policy conditioned on the predicted states, and the states are predicted according to the sampled actions). To ensure this consistency, the policy is conditioned the current predicted denoised state sequence, and the denoising model is conditioned the current action sequence. Provided that the action update size is chosen correctly, during diffusion the actions are iteratively updated until they are similar to samples drawn from the policy. The state-reward soquence is predicted conditioned upon those actions. Therefore, PolyGRAD is able to generate on-policy synthetic trajectories.
Figure 2: Step-by-step illustration of PolyGRAD trajectory generation (Algorithm 2). Bottom left: Illustration of policy action distribution through- out state space. Top left: Trajectory initialised with random states and actions (Line and 3). The initial state (dark purple) is sampled from the dataset (Line 4). Solid black arrows: Action sequence is updated according to score of policy conditioned on current state sequence (Line 10). Hollow black arrows: State sequence is updated conditioned on current actions (Line 13). Bottom right: Final trajectory returned.

---

Published in Transactions on Machine Learning Research (02/2024)
Now that have described the mechanics of our algorithm, provide theoretical motivation for PolyGRAD by comparing it to score-based generative models and classifier-guided sampling.
Connection to Score-Based Generative Models Langevin dynamics (Rossky al., 1978; Neal, 2011) is Markov chain Monte Carlo method that produces samples from probability density p(x) using only the score function Vx lologp(x pxx.. Given a fixed step size 3 > 0, and arbitrary initial value xo, the Langevin method recursively computes
8 vx logp(xx-1) FB2 2~ N(,, X = -
(5)
2
When 3 0 and t oo, Langevin dynamics converges to sample from p(x) (Song & Ermon, 2019; Neal, 2011). Score-based generative models attempt to learn an approximation to the score function, sx(x) 2 vx p(x), and then perform sampling according to Equation 5 using sp. As noted in Section 3, diffusion models are in instance of soore-based generative models, where the score of the noisy data distribution is learnt across a range of noise levels.
In Line 10 of Algorithm 2, we update the action sequence using the same updale rule as the Langevin method (Equation 5), but with few changes. Unlike standard score-based generative model, where approximation to the score function is learned, in PolyGRAD we directly use the score function of the policy action distribution, V. log (als). Furthermore, the policy score function is conditioned upon the current predicted denoised state sequence, T6 (which also changes throughout the diffusion process). The final modification that we allow the update size to be a tuneable parameter, d. Because the action sequence is updated according Langevin dynamics, we might expect it should converge a sample near the policy action distribution, provided that: a) the state sequence 76 converges, b) the number of diffusion steps is large, and c) and are sized appropriately. A formal analysis of convergence is outside the scope of this work.
Connection to Puasifier-Guidance In classifier-guided sampling (Dhariwal Nichol, 2021), the aim is to sample from the conditional distribution p(x y), where y is a label. Rocall that Pa, denotes the noisy data distribution in the 7th step of diffusion model. To enable sampling from p(x v) using diffusion (Equation 1), instead of approximating the score function Vx, logPo. (x), would like to approximate the score function Vx, Pog (x; y). From Bayes' rule, we have that
Vx, Po, Po: (x: y) Vx, logPo, (x:) + Vx; logp(y x:).
Therefore, if we have a classifier pvy x;), the gradient the classifier with respect to x, can be used to guide the diffusion process to sample from the conditional distribution p(x y). Note that the classifier is evaluated on noisy samples, Xi, so the classifier typically also trained on noisy samples.
In our work, we are interested in sampling on-policy trajectories from the distribution p(T | T&), where va is the current policy. To directly apply maassffcr-guidee sampling to this problem, we would need to learn a differentiable classifier P(T, T:) so that we can evaluate V,, log | T:). However, it is unclear how to train such a classifier. To overcome this issuc, let us consider т, a trajectory that has had no noise added. Applying Bayes' rule, we have that
p(T r) = T) T) (aa 50) R( .a...00 R( P(T)
(a (at(st))))....
(6)
{ ( So (s" ss0
p(a0|50) )p(pp |s!)
where p(als) represents the probability of action a being selected at state s under any policy. Taking the gradient with respect to T, we have
V log
T) V, log (a'ls")
Σ v logp(a'ls").
1=0
1=0
If we assume that p(a[s) is uniform distribution over actions for all s, then the second term zero. Therefore, under this assumption that by default all actions are equally likely, as well the assumption that noise
7

---

Published in Transactions on Machine Learning Research (02/2024)
has been added to the trajectory, arrive at the final classifier gradient:
V l(T | () Σ V. lo'gs(""'ss))
(8)
Deriving the correct classifier glassiier-guiddnce gradient for the case where noise has been added to the trajectory is a subject for future work. To avoid conditioning the policy on noisy trajectories, in PolyGRAD we use the denoising model to predict the fully denoised state sequence, Fis and condition the policy on this sequence to compute the score function (Line of Algorithm 2).
Thus, by considcring classificr guidance arrive at similar update rule for the actions to that used by PolyGRAD. However, the classifier-guidant inspired update in Equation 8 indicates that both actions and states in the trajectory should be updated in the direction of the score of the policy distribution. Meanwhile, in Line 10 of the PolyGRAD algorithm (Algorithm 2), we update only the action sequence in this manner. The state sequence is updated using the denoising model only. In the experiments, we assess the performance of PolyGRAD when both the action and state sequences are updated according to Equation 8.
4.3 Imagined RL in PolyGRAD World Models (Algorithm 3)
Following previous works world models for RL (Halner et al., 2019; 2021), we optimise the policy by performing on-policy RL on the imag- ined trajectories generated by PolyGRAD. We as- sume that for each state, the policy Th outputs
Algorithm 3 Imagined RI. in PolyGRAD World Model
Require: environment, E;
2: Initialise: policy, T&द denoising model eg; action update scale, o; empty data buffer, D
a Gaussian distribution over actions: Ta(als) = ((pa(s), This is the most commonly used policy parameterisation in deep RL (Schulman et al., 2017; Haarnoja et al., 2018). For the most part, Algorithm 3 follows a standard model-based RL training loop. Data is gathered from the real environment, and used to train the denoising model 6p. PoAyGRAD is then used to generate a batch of on-policy synthetic trajectories. These trajectories are used to update the policy via on-policy RL. The key difference with respect to standard based RL is that Line 9 of Algorithm updates 6, which controls the magnitude of the action updates in PolyGRAD.
3: while training:
4:
T t F(TT)
D Sample real trajectory
5: D. Dddd((T)
6: Train eg D
Algorithm
7: Generate synthetic trajectories, {f} & Algorithm 2
8: Train To on { via on-policy RL
9: Update action update scale,
b. Eqqatiin 10
To perform the update for 6, we consider the set of state-action pairs in the synthetic trajectories generated by PolyGRAD, (f). For each state-action pair, (si, ai), we standardise the action according to the mean and standard deviation of the policy action distribution at that state:
0: (s:
02
(x1))
We then compute da, the standard deviation of the set of standardised actions (a). If the actions are drawn correctly from the policy distribution, then the standardised actions should be distributed according to a standard normal distribution. Therefore, we update the action update scale to obtain a standard deviation of in the standardised actions:
8- 5 17 (5.- - 1),
(10)
where 17 is a learning rate. The intuition for Equation 10 is follows. If the policy guidance is too strong, all of the actions will be guided to be very near to the mean of the policy distribution. Therefore, the standardised actions will have a variance that is too low (i.c. below 1) and Equation 10 will reduce the action update scale. Likewise, if the guidance is too weak, the actions will remain spread out far from the policy mean. Thus, the standardised actions will have a variance that is too high and Equation 10 will increase the action update scale. As we shall show the experiments, we found that tuning the action update sizing this manner is sufficient to ensure the PolyGRAD generates a good approximation of the correct on- on-policy action distribution.
8

---

Published in Transactions on Machine Learning Research (02/2024)
Implementation Details Here, outline some key implementation details. A detailed description of our implementation is in Appendix A. After performing each update to the actions in Line 10 of Algorithm 2, we clip the actions to be within 3 standard deviations of the mean of the policy action distribution. We found that this helps to ensure that the action distribution can be reliably guided to the correct on-policy distribution. For imagined RL training Algorithu 3, we used Advantage Actor Critic (A2C) (Mnih el al., 2016) with Generalised Advantage Estimation (GAE) (Schulman et al., 2016). We observed that RL training with PolyGRAD world models is unstable when there are large updates to the policy: if the policy is changed drastically, PolyGRAD may not consistently produce the correct action distribution for on-policy RL training, leading to policy collapse. Therefore, we decay the learning rate for the policy to maintain a constant update size. For the denoising network, we use either a residual MLP or a transformer. As described in Algorithm 1, the denoising network is trained using actions that have no noise added, despite the fact that it is evaluated on (initially) noisy actions in Algorithm 2. However, we found this obtained much better prediction errors and RL. performance than training the denoising network on noisy actions.
5 Experiments
In our experiments, seek to answer the following questions: (a) Does PolyGRAD produce the correct action distribution? (b) How accurate are the trajectories produced by PolyGRAD compared to autoregressive world models? (c) Can the synthetic data produced by PolyGRAD be used to train performant policics? (d) Which implementation details influence the performance of PolyGRAD? To answer these questions, we run experiments using the MuJoCo environments in OpenAl Gym (Brockman et al., 2016). The code for our experiments is available at cechuu..omm/nne /itthubb.cm//macrrriigeee//// pataarrriigerpollyyr---woooooooo
5.1 Does PolyGRAD produce the correct action distribution?
We wish to evaluate whether PolyGRAD generates synthetic trajectories with the correct on-policy action distribution. To evaluate the action distribution produced, we train an RT. policy using synthetic data generated by PolyGRAD (Algorithm 3). We linearly decay the standard deviation of the Gaussian policy, dg from 1 0.02 over the course of training. In Figure 3, we plot the distribution of the difference between the action and the mean of the action distribution, (s), over all state-action pairs (s, a) in batches of data produced by PolyGRAD with h = 50 in Walker2d.
We observe that for Jg 2 0.1, the action distribution closely matches the target Gaussian distribution. This demonstrates that PolyGRAD produces the correct action distribution aggregate provided that the policy entropy is not too low. However, we observe that as the policy standard deviation decreases below 0.1, the action distribution begins to deviate from the larget distribution. Thus, if the policy has very low entropy, it is difficult for PolyGRAD to guide the initially randomly sampled actions to the target distribution. In the plot with 0 0.02, the distribution has additional modes +3 standard deviations. This artifact of the action clipping described in the implementation details. Plots for the other MuJoCo environments are provided in Appendix B.3 of the supplementary material.
Probability Donaty
Probabiltt Detsily
1.00
Probatolltt Denslyy
2.0
probabiit Density
Coneaaooo Denaty
Density
0.4
0.75
1.5
0.50
20
0.2
0.25
0.5
0.0--2 -2
0.00
-1
0
0.0
-0.5
0.0
0.5
-0.25
0.00
0.25
0
0.0 Action
0.2
00.05
0.00 Action
0.05
Action
Action
Action
Action
(a) 0o = 0.8
(b) 0o 0.4
(c) 0e 0.2
(d)
0e 0.1
(e)
0 0.06
(f) 0 = 0.02
Figure 3: Plots of action distributions produced by PolyGRAD. Blue line illustrates the distribution of (s) for batch of synthetic data. Each subplot is for policy with different entropy level that is constant throughout the state space. Dashed black line indicates the action distribution output by the policy.
Data is generated by running Algorithm 3 in Walker2d with h = 50.
9

---

Published in Transactions on Machine Learning Research (02/2024)
PolyGRAD (h- 10)
PolyGRAD Uh - 200)
Transformer
Autoregressive Diffusion
PolyGRAD (b= = 50)
Probabilistic MLP Ensemble
Dreamenva
10
10a
Predction MSE
10
Prediction MSE
10'
Prodiction
10
100
10
10
10-1
10
50 Trajectory Length
10
20
100
Trajectory Length
Trajectory Length
(a) HalfCheetah-v3
(b) Hopper-v3
(c) Walker2d-v3
Figure 5: Plots of mean squared error (MSE) of predicted states vs ground truth states for each world model trained on the same dataset for each environment. Shaded regions indicate standard deviation over seeds. For PolyGRAD, we use the transformer denoising network trained trajectories of length h 10, 50, or 200.
5.2 How accurate are the trajectories produced by PolyGRAD compared to existing world models?
To assess the accuracy of the trajectories produced by PolyGRAD, we train several different world models using the same dataset of 1M transitions in each environment. We then generate synthetic rollouts using the same performant policy in each world model. We replay the same actions in the MuJoCo simulator, and evaluate the errors in the trajectory predictions. We compare the accuracy of the trajectories against the following baselines which all generate trajectories autoregressively:
Probabilistic ML. Ensemble, an ensemble of MLPs outputting a Gaussian over the next state (Chua et al., 2018; Yu et al., 2020; Janner et al., 2019).
Transformer, a transformer-based world model (Micheli et al., 2023).
Dreamer-v3 (Hafner al., 2023), a world model based on sequential VAE.
Autoregressive Diffusion, a diffusiin model trained to generate one-step action-condittoned predictions, rolled out autoregressively (Anonymous, 2023; Zhang et al., 2023).
Detailed descriptions of the implementation of each baseline can be found in Appendix C.3, and further details on the experimental setup can be found in Appendix C.1.
Figure shows the error for each world modelling method for varying trajectory lengths. We train PolyGRAD for trajectories of h = 10, 50, and 200. Goly with h = 10 obtains the best trajectory errors for Walker-2d. For HalfCheetah and Hopper, PolyGRAD with h = 10 obtains the second-best errors to Autoregressive Diffusion. When PolyGRAD is trained on longer trajectories (h = 50 and h = 200) we observe that larger errors are obtained, and PolyGRAD performs less well relative to the basclines. This indicates that is more challenging for the denoising model to accurately predict. the denoised states for longer trajectories.
The errors for the Probabilistic MLP Ensemble increase quickly as the prediction horizon increases. This may be bocause this method outputs Gaussian distribution over the next state, and repeatedly sampling from this Gaussian can quickly lead to out-of-distrriuution states. The Transformer baseline obtains comparable performance to PolyGRAD with = 200, but PolyGRAD outperforms the transformer when trained and evaluated on shorter trajectories. Note that PolyGRAD uses the same transformer network architecture for the denoising model as the Transformer baseline. Therefore, this result demonstrates that there is advantage to iteratively refining the trajectory via diffusion rather than autoregressively predicting each successor state with transformer.
PolyGRAD Chanstormer
Tensformer
PolGRRA (MLP)
Drcamery3
Protutelssic MLP Ensembble
indoregrenazer Difluvon
102
Time (seconda)
1.0
10a
30-1
10-
200
Trajectory Length
Figure 4: Computation times to produce a batch of 1000 trajectories on a V100 GPU in Walker2d.
Interestingly, while Dreamer-v3 performs well on Walker2d, it obtains the poorest prediction errors at short horizons for HalfCheetah and Hopper despite being a stato-of-tho-art world model that produces strong RL agents. This may be because Dreamer-v3 performs dynamics predictions and RL training in latent space, and the reconstruction error to the original state space forms only one part of the training objective. Finally,
10

---

Published in Transactions on Machine Learning Research (02/2024)
Autoregressive Diffusion tends to be the most performant approach, especially on HalfCheetah. This indicates that for these MuJoCo environments, training very accurate one-step prediction model leads to accurate trajectories even at longer horizons. However, the downside of Autoregressive Diffusion is that it is the slowest world modelling method (Figure 4) as requires running diffusion at every step of the trajectory rollout, whereas PolyGRAD requires only a single pass of diffusion.
5.3 Can the synthetic data produced by PolyGRAD be used to train performant policies?
Following from previous works on model-based RL (Janner et al., 2019; Hafner et al., 2021; Yu et al., 2020), we use short imagined rollouts for model-based RL training (h = =0). For RL training, we use the residual MLP denoising network as it oblains similar error to the transformer denoising network at h = 10 (see Figure 10 Appendix B.1) but it is faster for both training and inference (Figure 4).
To assess the performance of RL training in imagination for PolyGRAD (Algorithm 3), we compare the rewards obtained against model-free on-policy RL algorithms (PPO (Schulman et al., 2017), TRPO (Schulman el. al., 2015), and A2C (Mnih et al., 2016)). We compare against on-policy model-free RL algorithms because we perform imagined on-policy training in the PolyGRAD world model. However, note that off-policy model-free RL algorithms obtain stronger performance for these environments (Haarnoja et al., 2018). To compare against a state-of-the-art world model-based agent, we also compare against Dreamer-v3 (Hafner et al., 2023). Reward curves for each algorithm are shown in Figure 6. Figure provides 95% confidence intervals of the interquartile mean of the normalised final performance aggregated across all three environments (see Appendix C.2 for details). Figure shows that PolyGRAD outperforms on-policy model-free algorithms, This is likely because PolyGRAD leverages all of the training data available to train the world model and generate large quantities of synthetic training data. On the other hand, on-policy model-free algorithms do not reuse training data. Confidence intervals in Appendix B.2 show that PolyGRAD significantly outperforms model-free algorithms in terms of mean performance in addition to interquartile mean.
PolyGRAD oblains worse sample-efficiency and final performance compared to Dreamer-v3, despite the fact that PolyGRAD obtains better prediction errors than Dreamer-v3 (Figure 5). This may be because Dreamer optimises policies by directly backpropagating through the dynamics model, while we train policies using policy-gradient RL. Alternatively, it may be because by optimising policies in latent space, Dreamer is able to produce more performant policies despite having less accurate reconstructions of the original state.
PolyGRAD
Dreamer-v3
PPO
A2C
TRPO
1OOD
3000
8
5000
2000
4500
0.0
1000
2000
0.4
5.00 0.5 0.50 3.55
1.00
1.25
3,50
0.5
0.2
0.5
0.8
0.00
0.22
0.50
0.75
1,00
1.5
1.50
0.2
Drnomment Siose
100
Drstunmant Stnzza
Ersearmant Mkgg
0.0
(a) HallCheetah-r3
(b) Hopper-v3
(c) Walker2d-v3
Figure 7: 95% CIs of interquartile Figure 6: Reward curves for RL training. Shaded regions indicate standard mean of normalised final perfor- deviation over 5 seeds.
mance aggregated across envs.
5.4 Which implementation details influence the performance of PolyGRAD?
To assess which implementation details are important, we compare the performance of the base PolyGRAD algorithm against the following ablations and modifications:
Random Actions: We sample random actions at the beginning of the diffusion process from a standard normal distribution, and they are not updated during diffusion.
Policy Sampling:
During each step of the diffusion process, we obtain new actions by directly sampling them from the policy conditionod upon the current predicted state sequence.
No Clipping: The actions are not clipped during the diffusion process.
11

---

Published in Transactions on Machine Learning Research (02/2024)
Add State Update: In addition to updating the actions according to the score of the policy distribution, we also update the states to increase the likelihood of the actions according to Equation 8.
Noisy State Conditioning: Instead of conditioning the policy on the denoised state prediction (76 in Line 10 of Algorithm 2) we condition the policy on the noisy states (7
A detailed description of each of these modifications is provided Appendix C.4. Figure 8 shows the reward curves for each variant of the algorithm, and Figure 9 shows 95% confidence intervals for the interquartile mean of the final normalised performance across all environments. We observe that Random Actions obtains the poorest performance. This is unsurrisiing randomly sampled off-policg act ions result. in incorrect policy gradient updates. Policy Sampling also performs very poorly. This method also fails to generate on-policy trajectories due to the following reasoning. Imagine that we sample an entire sequence of actions, conditioned on the current state sequence. Then, conditioned on these new actions, new prediction of the state sequence is generated at the next diffusion slep. However, the policy has a different action distribution at the new state sequence. Thus, the final trajectory of states and actions produced is not on-policy. We also observed that Policy Sampling results in worse prediction errors. This is likely because diffusion model, which gradually refines its predictions, is ill-suited to making an accurate prediction when the conditioning is completely changed at each diffusion step. Meaningful policies are still learnt when we ablate the action clipping (No Clipping), however the performance is worse, especially for Walker2d. The plots in Appendix B.4 show that without the action clipping, the action distributions produced are more heavy-tailed. Conditioning the policy on the noisy states during diffusion, rather than the denoised prediction (Noisy State Conditioning), also results worse performance. Additionally updating the states during diffussin according to Equation 8 results in comparable performance the base PolyGRAD algorithm. However, this modification increases the computation time required to about 80 hours for 1M environment steps (compared to 54 hours for the base algorithm), it requires backpropagation through the policy at each diffusion step.
we use on-policy RL to optimise the policy, and therefore
PolyGRAD
Random Actions
Clipping
Palicy Sampling
Add State Update
Noisy State Conditioning
1.0
3200
3500
2000
2000
3200
1000
0.00
103 Enviermenr Staps
0.2
Ervtrrment Sieps
103
Eroticemant Stages
105
100
(a) HalfCheetah-v3
(b) Hopper-v3
(c) Walker2d-v3
0.0
Figure 9: 95% Cls of interquartile mcan of normalisod final perfor- mance aggregated across envs.
Figure 8: Reward curves for PolyGRAD, and ablations and modifications to PolyGRAD. Shaded regions indicate standard deviation across 5 seeds.
6 Discussion
Limitations There are two key limitations of our approach. The first that to maintain training stability, we found that it was necessary to update the policy slowly to avoid incorrect action distributions being generated. This likely increases the number of RL training updates required to achieve good performance. The second is that PolyGRAD struggles to achieve the correct action distribution if the policy has very low entropy. This may prevent PolyGRAD from achieving strong RL performance on some domains.
Future Work To improve the stability issues mentioned above, we plan to investigate alternative algorithms for diffusing on-policy trajectorics, with the aim of finding an algorithm that more reliably generates the correct action distribution. We would like to scale PolyGRAD to more complex environments such as image-based environments by utilising latent diffusion (Rombach et al., 2022). We would also like to test PolyGRAD on Mar-Markovian environments. PolyGRAD is well-suited to non-Markovian predictions due to the fact that it is trained to predict entire trajectories, rather than making single step predictions under Markovian assumption. Another direction for future work is to perform theoretical analysis of the
12

---

Published in Transactions on Machine Learning Research (02/2024)
convergence properties of PolyGRAD to better understand if there are conditions under which it is guaranteed
to converge to on-policy trajectories.
Finally, and perhaps most importantly, we would like to investigate whether there are situations in which PolyGRAD obtains better prediction errors at long horizons compared to existing autoregressive baselines. The results in Section 5.2 show that while PolyGRAD obtains strong performance when trained and evaluated on short trajectories (h 10), performance deteriorates relative baselines for longer trajectories. We hypothesise that PolyGRAD may be more robust than autoregressive models when trained small datasets, where the predictions of single-step autoregressive models may be prone to quickly leaving the data distribution, resulting in erroncous predictions. Another line of investigation could be to train the denoising network on short trajectories, and use it to generate longer trajectories. Perhaps this may result in stronger performance and generalisation to longer trajectories. A final approach could be to investigate whether latent diffusion is better suited accurate generation of long trajectories.
Conclusion We have presented PolyGRAD, the first world model that can generate on-policy trajectories without autoregressive sampling. To enable this, instead of sampling actions directly from the policy, PolyGRAD gradually updates the trajectory of states and actions using diffusion. We have shown that PolyGRAD generates a good approximation to the correct on-policy action distribution, and therefore enables on-policy RL in imagination. PolyGRAD achieves solid performance in terms of trajectory prediction accuracy compared to statc-of-tho-art autorogressive basclines. Thus, PolyGRAD introduces a promising new paradigm for world modelling, with many possible directions for extensions and improvements future work.
13

---

Published in Transactions on Machine Learning Research (02/2024)
References
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G Bellemare, Deep reinforeement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 2021.
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making" International Conference on Learning Representations, 2023.
Anonymous. Diffusion world models. In Under review at ICLR, 2023.
Arthur Argenson and Gabriel Dulac-Arnold. Model-based offiin planning. International Conference Learning Represenlations, 2021.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAl gym. arXiv preprint arXiv:1606.cm 015660 2016.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09480 094811, 2022.
Ting Chen. On the importance of noise scheduling for diffusion models. arXiv preprint orXic:2301.1097: 2023.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. Robotics: Science and Systems, 2023.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in Neural Informtiio Processing Systems, 31, 2018.
Mare Deisenroth and Carl E Rasmussen. PILCO: model-based and data-efficient approach to policy search. In International Conference on Machine Learning, 2011.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems, 2021.
Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng. Diffusion world model. arXiv preprint (xXi::402.033700 2024.
Andreas Doerr, Christian Daniel, Martin Schiegg, Nguyen-Tuong Duy, Stefan Schaal, Mare Toussaint, and Trimpe Sebastian. Probabilistic recurrent state-space models. In International Conference on Machine Learning, 2018.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in Neural Information Processing Systems, 2018.
Tuomas Haarnoja, Aurick Zhou, Picter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. International Conference on Learning Representations, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. ar Xiv preprint xiv:22301.04 2023.
14

---

Published in Transactions on Machine Learning Research (02/2024)
Philippe Hansen-Estruch Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit Q-learning as actor-critic method with diffusion policies. orXiv preprint arXiv:2304.1057 10573, 2023.
Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is effective planner and data synthesizer for multi-task reinforcement learning. Advances in Newral Information Processing Systems, 2023.
Jonathan Ho and Tim Salimans. Classificr diffusion guidance. NeurIPS 2021 Workshop on Deep Generalive Models and Downstream Applications, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 2020.
Matthew Jackson, Michael Matthews, Cong Lu, Jakob Foerster, and Shimon Whiteson. Policy-guided diffusion. NewrIPS 2023 Workshop on Robot Learning: Pretraining, Fine- Saaen,,nnng, and Generalization with Large Scale Models, 2023.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. International Conference on Machine Learning, 2022.
Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. International Conference on Learning Hepresentations, 2020.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes arXiv preprint ar1in:3312.6114, 2013.
llya Kostrikov. Pytorch implementations of reinforcement learning algorithms. https://githubb.cml ikostrikor/pytooch--aaa -apercch-accpppooo sa2cc-po-acktr-eee tocktr-gail,, 2018.
Nathan Lambert, Albert Wilcox, Howard Zhang, Kristofer SJ Pister, and Roberto Calandra. Learning accurate long-term dynamics for model-based reinforcement learning. In IEEE Conference on Decision and Control, 2021.
Nathan Lambert, Kristofer Pister, and Roberto Calandra. Investigating compounding prediction errors in learned dynamics models. arXie preprint Xiw:22033009677 2022.
Cong Lu, Philip Ball, and Jack Parker- Holder. Synthetic experience replay. Advances in Neurol Information Processing Systems, 2023.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofto, and Luc Van Gool. Repaint: Inpainting using denoising diffussin probabilistic models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. Interna- tional Conference on Learning Representations, 2023.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International Conference on Machine Learning, 2016.
Radford Neal. MCMC using Hamiltonian dynamics. In Handbook of Markov Chain Monte Carto, chapter 5. Chapman and Hall / CRC Press, 2011.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, 2021.
15

---

Published in Transactions on Machine Learning Research (02/2024)
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Tda Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. International Conference on Learning Representations, 2023.
Antonin Raffin, Ashley Hill, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, and Noah Dormann. Stable basclines 2019.
Mare Rigter, Minqi Jiang, and Ingmar Posner. Reward-free curricula for training robust world models. arXie preprint arXie:2306. 09205,, 2023.
Jan Robine, Mare Höftmann, Tobias Uclwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. International Conference on Learning Represenlations, 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In IEEE/CVF Conference on Compuler Vision and Pallern Recognition, 2022.
Peter J Rossky, Jimmie D Doll, and Harold Friedman. Brownian dynamics as smart monte carlo simulation. The Journal of Chemical Physics, 69(10))4628-4800 1978.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588:604 609, 2020.
Ingmar Schubert, Jingwei Zhang, Jake Bruce, Sarah Bechtle, Emilio Parisotto, Martin Riedmiller, Jost Tobias Springenberg, Arunkumar Byravan, Leonard Hasenclever, and Nicolas Heess. A generalist dynamics model for control. arXiv preprint ar Xiv:223005.1 10912, 2023.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning, 2015.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continu- ous control using generalized advantage estimation. International Conference on Learning Representations, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alee Radford, and Oleg Klimov. Proximal policy optimization algorithms. Xrxiv preprint arXiv:1707.063//3 2017.
Ramanan Sekar, Oleh Rybkin, Kostas Danilidiss, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-suporvised world modols. In International Conference on Machine Learning, 2020.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using noncquilibrium thermodynamics. In International Conference on Machine Learning, 2015.
Yang Song and Stefano Ermon. Gencrative modcling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 2019.
Yang Song, Jascha Sohl-Dickstein, Diedorik Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score- based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021.
Richard Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigort Bulletin, 2(4):160-163, 1991.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Newral Information Processing Systems, 30, 2017.
Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. International Conference Learning Representations, 2023.
16

---

Published in Transactions on Machine Learning Research (02/2024)
Jialong Wu, Haoyu Ma, Chaoyi Deng, and Mingsheng Long. Pre-training contextualized world models with in-the-wild videos for reinforcement learning. Advances in Newral Information Processing Systems, 2023.
Mengjiao Yang, Yilum Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. aiv preprint arXiv:23100061 14, 2023.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 2020.
Lunjun Zhang, Yuwen Xiong, Ze Yang, Sergio Casas, Rui Hu, and Raquel Urtasun. Learning unsupervised world models for autonomous driving via discrete diffusion. arXiv preprinl arxi:2331.00011,, 2023.
Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Yong Yu, and Weinan Zhang. Diffusion models for reinforcement learning: survey, arxie preprint Xiv:22311 11,2222,,,1 2023.
17

---

Published in Transactions on Machine Learning Research (02/2024)
A Implementation Details
Policy Parametcrisation To make it casier to analysc the behaviour of PolyGRAD with policies of different entropy levels, we defined the policy so that the slandard deviation of the policy is controlled by a single learnable parameter Js which is independent of the state. Therefore, the Gaussian policies that we consider are of the form (a(s)s) = (s)(ss), oΦ), where & indicates the parameters of an MLP. This policy parameterisation commonly used in other deep RL implementations (Kostrikov, 2018).
Initial State Conditioning We generated branched synthetic rollouts from initial states within the dataset. To achieve this, we sample the initial state uniformly at random from the dataset. To condition trajectory generation on the initial state, we perform inpainting by replacing the initial state in the trajectory with the sampled initial state during both training and inference. This the same procedure used by Janner et al. (2022).
Action Clipping During Diffusion During diffusion, we clip the action sequence so that it is within three standard deviations of the mean action output by the policy. Therefore, the update in Line 10 of Algorithm 2 is implemented as:
T-11 + clip(?:' + V log (T 7)), min = Ho - 30 max = (ss + 30 + v B,z
(11)
Imagined RL Training For imagined RL training in Algorithm 3, we used Advantage Actor Critic (A2C) (Mnih et al., 2016) with Generalised Advantage Estimation (GAE) (Schulman et al., 2016). We performed one update to the policy and value function for every four steps of data collection in the real environment. During each A2C training update, generated 1024 synthetic on-policy rollouts of length 10. We restricted the minimum standard deviation of the policy, Jg, to be 0.1.
We observed that RL training with PolyGRAD world models is unstable when there are large updates to the policy: if the policy is changed by large amount, PolyGRAD may not consistently produce the correct action distribution. To address this, we defined the update magnitude as the average change in log-likelihood log (als) - log (ata(a]s)) s)] across all state-action pairs in the training batch. For each update to the policy, tuned the learning rate via linescarch so that the average update magnitude target update magnitude, larget loga The hyperparameters used for A2C are shown in Table 1.
within 20% of the
Table 1: A2C Hyperparameters
Parameter
Value
Number of imagined trajectories per update
1024
Imagined trajectory length, h
10
Generalised advantage estimation A
0.9
Critic learning rate Optimiser
3e-4
Adam
Discount factor, Y
0.99
Target policy update, targel a logar Entropy bonus weight Minimum policy std. dev.
0.01
1e-5
0.1
Training steps per environment steps
0.25
Denoising Model The hyperparameters used for the denoising networks and diffusion process are sum- marised in Tables 2, 3 and 4. The noise schedule 3r is defined using a cosine noise schedule (Nichol & Dhariwal, 2021). We used the implementation (2023). We used the default parameter of Thoise to control the noise schodule, with the exception that for Hopper-v3 with the residual MLP denoising network we used Tnoise = 0.1 (which reduces the noise level throughout the early steps of the forward process) as we found this obtained better prediction errors.
cosine noise schedule defined in Algorithm 1 of Chen
18

---

Published in Transactions on Machine Learning Research (02/2024)
For the denoising network, considered both a Residual MLP and Transformer architecture, both trained to minimise the 1.2 loss. The Residual MLP has skip connections at each layer. To add conditioning on the step of the diffusion process, we add a learnable embedding of the diffusion step, i. Therefore, each layer has the form:
LL-1 = (inear(activaaion((z))) + x - embed(1)
To add action conditioning we concatenate the actions to the inpuls so that the input size is (1S) + +A)) x h where h is the length of the trajectory. The final layer projects the output to the correct dimensionality of [S] x h.
In the transformer denoising network, we first. embed each noisy state and action pair using learnt embedding function. The context length is equal to the length of the trajectory, h. We add conditioning upon the timestep in the diffusion process and the timestep in the trajectory by adding learned embedding of each. Each transformer block consists of LayerNorm followed by multi-head causal self-attention layer and a 2-layer MLP. The linear output layer projects each of the h embeddings to the correct dimensionality of (S).
Table 2: Residual MLP Denoising Network Hyperparameters
Parameter
Value
MLP Width
1024 (h 10) or 2048 (h 10)
Batch size
256
Number of layers
Learning rate
3e-4
Optimiser
Adam
Training steps per environment steps
Table 3: Transformer Denoising Network Hyperparameters
Parameter
Value
Embedding dimension
312
Batch size
256
Number of layers
Self-attention heads
Learning rate
1e-4
Optimiser
Adam
Training steps per environment steps
Table Diffusion Hyperparameters
Parameter
Value
Number of diffusion steps, N
128
Noise schedule
cosine
Noise schedule paramcter Tnoise
1 (0.1 for Hopper-v3 + residual MLP)
Hyperparameter Tuning We reduced the cosine noise schedule parameter Tuoise to 0.1 for Hopper-v3 with the Residual MLP denoising network as we found this oblained better prediction errors. We used a larger MLP width of 2048 for the Residual MLP when the trajectory length was greater than 10. All other hyperparamoters were kept constant across all environments. We found that the most critical hyperparameter for RL training was the size of the target policy update, target lo To tune this parameter, we started with large value and decreased it until RL training was stable for all MuJoCo environments.
19

---

Published in Transactions on Machine Learning Research (02/2024)
B Additional Results
B.1 Prediction Error Plots with Transformer and MLP Denoising Networks
Figure 10 compares the trajectory errors for PolyGRAD between the transformer denoising network and the residual MLP denoising network. We observe that for short (rajectories (4 = 10), similar prediction errors are obtained. However, for the longer trajectories the transformer denoising network obtains considerably better performance.
PolyGRAD |Transtormer, 10)
PolyGRAD (Transtormer, h= 200) PolyGRAD (MLP, / 10)
PolyGRAD (MUP. 50)
PolyGRAD [Transformer, h= 50)
PolyGRAD (MLP, h= 200)
10
10
Prodttoon MEE
30P
10h
20
:00
200
10
20
50
100
200
10
20
50
100
(a) HalfCheetah-v3
Hopper-v3
Walker2d-v3
Figure 10: Plots of mcan squared error (MSE) of predicted states ground truth states for PolyGRAD using either transformer residual MLP denoising network. Shaded regions indicate standard deviations over 5 seeds.
B.2 Confidence Intervals for Mean Performance
In the main body of the paper, we reported 95% confidence intervals for the normalised interquartile mean (IQM) generated using the rliable framework (Figures 7 and 9). To provide a more complete picture, here we provide equivalent confidence intervals for the mean, rather than the IQM. We observe that the confidence intervals for mean are similar to those for IQM. This indicates that there are not many outliers in the results, as the mean is effected by outliers while TQM is not.
The intervals in Figure 12 compare PolyGRAD against other model-based and model-free algorithms. We observe that PolyGRAD outperforms the model-free algorithms for mean performance, but does not perform as well Dreamer-v3. Figure compares the base PolyGRAD algorithm against each ablation and variant. of the algorithm. PolyGRAD significatly outperforms each variant abllaioo except for Add State Update.
PolyGRAD
Dreamerv33
PPO
A2C
TRPO
1.0
Nomalsed Mean
0.8
0.4
0.2
0.0
Figure 12: 95% confidence intervals of mean of normalised final performance aggregated across environments.
20

---

Published in Transactions on Machine Learning Research (02/2024)
PolyGRAD
Random Actions
Clipping
Policy Sampling
Add State Update
Noisy State Conditioning
1.0
Noralised Mean
0.8
0.6
0.2
0.0
Figure 14: 95% confidence intervals mean of normalised final performance aggregated across environments for PolyGRAD and each variant.
B.3 Extended Plots of PolyGRAD Action Distributions for All Environments
Figures 15-11 show plots of action distributions produced by PolyGRAD in each of the MuJoCo environments with h = 10. Blue line illustrates the distribution of a - (s(s) for a batch of synthetic data. Each subplot is for policy with a different entropy level. Dashed black line indicates the action distribution output by the policy.
Probabilty Density
Probabilty Density
Probability Density
1.00
Probabiiity Density
0.6
0.4
0.75
1.5
0.4
0.50
1.0
0.2
0.2
0.0
0.0
0.00
Action
Action
Action
Action
(a) 0,
(b) de
0.6
(c)
0.4
(d) o
0.2
Probability Density
Probability Density
Probability Density
Probability Density
20
10
2
0-
0.2
0.0 Action
Action
Action
Action
(e) 0e 0.1
(f) d 0.06
(g) 0e
0.04
(h) 0e 0.02
Figure 15: Walker2d-v3 action distribution plots for trajectories generated by PolyGRAD for 10.
21

---

Published in Transactions on Machine Learning Research (02/2024)
Probablitty Density
Probability Density
Probability Density
1.00
Probability Density
0.6
0.4
0.75
1.5
0.4
0.50
1.0
0.2
0.2
0.25
-2
0.0
0.00
00.5
0.0
0.5
Action
Action
Action
Action
(a) 0g = 0.8
(b) 0e = 0.6
(c) 0g = 0.4
(d) 0p.
0.2
Probability Demsity
Probability Density
Probability Density
Probability Density
20
7.5
15
2-
10
2
2.5
-0.25
0.00
0.25
0.2
0.0 Action
0.2
0.1
0.0
-0.05
0.00
0.05
Action
Action
Action
(e) 06 = 0.1
(f) 0.
0.06
(g) 0, = 0.04
(h) 0e = 0.02
Figure 16: Hopper-v3 action distribution plots for trajectories generated by PolyGRAD for h = 10.
Probability Density
Probability Density
Probability Density
1.00
Probability Density
0.4
0.75
1.5
0.4
0.50
0.2
0.2
0.25
0.5
0.0
-2
0.0
2
0
0.00
-1
0.0
00.5
0.0
0.5
Action
Action
Action
Action
(a)
0.8
(b) dee
0.6
(c) 0
0.4
(d)
0.2
Probabiitty Density
Probability Density
Probability Density
15
5
10
0.25
0.00 Action
0.25
9..2
0.0 Action
0.2
0..1
0.0 Action
0.1
(e)
0.1
0
0.06
(g) 0
0.04
Figure 17: HalfChoctah-v3 action distribution plots for trajectories generated by PolyGRAD for h = 10. NaN values were produced before the minimum policy standard deviation of 0.02 was reached.
B.4 PolyGRAD Action Distributions with Action Clipping Ablated
Figure 18 shows plots of the action distributions in Walker2d when the action clipping during diffusion is removed. Blue line illustrates the distribution of a (o(s) for batch of synthetic data and dashed black line is the policy distribution. We observe that when the action clipping is ablated, PolyGRAD tends to produce a heavy-tailed action distribution with some actions very far from the mean. This means that it still obtains the same standard deviation over actions the policy, but the distribution over actions no longer Gaussain. Thus, PolyGRAD less effective at produce the correct action distribution when the action clipping is removed.
22

---

Published in Transactions on Machine Learning Research (02/2024)
Probability Density
Probability Density
Probability Densily
A
Probability Density
Probabllity Density
0.4
0.6
N
1.0
0.4
0.2
0.2
0.5
-2
0
2
0.0
0.5
0.0 Action
0.5
-0.25
0.00 Action
0.25
Action
Action
Action
(a)
0o = 0.8
(b) 0s = 0.6
(c) dg = 0.4
(d) 0e = 0.2
(e) 0 = 0.1
Figure 18: Walker-v3 action distribution plots when the action clipping has been ablated. Without the action clipping, PolyGRAD tends produce a heavy-tailed action distribution with some actions very far from the mcan.
C
Additional Experiment Details
C.1
Experimental Setup
Error Evaluation Experiments To collect the datasets, we ran Algorithm 3 until we had collected 1M transitions in each environment. Each world model was then trained on the same 1M transitions collected. The final policy produced by Algorithm was used the policy for sampling actions in each world model.
To generate synthetic rollouts, we sampled 500 initial states uniformly at random from the dataset and generated synthetic rollout starting from each of these initial states. We computed the average mean squared error of the prediction across all 500 rollouts for each possible horizon length. We continued to train each world model until obtained the best prediction error evaluation at a. 5 step horizon, up to a maximum of cither 1M gradient steps or 72 hours of training on an RTX 3090 GPU.
C.2 Results Processing
Smoothing of Reward Curves To generate the reward curves in Figures 6 and 8, the average return over 10 episodes of the current policy is evaluated after every 10,000 environment steps. To generate the plots we perform smoothing by computing moving average with window width of 10 evaluations.
Confidence Intervals Figures and present 95% confidence intervals of the interquartile mean (IQM) of the normalised final performance aggregated across all three environments. Figures 12 and are equivalent plots for the mean rather than IQM. To generate the confidence intervals, we use the rliable framework (Agarwal et al., 2021) which is designed compute confidence intervals with limited seeds by aggregating the runs across all environments, and generating the confidence intervals using stratified bootstrapping. IQM is the performance on the middle 50% of combined runs, and is the metric recommended by rliable as it is robust to outliers (unlike the mean) and is more stastically efficient than the median.
To generate the confidence intervals, we first aggregate the final smoothed rewards obtained at the maximum number of environment steps for each run of each algorithm in Figures 6 and Following Agarwal et al. (2021), we first. normalise the rewards in each environment. We normalise the rewards by dividing by the best average final total reward obtained by any algorithm for each environment. We then use the normalised final rewards for each run lo compute the aggregated TQM confidence intervals using the rliable package with 50,000 bootstrapping repeats.
C.3 Baselines
Probabilistic MLP Ensemble We train an ensemble of MLPs that output a Gaussian distribution over the next state. Following Janner al. (2019) and Yu et al. (2020), we train an ensemble of 7 MLPs with 4 layers and 200 hidden units per layer. Each model outputs the mean and variance over the next state using a two-head architecture. When sampling from the model, sample uniformly from the 5 MT.Ps that obtain
23

---

Published in Transactions on Machine Learning Research (02/2024)
the lowest prediction loss a held-out test set. We use the PyTorch implementation publicly available at Xingyuu-iin/mbpo http:uuoiiimmmopppppoo.chhh pytorch.
Transformer We use the same transformer architecture and hyperparameters as for the PolyGRAD transformer denoising model described in Appendix A. Like the denoising model, the transformer is trained using the L, loss. However, instead of the denoising objective, the transformer is trained to predict the next state given the context of a sequence of previous states and actions. We use a maximum context length of 15 stale-action pairs. To generate trajectories, we query the transformer autoregressively.
Autoregressive Diffusion We use the same diffusion model as Synther (Lu et al., 2023), except that we add action conditioning by concatenating the action to the input. This diffusion model is used to make a one-step prediction of the next state, conditioned on the current action. To generate synthetic trajectories, we sample actions from the policy and autoregressively generate one step of the trajectory at a time by completing the reverse diffusion process individually for each step.
Dreamer-v3 We use the implementation of Dreamer-v3 (Hafner For RL training, we perform one step of RL training by
al., 2023) available
at github.com directly performing backpropagation through the dynamics model for every four steps of data collection in the environment. All hyperparameters are set to the defaults used by Dreamer-v3 for MuJoCo. To evaluate the prediction errors, we first initialise the latent state by observing a sequence of five states and actions from the real environment. We then perform open-loop predictions by predicting the next latent states and decoding these using the decoder.
Model-Free RL Algorithms We the implementations of PPO, TRPO, and A2C from Stable Baselines 3 (Raffin et al., 2019). For PPO and TRPO used the default hyperparameters, For A2C we used the default hyperparamcters, with the exception that we set the learning rate to 3e - 4 and the entropy coefficient to le - 5 as we found this worked better for the MuJoCo benchmarks than the default. parameters.
C.4 Ablations and Modifications
We sample random actions at the beginning of the diffusion process from standard normal distribution, and they are not updated during diffusion.
Random Actions:
Policy Sampling: During each step of the diffusion process, we obtain new actions by directly sampling from the policy action distribution TT(-1 7.).
No Clipping: The actions not clipped during the diffusion process.
In addition to updating the actions according to the score of the policy distribution, we also update the states to increase the likelihood of the actions. For this modification, we first compute the noise prediction per Line 7 of Algorithm 2. We use this to compute the estimate of the denoised states, +/:
Add State Update:
Oi
ē
(12)
Vo
We then update the estimate of the denoised states according to by using the score of the policy distribution with respoct to the states:
70 + 8.
log π. ( 9J)
Computing this gradient requires performing backpropagation through the policy. By rearranging Equation 12 we then compute an updated noise prediction based on the modified denoised state prediction
04
E+
VI
24

---

Published in Transactions on Machine Learning Research (02/2024)
This final noise prediction is then used to update the states using the diffusion update in Line 13 of Algorithm 2. Noisy State Conditioning: Instead of conditioning the policy on the denoised state prediction (fī in Line 10 of Algorithm 2) we condition the policy on the noisy states (f,").
C.5 Computational Requirements
For PolyGRAD, one run of RL training up to 1M environment steps (Algorithm 3) requires 54 hours of computation time on an RTX 3090 GPU.
25


Reviews:
[{"review_id": "n1qvR9Z4U9", "reviewer": "Reviewer_bqnZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "As a disclaimer, I am coming from the RL field with very limited knowledge of generative models and diffusion especially. Thus, I cannot fully judge the novelty with respect to related works, or whether the technical solutions to adapt diffusion to on-policy sampling is really substantial. Having said that, from an RL perspective the paper seems to introduce interesting ideas, although the experiments do not fully support the claim that generating long-horizon trajectories in one diffusion pass benefits accuracy w.r.t. autoregressive models of prior works. I summarize below what are, in my opinion, the main pros and cons of the paper. Then, I report some more detailed comments.\n\n*STRENGTHS*\n- How to limit the error propagation across long-horizon trajectories looks like a relevant and important direction;\n- The paper is very well written and easy to follow;\n- The empirical analysis is thorough in terms of validation and ablations;\n- The proposed world modelling algorithm can work in non-Markovian environments (including POMDPs?)\n- The discussion section at the end is upfront in outlining how the method could be improved and interesting directions.\n\n*WEAKNESSES*\n- Experiments does not seem to fully back the claim that generating trajectories in one pass improves over autoregressive models over long horizons;\n- Theoretical analysis is left as future work;\n- The paper does not explore the use of the trained world model for planning;\n- Statistical significance of the results is questionable, as very few seeds are considered (mostly 4) and confidence intervals are not reported (though it is hard to get sensible c.i. without adding more runs);\n- Image generation is arguably the most popular application of diffusion, but the paper only consider physical states instead of learning from images.\n\n*COMMENTS*\n1) With the proposed method, how can we ensure that the model does not generate unfeasible trajectories? E.g., from my understanding the initial state s^0 is forced through, but how can we guarantee that the transition s^0 to s^1 remains feasible? The empirical validation consider (mostly) unconstrained domains, where the latter may be less of an issue.\n2) Can the model be used for planning as is, or that would require some non-trivial modifications?\n3) One drawback that is often mentioned about diffusion is the computational cost of doing inference. Can the authors comment on the computational complexity of the inference, especially w.r.t. with autoregressive models in prior works (computational times are reported in the experiments)? Do the authors believe that this may be a major issue for employing their method for doing planning?\n4) In Algorithm 3, the generator of action sequences is trained at every iteration to keep the data on-policy, but the denoising network for state sequences is also re-trained. I guess this is crucial to avoid using the generator of action sequences out-of-distribution. However, I was wondering whether one could try to collect a dataset with a nice covering of  trajectories, then train the state-sequences generator once and only train the generator for actions during the loop (this may be a dumb idea as I am not knowledgeable, as I said before).\n5) Another perhaps interesting question to ask empirically is whether PolyGRAD trained to generate short trajectories display some generalization when used to generate longer trajectories. Could this experiment be done technically?\n6) The paper mentions that the method suffer from severe instability when paired with policy learning, which is mitigated with very slow policy updates. Providing a theoretical characterization of this behavior seems extremely interesting.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper addresses the problem of learning world models for model-based RL. Differently from previous methods, which are generally autoregressive and generate trajectories one step at a time, the paper proposes to generate full trajectories (state-action-reward sequences) in one-pass through diffusion. The key technical hurdle to adapt diffusion to this setting is to make the model generate on-policy data, i.e., to generate actions that are consistent with the policy conditioned on the generated states. In the paper, the latter is overcame by training a denoising model for state sequences (conditioned on actions sequences) and then interleaving state sequence generations and action sequence updates in the directions of the policy score. The method is plugged into an on-policy policy optimization algorithm and evaluated on various Mujoco tasks (on proprioceptive states).", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think that the changes I am mentioning below can improve the current submission, although a more convincing empirical analysis is what could dramatically improve the value of the paper.\n- I would suggest the author to drop the claim that generating trajectories in one pass improves the accuracy over autoregressive models, which is not fully supported by the experiments, and rework the narrative accordingly;\n- I would refrain to say that Algorithm 3 outperforms PPO/TRPO, as the statistical significance of the improvement is unclear from the given experiments;\n- I would suggest the author to at least explain how the approach can be employed for planning other than policy optimization. Having experiments on planning would be great, but at least discussing how that shall work is already valuable.", "broader_impact_concerns": "This paper fits into the category of fundamental research, as all the experiments are in simulation and directly applying the proposed method in the real-world seems non-trivial.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "CoxEybkUo8", "reviewer": "Reviewer_ZEpn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n1. This is a timely paper that to the best of my knowledge has\n   a reasonable and sound experimental evaluation.\n   Figure 3 clearly demonstrates the ability of fitting\n   the synthetic data.\n   Figure 5 clearly demonstrates the ability of modeling\n   real systems and how will it compares to existing\n   methods as the generated trajectory lengths are increased,\n   and Figure 6 seems like a reasonable attempt to document on\n   using the generations for policy learning.\n2. The idea of policy conditioning is reasonable and\n   clearly motivated throughout the main paper, and the\n   connection to classifier guidance makes a lot of sense.\n\n# Weaknesses\nI do not have any significant concerns with this paper\n--- it is a well-documented idea and investigation.\nMy one concern is that it is difficult to evaluate the contribution in comparison\nto existing approaches as there are no standardized benchmarks\nfor a) learning dynamics models and b) learning policy-conditioned\nmodels.\nFor example, while the experimental results in Figure 5 \ncompare to other methods on MuJoCo experiments, all of the\nbaselines are newly reported in this paper and not directly comparable\nto any existing published results in the literature.\nThe paper seems to do a reasonable job at properly describing\nand setting up the experimental settings, so I do not see\n  this as a strong issue for this paper", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes to use a policy to condition a diffusion\nprocess for generating state sequences in an MDP.\nThis is an advantageous way of sampling trajectories from a\npolicy in an MDP that combines the capabilities of a diffusion\nmodel to generate a temporally and spatially coherent trajectory\nwith the space of feasible trajectories induced by some policy.\nThe main methodological and algorithmic contributions are in\nSection 4 for training/sampling from the policy-conditioned models.\nThen Section 5.1 evaluates the ability of sampling from synthetic\ntrajectories for matching Gaussians. Section 5.2 compares\nmethods for MuJoCo environments (HalfCheetah, Hopper, Walker2d)\nin terms of MSE of the predictions. Section 5.3 uses the\nsampled trajectories for policy learning,. And Section 5.4\nablates design choices.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The paper seems to have been done concurrently with the paper\n[Policy-Guided Diffusion](https://www.robot-learning.ml/2023/files/paper13.pdf),\npresented at the NeurIPS 2023 workshop on robot learning.\nI see this as an indication of the relevant of the idea of\npolicy-guided diffusion and believe it would be helpful to\nconnect to their methodology and experiments in the final version of the paper.", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "v1KuH54LC9", "reviewer": "Reviewer_Rzwm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* interesting, novel formulation with nice proof-of-concept results\n* Interesting discussion of theoretical connections of the proposed approach to Langevin dynamics and classifier guidance.\n\nWeaknesses:\n* The paper claims scalability but only evaluates on 3 simple environments. I think the claims in the paper should be tempered a bit.\n* Typically world models are motivated by a desire for more sample efficiency. In that context the on-policy baselines are not very informative. A comparison with more model-based methods or off-policy RL algorithms would be more meaningful.\n* No improvement on DreamerV3", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Policy-Guided Trajectory Diffusion, an approach that uses a world model to perform on-policy RL on trajectories “in imagination”. Crucially, the proposed approach does not rely on autoregressive sampling but denoises full trajectories. To keep the action compatible with the current policy the algorithm updates them using the policy score function. \n\nThe proposed method is evaluated on 3 simple continuous control environments and appears to outperform on-policy RL baselines but not DreamerV3, a strong world model baseline. Furthermore the paper presents a number of ablations.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* page 6: I think there’s a log missing in line 2 \n* Consider citing Schubert et al. \"A Generalist Dynamics Model for Control.\" arXiv preprint arXiv:2305.10912 (2023) as another recent paper on transformer dynamics models.\n* It is claimed that PolyGrad outperforms on-policy RL. However from the figure this could be due to statistical noise. Could you include a test?\n* Not sure if Welling & Teh is the right reference for Langevin dynamics, since that paper is about using it with stochastic gradients. \n* Give smoothing details in figures in e.g. figure 6.\n* Please give intuition for why eq (10) makes sense. I understand that there should be a fixed point but why this update?\n* Why A2C? Why not PPO as a base algorithm?\n* I’m not sure the paragraph on diffusion augmentation methods in the related works section needs to be there. These papers are quite different.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is easy to read and the reviewers generally agreed on acceptance. The decision is Accept with Minor Revisions, as I would like to make two requests that I can check without further review.

Primarily, the number of runs should be consistent between methods. But some have 4 runs and others have 5. Please fix this, or explain why it is difficult to fix. 
(And once you are fixing this, I highly encourage you to do more runs, even just 3-5 more runs since it will make results more significant). 

Second, please include another figure like Figure 7, but showing 95% CIs for more than just the interquartile. It is not required that the outcome is significant in this case, but it provides a more complete picture, especially because it better matches the mean reported in learning curves. Also note that if it is significant, then you can claim that it is significant according to both metrics. 

Minor question that does not have to be answered in the paper but I am curious about: why not also learn on real trajectories? Namely, learn on a mix of synthetic and real trajectories? Is there a complication with doing this that I am missing? Or did you find it performs more poorly?
Decision: Accept (poster)