Reasoning with Latent Diffusion in Offline Reinforcement Learning
Authors: Siddarth Venkatraman
Shivesh Khaitan
Ravi Tej Akella
John Dolan
Jeff Schneider
Glen Berseth
Venue: ICLR 2024 poster (2024)
Keywords: Reinforcement Learning
Diffusion Models
Primary area: reinforcement learning
TL;DR: We leverage latent diffusion models to learn skill representations with which we learn high value policies from offline datasets.

Abstract: Offline reinforcement learning (RL) holds promise as a means to learn high-reward policies from a static dataset, without the need for further environment interactions. However, a key challenge in offline RL lies in effectively stitching portions of suboptimal trajectories from the static dataset while avoiding extrapolation errors arising due to a lack of support in the dataset. Existing approaches use conservative methods that are tricky to tune and struggle with multi-modal data or rely on noisy Monte Carlo return-to-go samples for reward conditioning. In this work, we propose a novel approach that leverages the expressiveness of latent diffusion to model in-support trajectory sequences as compressed latent skills. This facilitates learning a Q-function while avoiding extrapolation error via batch-constraining. The latent space is also expressive and gracefully copes with multi-modal data. We show that the learned temporally-abstract latent space encodes richer task-specific information for offline RL tasks as compared to raw state-actions. This improves credit assignment and facilitates faster reward propagation during Q-learning. Our method demonstrates state-of-the-art performance on the D4RL benchmarks, particularly excelling in long-horizon, sparse-reward tasks.

Published as a conference paper at ICLR 2024
REASONING WITH LATENT DIFFUSION IN OFFLINE REINFORCEMENT LEARNING
Siddarth Venkatraman' Shivesh Khaitan 2*
Ravi Tej Akella?
John Dolan 2
Jeff Schneider"
Glen Berseth
Mila, Université de Montréal
CCamegie Mellon University
"qqual Contribution
siddarth. venkat raman .uqubbb
ABSTRACT
Offline reinforcement learning (RL) holds promise as reward policies from a static dataset, without the need for further environment interactions. However, key challenge in offline RL lies in effectively stitching portions of suboptimal trajectories from the static dataset while avoiding extrapo- lation errors arising due to lack of support in the dataset. Existing approaches use conservative methods that are tricky to tune and struggle with multi-modal data (as we show) or rely on noisy Monte Carlo return-to-go samples for reward condition- ing. In this work, we propose a novel approach that leverages the expressiveness latent diffusion to model in-support trajectory sequences as compressed latent skills. This facilitates learning Q-function while avoiding extrapolation error via conttroontiainnng. The latent space is also expressive and gracefully copes with
learn high-
data. We show that the learned encodes richer task-specific information for offline RI. tasks as compared to raw atate-actions. This improves credit assignment and facilitates faster reward propa- gation during Q-learning. Our method demonstrates state-of-the-att performance across the D4RL suite, particularly excelling in long-horizon, sparse-reward tasks.
sapporrlly-abstraae latent space
INTRODUCTION
Offline reinforcement learning (RL) offers a promising approach to learning policies from static datasets. These datasets are often comprised of undirected demonstrations and suboptimal sequences collected using different behavior policies. Several methods (Fujimoto et al. (2019); Kostrikov et al.; Kumar et al. (2020)) have been proposed for offline RL, all of which aim to strike a balance between constraining the learned policy to the support of the behavior policy and improving upon it. At the core of many of these approaches an attempt to mitigate the extrapolation error which arises while querying the learned Q-function in order to extract the best policy from the data, Q-lcarning uses an argmax over actions to obtain the temporal- -iffeeeeee target. However, querying the Q-function on out-of-support state-actions can lead to errors via exploiting an imperfect Q-function (Fujimoto et al. (2019)).
out-of-support samples for policy improvement. For example,
Framing offline RL generative modeling problem has gained significant traction (Chen et al. (2021); Janner et al. (2021)); however, the performance is dependent on the power of the gener- ative models used. These methods either avoid learning Q-function or rely on other offline Q- learning methods. Recently diffusion models (Sohl-Dickstein et al. (2015); Song & Ermon (2019)), have emerged as state-of-the-ar generative models for conditional mmage-generation (Ramesh al. (2022); Saharia et al. (2022a)). Rather than avoiding Q-learning, we model the behavioral policy with diffusion and use this to avoid extrapolation error through catch-constraining. Previous diffusion- based sequence modelling methods in offline RI. diffused over the raw state-action space. However, the low-level trajectory space tends to be poorly suited for reasoning. Some prior works (Pertsch et al. (2021); Ajay al. (2020)) have proposed to instead reason in more well-conditioned spaces composed of bigher-level behavioral primitives. Such temporal abstraction has been shown to result in faster and more reliable credit assignment (Machado et al. (2023); Mann & Mannor (2014)), particularly in long- morizooi sparse-reward tasks. We harness the expressivity of powerful diffu- sion generative models to reason with temporal abstraction and improve credit assignment.

---

Published as a conference paper at ICLR 2024
Inspired by the recent successes of Latent Diffusion Models (LDMs) (Rombach et al. (2022); Jun & Nichol (2023)), we propose learning similar latent trajectory representations for offline RL tasks by encoding rich high-level behaviors and learning a policy decoder to roll out low-level action sequences conditioned on these behaviors. The idea is to diffuse over semantically rich latent representations while relying on powerful decoders for high-frequency details. Prior works which explored diffusion for offline RL (Janner et al. (2022), Ajay et al.) directly diffused over the raw state-action space, and their architectural considerations for effective diffusion models limited the networks to be simple U-Nets (Ronneberger et al. (2015)). The separation of the diffusion model from the low-level policy allows us to model the low-level policy using a powerful autoregressive decoder. We perform state-conditioned latent diffusion on the learnt latent space and then learn a Q-function over states and corresponding latents. During Q-learning, we batch-constrain the candidate latents for the target Q-function using our expressive diffusion prior, thus mitigating extrapolation error. Our final policy samples latent skills from the LDM, scores the latents using the Q-function and executes the best behavior with the policy decoder. We refer to our method as Latent Connstrine Q-learning (LDCQ).
There proposed latent diffusion skill learning method offers several advantages:
Flexible decoders for high-fidelity actions. The latent diffusion framework allows us use pow- erful decoders for our low-level policy ny- Previous diffusion works for offline RI. (Janner et al. (2022), Ajay al.) directly diffused over the raw state-action space, and architectural considera- tions for effective diffusion models limited the networks to be simple U-Nets (Ronneberger et al. (2015)). The separation of the diffusion model from the low-level policy allows us to model my using an expressive autoregressive decoder. (Model architecture discussed in Appendix A.2). We also note that LDMs can be easily used to model trajectories with discrete action spaces, since the decoder and diffusion models are separated.
Temporal Abstraction with information dense latent space. Prior works (Pertsch et al. (2021); Ajay et al. (2020)) have learned latent space representations of skills using VAFs. Their use of weaker Gaussian priors forces them to use higher values of the KL penalty multiplier to ensure the latents are well regularized. However, doing restricts the information capacity of the latent, which limits the variation in behaviors captured by the latents. As we show in section 5.1, increasing the horizon H reveals clear separation of useful behavioral modes when the latents are weakly constrained. Our method allows modeling the dense latent space with diffusion.
Faster training and inference. Inference with LDMs is significantly faster than having to recon- struct the entire trajectory every forward pass with a raw trajectory diffusion model. The training process is also more memory efficient since the networks can be much smaller.
Our method excels at long-horizon credit assignment through temporal abstraction, which allows it outperform prior offline RL algorithms in the challenging sparse reward antmaze and franka- kitchen tasks. Further, the expressiveness of diffusion models also facilitates greatly improved batch- constrained Q-learning outperforming prior work in the Adroit suite and the image-based Carla lane driving task. Through these strong results, we show that Catch-Constrained Q-learning is a much stronger method than prior work has indicated, when paired with more expressive generative models and temporal abstraction. We also demonstrate how LDCQ can be extended to goal- conditioned reinforcement learning.
2 RELATED WORK
Offline RL. Offline RL poses the challenge of distributional shift while stitching suboptimal trajec- tories together. Conservative Q-Learning (CQL) (Kumar et al. (2020)) tries to constrain the policy to the behavioral support by learning a pessimistic Q-function that lower-bounds the optimal value function. Implicit Q-Learning (IQL) (Vieillard et al. (2022)) tries to avoid extrapolation error by performing a trade-off between SARSA and DQN using expectile regression. Inspired by notable achievements of generative models in various domains including text-generation (Vaswani et al. (2017)), speech synthesis (Kong et al.) and (2022a)), Chen et al. (2021) proposed to generative model for offline RL and bypass the need for Q-learning or bootstrapping altogether with return catur-cconitioonnn (Srivastava et al. (2019); Ku- mar et al. (2019)). Our method instead formulates a solution with batch-constraining which uses
(Ramesh et al. (2022); Saharia et al.
2

---

Published as a conference paper at ICLR 2024
generative models to model the data distribution and use it to generate candidate actions to learn Q-function without extrapolation-ena erative model does not sample out-of-support samples, which has been difficult to achieve with previously used generative models in offine RI.. This is a form of Batch-Constrained Q-Learning (BCQ) (Fujimoto al. (2019)). Further, to effectively address the problem of stitching, Pertsch et al. (2021) and Ajay et al. (2020) proposed learning policies in latent-traeectory spaces. However, they have to rely highly constrained latent space which is not rich enough for the downstream policy. Our proposed method to use latent diffusion, which can model complex distributions, allows for the needed flexibility in the latent space for effective Q-learning and the final policy.
This relies on the assumption that sampling from the gen-
et al. (2015); Song & Er- mon (2019)) have emerged as state-of-the-art generative models for conditional image-generation (Ramesh et al. (2022); Saharia et al. (2022a)), super-resolution (Saharia et al. (2022b)) and inpaint- ing (Lugmayr et al. (2022)). Recent offline RL works (Janner et al. (2022), Ajay et al.) have proposed using diffusion to model trajectories and showcased its effectiveness in stitching together behaviors to improve upon suboptimal demonstrations. However, Janner et al. (2022) make the assumption that the value function is learnt using other offline methods and their classifier -uiieed diffusion requires querying the value function on noisy samples, which can lead to extrapolation- error. Similarly, Ajay et al. suffer from distributional shift, it relies on return-conditioning. and maximum returns from arbitrary states can be unknown without access to a value function. We propose method for learning Q-functions in latent trajectory space with latent diffusion while avoiding extrapolation-erer -eroor and facilitating long horizon trajectory stitching and credit assignment.
Diffusion Probabilistic Models. Diffusion models
3 PRELIMINARIES
3.1 DIFFUSION PROBABILISTIC MODELS
Diffusion models (Sohl-Dickstein et (2015); Song & Ermon (2019)) are a class of latent variable generative model which learn to generate samples from probability distribution p(x) by map- ping Gaussian noise the target distribution through an iterative process. They are of the form Pe (Xo) == f pv (xo.)) )dx:) where Xo... Xr are latent variables and the model defines the ap- proximate posterior q(xIT Xo) through fixed Markov chain which adds Gaussian noise to the data according variance schedule $ 3r. This process is called the forward diffusion pro- cess:
T II (xx x)) I=l
((xi xo)
g(xe xx-1) == N(x;: v Bix
(1)
The forward distribution can be computed for an arbitrary timestep t in closed form. Let a, = 1 3, and o = I O;. Then q(x, xx) = (x:: Varxo, (1 a.))).
Diffusion models learn sample from the target distribution p(x) by starting from Gaussian noise p(xT) ~ (NO, and iteratively denoising the noise to generate in-distribution samples. This is defined as the reverse diffusion process Po (X x:):
Ps (Xo.T) == p(xT) II P (x.-1 xt),
Pe (xe-1 := (<x==:((xx1 1), (xe, ()) (2)
The reverse process is trained by minimizing surrogate loss-function (Ho et al. (2020)):
((v)= = (
cy. (xe, I|2
(3)
Diffusion can be performed in a compressed latent space z (Rombach et al. (2022)) instead of the final high-dimensional output space of x. This separates the reverse diffusion model py (Z, 2r)
from the decoder (x z). The training is done in two stages, where the decoder is jointly trained B- V-Variational Autoencoder (Kingma & Welling; Higgins et al. (2016))
with encoder, similar with a low 3. The prior is then trained to fit the optimized latents of this model. We explain this two-stage training in more detail in section 4.1.
3

---

Published as a conference paper at ICLR 2024
3.2 OFFLINE REINFORCEMENT LEARNING
The reinforcement learning (RL) problem can be formulated as a Markov decision process (MDP). This MDP is tuple (pu. A.A. P. ?), where Po is the initial state distribution, S is a set of states, A is set of actions, r S x A IR is the reward function, P S x A x S + [0. 1] is the transition function that defines the probability of moving from one state to another after taking an action, and E [0, 1) is the discount factor that determines the importance of future rewards. The goal in RI. to learn a policy, i.e., a mapping from states to actions, that maximizes the expected cumulative discounted reward. In the offline RL setting (Levine et al., 2020), the agent has access to a static dataset = {s} ai, $,41 ri} of transitions generated by a unknown behavior policy T s) and the goal is to learn new policy using only this dataset without interacting with the environment. Unlike behavioral cloning, offline RL methods seek to improve upon the behavior policy used to collect the offline dataset. The distribution mismatch between the behavior policy and the training policy can result in problems such as querying the target Q-function with actions not supported in the offline dataset leading the extrapolation error problem.
LATENT DIFFUSION REINFORCEMENT LEARNING
In this section, we elaborate on our latent diffusion- method for offline RL.
4.1 TWO-STAGE LDM TRAINING
Latent Representation and Low-Level Policy. The first stage in training the latent diffusion model is comprised of learning a latent trajectory representation. Given a dataset D of /I-length trajectories Th represented as sequences of states and actions, So, ag, , ... low-level policy (a ss, z) such that z represents high-level behaviors in the trajectory. This is done using a 3- Variational Autoencoder (VAE) (Kingma & Welling; Higgins et al. (2016)). Specifically, maximize the evidence lower bound (ELBO):
- - - we want to learn a
[(8, w) E (Z)TT)
log No (a [s. z) (Dxz(g((z)) TH) |I p (z so))]
(4)
where is represents our approximate posterior over z given THe and Pw represents our conditional Gaussian prior over z, given Sp- Note that unlike BCQ (Fujimoto et al. (2019)), which uses the VAE as the generative model, we only use the B-VAE to learn a latent space diffuse over. As such, the conditional Gaussian prior Pw is simply a loose regularization of this latent space, and only weakly constrains the posterior. This is facilitated by the ability of latent diffusion models to later sample from such complex latent distributions. Prior works (Pertsch et al. (2021); Ajay et al. (2020) have learned latent space representations of skills using VAEs. Their use of weaker Gaussian priors forces them are well regularized. However, doing so restricts the information capacity of the latent, which limits the variation in behaviors captured by the latents. As show in Section 5. increasing the horizon H reveals a clear separation of useful behavioral modes when the latents are weakly constrained.
higher values of the KL penalty multiplier 3 to ensure the latents
The use of latent diffusion gives us flexibility make the decoder more powerful. The low-level policy тæ is represented as an autoregressive model which can capture the fine details across the action dimensions, and is similar to the decoders used by Ghasemipour et al. (2021) and Ajay et al. (2020). While the environments we test in this work use continuous action spaces, the use of latent diffusion allows the method to easily translate to discrete action spaces too, since the decoder can simply be altered to output a categorical distribution while the diffusion process remains unchanged.
Latent Diffusion Prior. For training the diffusion model, we collect a dataset of state- pairs (So. z) such that ти D is H H-length trajectory snippet, 2 ~ ( TH) where go is the VAE encoder trained earlier, and So is the first state in TH. We want to model the prior p(z So). which is the distribution of the learnt latent space z conditioned different behaviors possible from the state Su as supported by the behavioral policy that collected the dataset. To this end, learn conditional latent diffusion model p So) by learning the dependent denoising function Hy (z,, Sp, t), which takes as input the current diffusion latent estimate 2, and the diffusion timestep & to predict the original latent 2g. Like Ramesh et al. (2022)
state So. This effectively represents the
4

---

Published as a conference paper at ICLR 2024
and Jun & Nichol (2023), found predicting the original latent Zo works better than predicting the noise e. We reweigh the objective based on the noise level according to Min-SNR-y strategy (Hang et al. (2023)). This re-balances the objective, which is dominated by the loss terms corresponding to diffusion time steps closer to T. Concretely, we modify the objective in Eq. 3 to minimize:
((6) [1,T) go (x/mw) ( (2.) [min (SNR((), (]l 2o (w Su, t) [1²)] (5)
Note that (2 TH) is different from g(z; Zn), the former being the approximate posterior of the trained VAE, while the latter is the forward Gaussian diffusion noising process. We use DDPM (Ho et al. (2020)) to sample from the diffusion prior in this work due to its simple implementation. As proposed Ho & Salimans, we use classífier-free guidance for diffusion. This modifies the original training setup to learn both conditional Me (2,, So, I) and an unconditional model. The unconditional version, is represented as pu (z,,0, t) where a dummy token Ø takes the place so- The following update is then used to generate samples: p (z, 0, +w(wp (Z,, (Z,, So. t) p (z, (z))). where w a tunable pyper-parameter. Increasing w results in reduced sample diversity, in favor of samples with high conditional density. We summarize the two stage LDM training in Appendix 2.
Nolsed Latent
a
11-1
Forward Diftusion Process
x z)
8
(x(x)x)
argmax Q(.
times
8
Dencising
Denolxing Network
candidates
xy(a (a s, x)
Z
Z
Diflusion (2)x)
Policy Decoder
H-11
a) Latent Space Representation Training
b) Latent Diffusson Model Training
Policy Execution
Figure 1: Latent Diffusion Reinforcement Learning Overview a) We first learn the latent space and low-level policy decoder by training 8-VAE over H-length sequences from the demonstrator dataset. b) We train a latent diffusion prior conditioned on Sy to predict latents generated by the VAE encoder. e) After learning Q function using LDCQ (Algorithm 1), score latents sampled by the prior with this Q function and execute the low-level policy To conditioned on the argmax latent.
4.2 LATENT CHIFUSION-CONSTRAINNG Q-LEARNING (LDCQ)
In batch-constrained Q-learning (BCQ), the target Q-function is constrained to only be maximized using actions that were taken by the demonstrator from the given state (Fujimoto al. (2019)).
(s) argmax Q(s, a)
In deterministic MDP setting, BCQ is theoretically guaranteed converge the optimal batch- constrained policy. In any non-trivial setting, constraining the policy to actions having support from given state in the dataset is not feasible, especially if the states continuous. Instead, behavior model (a s) must he learned on the demonstrator data and samples from this model are used as candidates for the argmax:
T (s)
argmax Q(s, a,)
(als)
However, in many offline RL datasets, the behavior policy is highly multimodal either due to the demonstrations being undirected because the behavior policy is actually mixture of unimodal
5

---

Published as a conference paper at ICLR 2024
policies. making it difficult for previously used generative models like VAEs sample from the distribution accurately. The multimodality of this policy is further exacerbated with increases in temporal abstraction in the latent space, as we show in section 5.1. We propose using latent diffusion to model this distribution, as diffusion is well suited for modelling such multi-modal distributions. We propose to learn a Q-function in the latent action space with latents sampled from the diffusion model. Specifically, we learn a Q-function Q(s, z). which represents the aetion-value of latent action sequence given state S. At test time, we generate candidate latents from the diffusion prior (z s) and select the one which maximizes the learnt Q-function. We then use this latent with the low-level policy (a S, z) to generate the action sequence for H timesteps.
Training. We collect a replay buffer B for the dataset D of H Hennggth trajectories and store transition tuples (st, 2, Tt:t-ff; ;1; 8t+H)) from Tи ~ D, where 8, the first state in TH, z ~ ( TH) is the latent sampled from the VAE approximate posterior, T+++= represents the 7-discounted sum of rewards accumulated over the H time-steps in Tи, and 8 represents the state at the end of H--length trajectory snippet. The Q-function is learned with temporal-difference updates (Sutton & Barto (2018)), where we sample a batch of latents for the target argmax using the diffusion prior (z S ). This should only sample latents which are under the support of the behavioral policy, and hence with the right temporal abstraction, allows for stitching skills to learn an optimal policy grounded on the data support. The resulting Q update be summarized
( z) (Te4* + Q(S++H, arginax
(Q($L+H,
(8)
2e (x)
We use Clipped Double Q-learning as proposed in (Fujimoto et al. (2018)) to further reduce over- estimation bias during training. We also use Prioritized Experience Replay (Schaul al. (2015)) to accelerate the training in sparse-reward tasks like AntMaze and FrankaKitchen. We summarize our proposed L.DQ method in Algorithm 1.
Algorithm Latent Diffusion-CCnsttranne Q-Learning (LDCQ)
1: Input: prioritizeedrreply-bbuu B, horizon HI, target network update-rate p. mini-batch size
N, number of sampled latents n, maximum iterations M, discount-factor % latent diffusion denoising function Hig, variance schedule a1,..... Or,, ai..... "2....,,ār, ,,ār, .I...., .......,.. .r.
2: Initialize Q-networks Qe, and Qe with random parameters Qe Qo, and target Q-networks
Qetaraac and gtaara with 6far saarget + 1, O targgt
O
3: for iler = 1 to M do
4: Sample a minibatch of N transitions {(S,, 2, T 81+H)} from B
5: Sample n latents for each transition: 27 NO,
6: for = T to 1 do
> DDPM Sampling
7:
ub 0,t) wH (Ze, StHH, M (2, (,t)
2,-1 ~ N/ (a 1) 2, +
va.- 2, T(l > 1) l-āt
8:
l-a.
9:
end for
Compute the target values y = Tr.t + H (max min Qotaron 20)}} zo 1=1,2
11:
Update Q-networks by minimizing the loss: 1 lly Q 2)1B
12:
Update target Q-networks: Otarget + p + (1 p) Starget
13:
end for
Policy Execution. The final policy for LDCQ comprises generating candidate latents z for a par- ticular state using the latent diffusion prior z p / s). These latents are then scored using the learnt Q-function and the best latent Zmar is decoded using the VAE autoregressive decoder (a s, Zmaz) ) to obtain H-length action sequences which are executed sequentially. Note that the latent diffusion model is used both during training the Q-function and during the final evaluation phase, ensuring that the sampled latents do not go out-of-support. The policy execution algorithm detailed in the Appendix 3.
4.3 LATENT DIFFUSION GOAL CONDITIONING (LDGC)
Diffuser (Janner et al. (2022)) proposed framing certain navigation problems as a sequence inpaint- ing task, where the last state of the diffused trajectory is set to be the goal during sampling. We

---

Published as a conference paper at ICLR 2024
can similarly condition our diffusion prior the goal to sample from feasible latents that lead to the goal. This prior is of the form Pz(z So. sy), where Sy is the target goal state. Since with la- tent diffusion, the training of the low-level policy alongside the VAE is done separately from the diffusion prior training, we can reuse the same VAE posterior to train different diffusion models, such as this goal-conditioned variant. At test time, we perform classifer- freee guidance to further push the sampling towards high-density goal-conditioned latents. For tasks which are suited to goal conditioning, this can be simpler to implement and achieves better performance than Q-learning Also, unlike Diffuser, our method does not need to have the goal within the planning horizon of the trajectory. This allows our method to be used for arbitrarily long-horizer tasks.
5 EXPERIMENTAL EVALUATION AND ANALYSIS
In our experiments, we focus on 1) studying the helpfulness temporal abstraction has in distinguish- ing latent skills (Section 5.1) 2) evaluating the ability of diffusion models to sample from the latent space (section 5.2 and 5.3) and 3) evaluating the performance of our method in the D4RL offline RL. benchmarks (section 5.4).
5.1 TEMPORAL ABSTRACTION INDUCES MULTI-MODALITY IN LATENT SPACE
In this section, we study how the horizon length H affects the latent space and provide empirical justification for learning long-horizon latent space representations. For our experiment, we consider the kitchen-mixed-v0) task from the D4RL. benchmark suite (Fu et al. (2020)). The goal in this task is to control a 9-DoF robotic arm to manipulate multiple objects (microwave, kettle, burner and switch) sequentially, in single episode to reach desired configuration, with only sparse 0-1 completion reward for every object that attains the target configuration. As Fu et al. (2020) states, there is high degree of multi-modality in this task arising from the demonstration trajectories because different trajectories in the dataset complete the tasks in a random order. Thus, before starting to solve any task, the policy implicitly needs to choose which task solve and then generate the actions to solve the task. Given state, the dataset can consist of multiple behavior modes, and averaging over these modes leads to suboptimal action sequences. Hence, being able to differentiate between these tasks is desirable.
We hypothesize that as we increase our sequence horizon H, we should see better separation be- tween the modes. In Figure 2, we plot 2D (PCA) projection of the VAE encoder latents of the starting state-action sequences in the kitchen-mixed dataset. With a lower horizon, these modes are difficult to isolate and the latents appear to be drawn from Normal distribution (Figure 2). How- ever, as we increase temporal abstraction from H = 1 to = 20, we can see three distinct modes emerge, which when cross-referenced with the dataset correspond to the three common tasks exe- cuted from the starting state by the behavioral policy (microwave, kettle, and burner). These modes better capture the underlying variation in an action sequence, and having picked one we can run our low-level policy to execute it. As demonstrated in our experiments, such temporal abstraction facilitates easier Q-stitching, with better asymptotic performance. However, in order to train these abstract Q functions, it is necessary to sample from the complex multi-modal distribution and the VAE conditional Gaussian prior no longer adequate for this purpose, as shown in section 5.2.
5.2 LDMS ADDRESS MULTI-MODALIT IN LATENT SPACE
In this section, we provide empirical evidence that latent diffusion models are superior in modelling multi-modal distributions as compared to VAEs. For our experiment, we again consider the kitchen- mixed-vO task. The goal of the generative model here is to learn the prior distribution p(z | s) and sample from it such that we can get candidate latents corresponding to state s belonging to the support of the dataset. However, as demonstrated earlier, the multi-modality in the latent spaces increases with the horizon. We visualize the latents from the initial states of all trajectories in the dataset in Figure 3a using PCA with H/ 20. The three clusters in the figure correspond to the latents of three different tasks namely microwave, kettle and burner. Similarly, we also visualize the latents predicted by the diffusion model Figure 3b) and the VAE conditional prior Figure 3c) for the same initial states by projecting them onto the principal components the ground truth latents. We can that the diffusion prior is able sample effectively all modes from the ground truth
7

---

Published as a conference paper at ICLR 2024
Kett
Microweve
Durnee
Figure 2: Projection of latents across horizon. Latent projections of trajectory snippets with different horizon lengths H. From the initial state there are tasks (Kettle, Microwave, Burner) which randomly selected the start of each episode. These primary modes emerge increase H, with the distribution turning multi-modal.
latent distribution, while the VAE prior spreads its mass over the three modes, and thus samples out of distribution in between the three modes. Using latents sampled from the VAE prior to learn the Q-function can thus lead to sampling from out of the support, resulting in extrapolation error.
(a) Ground truth
(b) Diffusion prior
(c) VAE prior
Figure 3: Visualization of latents projected using PCA for kitchen-mixed with H 20. The dif- fusion prior models the ground truth much more accurately while the VAE prior generates out-of- distribution samples.
5.3 PERFORMANCE IMPROVEMENT WITH TEMPORAL ABSTRACTION
We empirically demonstrate the importance of temporal abstraction and the performance im- provement with diffusion on modelling temporally abstract latent spaces. method with a variant of BCQ which uses temporal abstraction (H > 1), which we re- fer to as BCQ-H. We use the same VAE architecture here as LDCQ, and fit the condi- tional Gaussian prior with network having comparable parameters We find that generally, increasing the horizon H results bet- ter performance, both in BCQ-H and LDCQ, and both of them eventually saturate and degrade, possibly due to the limited de- coder capacity. With H 1, the latent distribution is roughly Normal as discussed earlier and our diffusion prior is essen- tially equivalent the Gaussian prior in BCQ, so we see sim- ilar performance. As we increase H, however, the diffusion prior able efficiently sample from the more complex latent distribution that emerges, which allows the resulting policies to benefit from temporal abstraction. BCQ-H, while also see- ing performance boost with increased temporal abstraction, lags behind LDCQ. We plot D4RL score-vs- for BCQ-H and LDCQ evaluated on the fitchen-mixeds task in Figure 4. The benefit temporal abstraction is ablated for different tasks in Appendix J.
compare our
diffusion model.
LDCQ
BCQU
LDCQ
3
Figure 4: D4RL score of LDCQ and BCQ-H on kitchen-mixed-V with varying sequence horizon H
8

---

Published as a conference paper at ICLR 2024
5.4 OFFLINE RL BENCHMARKS
In this section, we investigate the effectiveness of our Latent Difftusion Reinforcement Learning methods on the D4RL offline RL benchmark suite (Fu et al. (2020)). We compare with Behavior Cloning and several state-of-the-at offline RL methods. Diffuser (Janner et al. (2022)) and Decision Diffuser (Ajay al.) are prior raw trajectory diffusion methods. We found that our method does not require much hyperparameter tuning and only had to vary the sequence horizon H across tasks. In maze2d and AntMaze tasks we use H 30, in kitchen tasks we use H and adroit tasks use H 10. We train diffusion prior with T 200 diffusion steps. The other hyperparameters which are constant across tasks are provided in the supplemental material. In Table compare performance across tasks in the D4RL suite. We would like highlight tasks in Maze2d, AntMaze and FrankaKitchen environments which are known to be the most difficult in D4RL, with most algorithms performing poorly. Maze2d and AntMaze consist of undirected demonstrations controlling the agent to navigate to random locations in a maze. AntMaze is quite difficult because the agent must learn the high-level trajectory stitching task alongside low-level control of the ant robot with 8-DoF. In the maze navigation tasks, we also evaluate the performance of our coalicondi (LDGC) variant. For Diffuser runs we use the condiiontt inpainting version proposed by the authors since the classifier-guii guided version yielded poor results.
and locomotion
Table Performance comparison on D4RL tasks. LDGC evaluated in goal-directed maze tasks.
Dataset
BC
BCQ
CQI.
IQL.
DyT
Diffuser
DD
LDCQ (Ours)
LDGC (Ours)
maze2d-umaze-vl
3.8
12.8
5.7
47.4
27.3
113.5
134.2 = 4.0
141.0 = 2.7
maze2d-medium-vl
30.3
8.3
5.0
34,9
32.1
121.5
125.3 2.5
139.9 4.2
maze2d-large-vl
5.0
6.2
12.5
58.6
18.1
123.0
150.1 = 2.9
206.8 = 3.1
antmazz-umaze-divrree
45.6
55.0
84.0
62.2
54.0
81.4 +
85.6 :: 2.4
untuuze-medium--dinnes--z
0.0
53.7
70.0
0.0
45.5
24.6
68.9
75.6
umtmuze-lange-diveres
0.0
2.2
14,9
47,5
0.0
22.0
7.5
57.7 + 1.8
73.6 + 1,3
kicchen-complete-rO
65.0
$2.0
43.8
62.5
62.5 2.1
Sitchen-partial-
38.0
50.1
46.3
42.0
57.0
67.8 0.8
kitchen-mixxd-v0
51.5
34.5
52.4
51.0
50.7
65.0
62.3 0.5
halfcheetah- halfchectah-meddun-eeppe
$5.2
64.7
91.6
86.7
86,8
88.9
90.6
90.2 - 0.9
walker2d- madirr22-mmmiiun--xeer--22222
107.5
108.8
109.6
108.1
106.9
108.8
109,3 0.4
hopper-medium-ee
52.5
110.9
105.4
91.5
107.6
103.3
111.8
111.3 0.2
42.6
40.7
44.0
47.4
42.6
42.8
49.1
42,8 = 0.7
walker2d (a-erddd--em--- v2 medper-medium-22
75.3
53.1
72.5
78.3
74.0
79.6
82.5
69.4 + 3.5
$2.9
54.5
58.5
66.3
67.6
74.3
79.3
66.2 = 1.7
halfcheetah- medium-replay--v
36.6
38.2
45,5
44,2
36.6
37.7
39.3
41.8 + 0.4
w meddr---mmiium-eppaaa22222
26.0
15.0
77.2
73.9
66.6
70.6
75.0
68.5 4.3
hopper-medium meddimmmeeeuuu---eeeeeee-
18.1
33.1
95.0
94.7
82.7
93.6
100.0
86.2 1: 2.5
pen hum-human
34.4
68.9
37.5
71.5
74.1 + 2.7
hammer- human
1.2
0.3
4.4
1.4
1.5 0.8
door-human
0.5
9.9
4.3
11.8 1.9
relocate-human
0.0
0.2
0.1
0.3 0.1
pen-cloned
37.0
44.0
39.2
37.3
47.7 - 1.9
hammer-cloned
2.1
door-cloned
0.0
0.0
0.4
1.6
1.1 0.4
relocate-cloned
-0.3
-0.3
-0.1
-0.2
121.2 .024 0.1
pen expert
85.1
107.0
hammer-expert
125.6
107.2
86.7
45,8
door-expert
34.9
99.0
1015
105.1 : 0.3
relocate-expert
101.3
41.6
95.0
104.7 1.4
carla-lane-v0
17.2
-0.1
20.9
18.6
24.7 = 3.2
Both LDCQ and LDGC achieve state-of-the-arr results all sparse reward D4RL tasks. The goal- conditioned variant outperforms all others in maze2d and AntMaze. This variant is extremely simple to implement through supervised learning of the diffusion prior with no Q-learning or online plan- ning and is ideal for goal-reaching tasks. We also provide evaluation of our method on the D4RL locomotion suite and the Adroit robotics suite. While these tasks are not specifically focused on trajectory-stitching our method is competitive in the locomotion tasks and stronger than baselines in Adroit. To extend our method to Carla's image input spaces, we compress the image using a 3- 3-VAE encoder such that our method operates compressed state space (more in Appendix G).

---

Published as a conference paper at ICLR 2024
6 CONCLUSION
In this work, we showed that offline RL datasets comprised of suboptimal demonstrations have expressive multi-modal latent spaces which can be captured with temporal abstraction and is well suited for learning high-reward policies. With a powerful conditional generative model to capture the richness of this latent space, we demonstrated that the simple aatch-constrained Q-learning frame- work can be directly used to obtain strong performance. Our biggest improvements come from long-horizon sparse reward tasks, which most prior offline RL methods struggled with, even previ- raw trajectory diffusion methods. Our approach also required task-specific tuning, except for the sequence horizon H. We believe that latent diffusion has enormous potential in offline RL and our work has barely scratched the surface of possibilities.
REPRODUCIBILITY
We provide link to our code in section A.I. We provide details of our model architectures in section A.2 and hyperparameters in A.3. Our experiments were conducted on the open D4RL benchmark datasets.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making In NeurlPS 2022 Founda- tion Models for Decision Making Workshop.
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. Opal: Offline prim- itive discovery for accelerating offline reinforcement learning. In International Conference on Learming Representations, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, P. Abbeel, A. Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Neural Information Processing Systems, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-
critic methods. In International conference machine learning, pp. 1587-1596. PMLR, 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp. 2052-2062. PMLR, 2019.
Seyed Kamyar Seyed Ghasemipour, Dale Schuurmans, and Shixiang Shane Gu. Emaq: Expected- max q-learning operator for simple yet effective offline and online rl. In International Conference on Machine Learning. pp. 3682-3691. PMLR, 2021.
Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy. arXiv preprint arXiv:2303.09556, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Irina Higgins, Loic Matthey, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representa- tions, 2016.
Jonathan Ho and Tim Salimans. Classifier-fre free diffusion guidance. In NewrIPS 2021 Workshop on Deep Generative Models and Downstream Applications.
10

---

Published as a conference paper at ICLR 2024
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 2020.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 14:1222-, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pp. 9902-9915. PMLR, 2022.
Heewoo Jun and Alex Nichol. Shap-e: Generating conditional 3d implicit functions. arXiv preprint arXiv:2305.02463, 2023.
Diederik P Kingma and Max Welling, Auto-encoding variational (Bayes}. In Int. Conf. on Learning Representations.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference Learning Representations.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In International Conference on Learning Representations.
Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-conditioned policies. arXiv preprint arXiv:1912.13465, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 13:1179-1191, 2020.
Sergey Levine, Aviral Kumar, G. Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. ArXiv, 0bss/2005001644 2020.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11461-11471, 2022.
Marlos Machado, André Barreto, Doina Precup, and Michael Bowling. Temporal abstraction in reinforcement learning with the successor representation. Journal Machine Learning Research, 24(80):1-69, 2023.
Timothy Mann and Shie Mannor. Scaling up approximate value iteration with options: Better poli- cies with fewer iterations. In International conference on machine learning, pp. 127-135, PMLR, 2014.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Ser- gio Valcarcel Maeua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin, Imi- tating human behaviour with diffusion models, 2023.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. Jens Kober, Fabio Ramos, and Claire Tomlin (cds.), Proceedings of the 2020 Conference Robot Learning, volume 155 of Proceedings of Machine Learning Research, pp- 188-204. PMLR, 16-18 Nov 2021. URL https: http:////ooceeee......c..o .mw.mlr........mm.m............m prwss/eeesss. pertsch2la. html.
/
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEF/CVF Con- ference Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. volume 9351, pp. 234-241, 10 2015. ISBN 978-3-719-24573-7. doi: 10.
28.

---

Published as a conference paper at ICLR 2024
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Informa- tion Processing Systems, 35:36479-36494, 2022a.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement, IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022b.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay, arXiv preprint 1tt:::1511..99900 2015.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei (eds.), Pro- ceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2256-2265, Lille, France, 07-09 Jul 2015. PMLR. URL ht httt:::/prrceeeee
www.ppess...//l/. (atppp/ss//S//sslllcicctttttetie wwwhsdiicsteeinnnii..e.. htmm....
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski, and Jürgen Schmidhu- ber. Training agents using upside-down reinforcement learning. arXiv preprint arx.v::9/2.00880 2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention you need. Advances in neural informa- fion processing systems, 30, 2017.
Nino Vieillard, Marcin Andrychowicz, Anton Raichuk, Olivier Pietquin, and Matthieu Geist. Im- plicitly regularized with implicit valuus.. Gustau Camps- Valls, Francisco R. Ruiz, and Is- abel Valera (eds.), Proceedings of The 25th International Conference Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 1380-1402. PMIR, 28-30 Mar 2022. URL https: http:///pocceeeig..e///. wtw.://s..................meee wttpppesswwwl..vieeeeeeee..eeeee html.
Yiqin Yang, Hao Hu, Wenzhe Li, Siyuan Li, Jun Yang. Qianchuan Zhao, and Chongjie Zhang. Flow to control: Offline reinforcement learning with lossless primitive discovery. Proceedings of the Thirty-Seventh AAAI Conference Artificial Intelligence and Thirty-Fifth Conference Inno- vative Applications of Artificial Intelligence and Thirteenth Symposium on Educational Advances in Artificial Intelligence, 880-0. doi: 20.16099aaai.vv371..0000 26286. URL https: /oii /ttp//...rr......0////aaa//////. aaai. v3719. 26286.
AAAI Press, 2023, ISBN 978-1-57735-
Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforce- ment learning. In Conference on Robot Learning, 2020.
12

---

Published as a conference paper at ICLR 2024
A TRAINING DETAILS
A.1 SOURCE CODE
The source code is available at:
wtt.:lu....mmm/////////eeee ldcq.
A.2
NETWORK ARCHITECTURE
A.2.1 VARIATIONAL AUTOENCODER
Encoder. For learning the latent trajectory representation, our VAE uses an architecture similar to Ajay et al. (2020). The encoder consists of two stacked bidirectional GRU layers, followed by mean and standard deviation heads which are each a 2 layer MLP with RELU activation for the hidden layers. The mean output head is linear layer. The standard deviation output head is followed by SoftPlus activation function to ensure it is always positive. The hidden layer dimension fixed to 256.
Decoder. For the low-level policy decoder, we use an auto-regressive policy network similar to that described in EMAQ (Ghasemipour al. (2021)), in which each element of the action vector has its own MLP network, taking as input the current state, latent representation, and all previously- sampled action elements. The complete action vector sampled element-by-elemen with the most recently sampled element becoming an input to the network for the next element. These MLP networks consists of layers followed by 2 layers of mean and standard deviation heads similar to the encoder network. The mean output head is a linear layer and the standard deviation output head followed by a SoftPlus activation. Again, ReLU activation is used after all hidden layer and the hidden dimension is fixed to 256
A.2.2 DIFFUSION PRIOR
The diffusion prior is deep ResNet (He et al. (2016)) architecture consisting of 8 residual blocks. It takes as input vector representing a latent trajectory and outputs denoised version of the latent. The hidden blocks of dimensions: [128, 32, 16, 8, 16, 32, 128]. Similar to U-Net (Ronneberger et al. (2015)), the initial blocks are connected by residual connections to the later blocks having the same hidden dimension. The diffusion timestep t is encoded with 256-dimensional sinusoidal embedding and then further encoded with a 2-layer ML.P The conditioning state s is also encoded by layer MLP. In each residual block, the state and time encodings are concatenated with the current layer activation for conditioning. When training the unconditional diffusion model for classifier-free guidance, the state input given vector of zeros represent null vector.
A.2.3 Q-NETWORKS
The Q-networks take as input the state s, latent z and consist of layer MLP in the first 3 layers, 32 hidden units in the third layer, and finally linear output layer. We use GELU activation function between hidden layers. LayerNorm is applied before each activation.
with 256 hidden units
A.3 HYPERPARAMETERS
The hyperparameters which constant across tasks for the different stages of our proposed method are given in Tables 2, 3 and 4.
A.4 HARDWARE
The models were trained on NVIDIA RTX A6000. Since different tasks have different dataset sizes, the model training times changes across tasks. Depending the task. training the 3-VAE took between 3-7 hours, the diffusion prior took between 4-12 hours and the Q-Learning took between 3-5 hours.
13

---

Published as a conference paper at ICLR 2024
Table 2: B-VAE hyperparameters
Parameter
Value
Learning rate
5e-5
Batch size
128
Epochs
100
Latent dimension (z)
16
B
0.05
Hidden layer dimension
256
Table 3: Diffusion training hyperparameters
Parameter
Value
Learning rate
le-4
Batch size
32
Epochs
300
Diffusion steps (I)
500
Drop probability (For unconditional prior)
0.1
Variance schedule
linear
Sampling algorithm
DDPM
y (For Min-SNR-y weighing)
5
Table 4: Q-Lenniing hyperparameters
Parameter
Value
Learning rate
5c-4
Batch size
128
Discount factor (7)
0.995
Target update (p)
0.995
PER buffer a
0.7
PER buffer 3
Linearly increased from 0.3 to 1, Grows by 0,03 every 3000 steps
Diffusion samples for batch argmax
500
14

---

Published as a conference paper at ICLR 2024
B Low LEVEL POLICY AND SKILL PRIOR TRAINING
We describe the training of the 3- B-V for learning the skill representations and the low level policy, and then training the diffusion skill prior in Algorithm 2
Algorithm 2 LL Policy and Skill Prior Training
1: Input: horizon H, KL regularization coefficient 3, number of steps M, diffusion timesteps T', for Min-SNR-y, sequence encoder qa, conditional Gaussian prior Pws policy decoder To. latent diffusion prior j.g, offline dataset D, variance schedule 01,...,,TT: 01,.... Oep, $1,..., .r.
2: for iler 1 to M do
D. Training 8- 3-VAE
3: Sample size minibatch of H length subtrajectories (1m)) (1:n) H. 1144 from D
4: Sample z' ~ 46 (2)s b++
(i)
zz()) + 6D (g (1l 5: C( k,w E Σ log To Ta (9) 2(0)
6: Take gradient step to minimize L
7: end for
8: for iter = 1 to M do
D Training diffusion skill prior (1m) (1:n)
9: Sample size minibatch of length subtrajectories {5aa a14 from
10: Sample ( 40 (2|s (2/
11: Sample diffusion time (() [[
Noise latents z N(va,z (1-
12:
13:
((v)= =
min (NR((())) 7 2
us (2) Sc. 11?)
14:
Take gradient step to minimize L
15: end for
C POLICY EXECUTION
After training the diffusion prior and Q-learning, we execute the policy as described in Algorithm 3.
Algorithm Policy Execution
1: Input: horizon H, number of latents to sample n,, Q-function Qe. policy decoder no. latent diffusion denoising function pv. variance schedule 01, QT, āi..... or, BI...., .r....rr.
2: done = False
3: while not done do
4: Observe environment state
5: Sample n latents: Zr (0.
6: for \ T to do
DDPM Sampling
7:
ua (2e (,,, + w (p (Ze, Sp, My (2,,0,1))
21-1 N( (1-(1- 2,
vai-Ta 1- 2, I[t 1)3,I)
8:
9:
end for
10: Find best skill by scoring them with Q-function: argmax Q0(5).2))
11:
h=0
12:
while h < Hand not done do Observe environment state Ss.
Execute Skill z
14:
Get action a, E Ta (a (Sn, 2
Execute action as
16:
Update done +
18:
end while
19:
end while
15

---

Published as a conference paper at ICLR 2024
D COMPARISON WITH OTHER LATENT SKILL LEARNING METHODS
In this section, we compare our method against existing algorithms which do latent skill learning with generative models in the D4RL tasks. Specifically, compare against the VAE based algo- rithms OPAL (Ajay et al. (2020)) and PLAS (Zhou et al. (2020)), and the Normalizing Flow based algorithm Flow Control (Yang al. (2023)). We were not able to replicate the OPAL results with the code provided to us by the authors. We list the best scores we were able to obtain with it. We did find code implement Flow to Control, and so use the results listed in their paper. The results are shown in table 5.
Table Performance comparison on D4RL tasks. Algorithms are evaluated only on the tasks listed in their respective papers.
Dataset
OPAL
PLAS
Flow2Control
LDCQ (Ours)
LDGC (Ours)
maze2d-lagge-vl
122.7
150.1 2.9
206.8 +: 3.1
antmaze-man memnzzz-mmdiumddiveeseee.
57.5
0.0
83.7
68.9 + 0.7
75.6 0,9
uatmaze-large-divess-->>
52.0
0.0
52.8
57.7 + 1.8
73.6 + 1.3
kitchen-parrial-v0
55.5
43.9
74.9
67,8 + 0.8
kitchen-mixed-v0)
50.2
40.8
69.2
62.3 + 0.5
halfcheetah-mmddin-ccpp----22
99.3
90.2 0.9
wallerr2d-mcciium-exprr-t-
97.2
109.3 + 0.4
nopper-mediuum--pperr-
111.0
111.3 0.2
sulfcheetta-mmeiun--ZZ
42.2
42.8 1 0.7
wralke22d-mediung
66.9
69.4 1 3.5
hopper-medium--2
36.9
66.2 1 1.7
halfichcttah-mmciiuu--pplaaaa walkeer2d-mcdiium--ppaayy
45.7
41.8 + 0.4
14.3
68.5 + 4.3
hoppe--mddium-rrplayyy
51.9
86.2 + 2.5
pen-human
67.3
63.1
74.1 +: 2.7
hammer-human
4.6
3.3
1.5 + 0.8
door-human
15.1
+ 1.9
relocate-humee
0.3
0.3 0.1
pen-cloned
49.0
65.8
47.7 - 1.9
hammer-cloned
2.1
2.8 1.2
door-cloned
3.3
8.1
1.1
relocate-cloned
-0.2
-022 0.1
E LATENT DIFFUSION-CONSTRAINEE PLANNING (LDCP)
In this section, we explore another method to derive a policy for offline RL with latent diffusion other than our proposed methods Latent Diffusion-Constrainnee Constrained Q-Learning (LDCQ) and Lateni Diffusion Goal Conditioning (LDGC). This is model-based method which learns a temporally abstract world model of the environment from offline data. Specifically, we learn temporally abstract world model P, [s,[5t S;, z) that predicts the state outcome of executing a particular latent behavior after H steps. That is, given the current state s, and latent behavior z the model predicts the distribution of the state S This is trained in supervised manner by sampling transition tuples (st, S from TH D and minimizing the objective:
[(n) Eyg-D (S Sc, S 1|
where 17 the parameters of the temporally abstract world model Py-
In goal-reaching environments, we leverage this model to do planning using the diffusion prior. We sample latents z' (1 < n) using the diffusion prior for the current state s,, and use the learnt dynamics model to compute predicted future state s148 for each latent z'. These final states are then scored using cost-function J and the latent corresponding to the best final state is chosen for
16

---

Published as a conference paper at ICLR 2024
execution. Note that sampling latents from the diffusion prior ensures that the world model is not queried on out-of-support data. We refer to this method as Latent Diffusioo Diffusion-Consrrainn Constrained Planning (LDCP). The planning procedure is summarized in Algorithm 4.
Algorithm 4 Latent Diffusion- CrffussonnCCossraaineee Planning (LDCP)
1: Input: horizon H, number of latents sample cost-function J, policy decoder πe. tem- porally abstract world model Py. latent diffusion denoising function Hug variance schedule Ul.....:0 oh,,...,07 OT,, 31, -1.-.., 3r.
2: done = False
while not done do
4: Observe environment state 50
5: Sample n latents: Zr 10.
6: for t T to 1 do
D DDPM Sampling
7:
Hu (z, wp (ze, Su. My (z),
- N( Var(l ( 2, va1-18 2, 1(t 1)111)
8:
9:
end for
10: Compute future states for each latent z): si = ( so. z)) Find best skill based on the cost-function: i = argmin [(}]]]
12:
h h=0 0
13:
while h < H/nd not done do Observe environment state 8/:
Execute Skill zb
14:
15:
Get action a E Ty (a|s), 26) Execute action ax
16:
17:
Update done h +1
18:
19:
end while
end while
The cost-function which we use for the goal-reaching environments is the Euclidean distance to the goal. We can also extend this planning to horizons greater than И by further sampling latents for each future state H (1 < < n1). This means, after sampling n latents for s; with the diffusion prior, we further sample more latents for each of the future states $ This increases the 'plan- ning depth' The final states at the last level of planning are then scored using the cost-function and the latent at the first level which led to that final state is chosen for exccution. This procedure complexity grows exponentially and thus the planning depth has to be restricted. For a planning depth of d, there nd model calls. We found planning- depth of navigation environments achieving state-of-the-at results. Thus, with latent sequence horizon of H 30, our total planning horizon is 60. The results are tabulated in Table 6.
2 to be sufficient for all
Table Performance comparison on D4RL navigation tasks with LDCP.
Dataset
BC
BCQ
CQL
IQL.
DT
DifTuser
DD
LDCQ (Ours)
LDGC (Ours)
LDCP (Ours)
maxx2d-largo-v1
5.0
12.5
58.6
18.1
123.0
150.1 2.9
206.8
184.3
antmaae medium diverse v2
0.0
0.0
53.7
70.0
0.0
45.5
24.6
68.9 1 0.7
75.6 1 0.9
77.0 1 1.1
canmaazellage-diveeee..
1.8
73.6
1.3
E VISUALIZING MODEL. PREDICTIONS
Learning world model also allows us to visualize the effect of executing any given latent behavior. This means, when the model is not used for planning, like in LDCQ, it can be used to compute the final state that will be reached for every latent behavior from a particular state. This information can be used understand if the model is learning reasonable behavior modes.
We plot the xy-coordinates of our abstract world model pa ( St. z) predictions at a T- intersection of AntMaze large environment for latents sampled from diffusion prior z (z s,) in Figure 5 to demonstrate this. The plot shows that the diffusion prior sampled latents which go in all the three directions the T-intersection.
17

---

Published as a conference paper at ICLR 2024
Current position
23
18
14
Figure 5: Visualizing model predictions: Visualization of future states with latents sampled from the diffusion prior T-intersection in antmaze-large-divrree D4RL task. We can see multi- modal future state predidctions corresponding to possible directions at the T-intersection.
F D4RL TASKS
(a) Ant-Maze
(b) Franka-Kitchen
(c) Adroit
Figure 6: environeets
The tasks of AntMaze and Franka Kitchen require long horizon credit assigmment and stitching trajectories in the behavior dataset. We highlight consistently strong performance in these tasks as the primary empirical result our method, since most other baselines perform poorly here, even other diffusion based methods. The Adroit suite also consists of narrow, but precision oriented manipulation tasks with sparse rewards. We also show very strong performance here. This task does require trajectory stitching, but needs filtering of low reward demonstrations while sticking close to the behavior support.
The locomotion datasets are collected from SAC agents of varying performance. Our method has average performance on the locomotion suite while being much stronger in the other tasks. We sus- pect the high periodicity of the walking gaits in the locomotion suite does not benefit much from reasoning with temporal abstraction. We also do not use perturbation function during Q-learning like Fujimoto et al. (2019), which makes difficult for us to improve over the poor controllers in medium and medium-replay locomotion datasets. Introducing perturbation function requires careful tuning to avoid extrapolation error, and the converged Q-learning wouldn't necessarily cor- respond to a high value policy, which is why other offline RL methods, which try balance this tradeoff, evaluate online during training and consider the best scores. We however only evaluate a policy once after training is fully complete.
18

---

Published as a conference paper at ICLR 2024
G CARLA AUTONOMOUS DRIVING TASK
To extend our method for tasks with nigh-dimensional image input spaces, we propose to compress the image space such that method operates this compressed state space. We create low- dimensional compressed representation using an autoencoder & before using the LDCQ framework. Note that this encoder operates single image and not on temporal sequence of images (Figure 7). The downstream framework of LDCQ however operates on the temporal compressed image sequences.
Reconstruction
9x(s/o)
s
1,(o'(s)
Figure 7: Autoencoder training for mmage-basedd task
We evaluate the performance of our method on the CARLA autonomous driving D4RL task. The task consists of an agent which has control the throttle (gas pedal), the steering, and the break pedal for the car. It receives 48 x 48 RGB images from the driver's 's perspective as observations. We use B- B-VAE architecture to create 32-dimensional compressed state for this task. The horizon for LDCQ is set to H 30.
H RANDOM WALK 1D
In this experiment, we construct simple toy problem to show how sampling effectively from the multimodal behavioral distribution helps the diffusion prior outperform Gaussian VAE prior during Q-learning. We construct simple toy MDP with one-dimensional state space S E [-10.0, 10.0). The agent starts the origin (0,0) and receives a reward of 10 if reaches either the far left (-10.0) or far right (10.0) state, and -1 reward every other step. The environment times out after 500 steps. The action is the distance moved in that timestep with a max distance length I, A [-1.0, 1.0). The dataset consists of episodes where the agent randomly selects actions from the uniform distribution a ~ (([-1.0, -088 J [0.8, 1l). This means the agent has a step size between 0.8 to 1.0 units either left right every timestep. We train VAE to try fit this action distribution, and use BCQ to learn a policy. We also train train a diffusion based policy with LDCQ, using results.
and compare the
Figure 8: 1D Random walk
We find that the VAE frequently samples actions never present in the dataset. This is because the Gaussian mean to the above action distribution is 0.0, but no actual actions are present between (-0.8, 0.8) where large proportion of probability mass is assigned by the Gaussian model. Mean- while, the diffusion prior is able to the 2 modes quite well. After 10000 steps of Q-learning, the diffusion constrained policy learns to navigate to either end perfectly and achieves an average reward of -2.2 while the VAE constrained policy still almost random, frequently taking actions with small step size and average reward of -66.
19

---

Published as a conference paper at ICLR 2024
I INCREASING DIFFUSION STEPS IMPROVES PERFORMANCE
We study the impact of the number of diffusion steps on the performance for LDCQ. We found that for the locomotion tasks, increase in diffusion timesteps T during evaluation generally corresponds to increase task performance. We plot these results in Figure
Half Cheetah Medium Expert
Hopper Medium Expert
90
88.
5d
1.5
110
104.6
98
8
D4RL Score
DARL Scure
100
78.89
90
90
B4.71
70
50
70
61.45
50
50
400
3001
Difiusion Steps
Diffision Steps
Walker2d Medium Expert
115
110
108.80
80
109
23
109 22
104.43
DARL Score
105
100
85
100
300
400
500
Diffusion Steps
Figure 9: D4RL score for LDCQ with varying diffusion steps T in locomotion tasks.
For the long horizon tasks, we found that increasing diffusion steps resulted in an initial trend upward in performance. Beyond this, the performance does not improve with additional diffusion steps (Figure 10).
Antmaze Large
70
57.7
57.5
D4RL Score
50
40.2
40
30
26.3
20
100
200
300
400
500
Diffusion Steps
Figure 10: D4RL score for LDCQ with varying diffusion steps T
We also used additional diffusion steps at time 0 similar to Diffusion-X (Pearce et al. (2023))
This means that after the DDPM sampling of diffusion from time T to 1, we run X additional diffusion steps to further denoise the sample, assuming time-step 1. Pearce et al. (2023) reasoned that this pushes the samples further towards higher-likelihov regions. We used 10 additional steps across experiments and found this to slightly improve performance.
20

---

Published as a conference paper at ICLR 2024
J PERFORMANCE IMPROVEMENT WITH TEMPORAL ABSTRACTION
We provided empirical evidence for improvement in performance as we increase temporal abstrac- tion horizon Il for the kitchen-mixed-vD environment. We see similar trends for the other long- horizon tasks as well (Figure 11). The performance in general improves with increasing temporal abstraction but beyond a certain point, it drops possibly because of the limited capacity of the policy decoder.
For the locomotion tasks, we did not observe any noticeable difference with increase in temporal abstraction, so we ended up using moderate sequence length H 10. This could be due to the high frequency periodicity of these tasks that does not require much look-ahead.
Kitchen Mixed
Kitchen Partial
60
D4RL Score
D4RL Sture
20
20
10
20
10
Antmaze Large
Antmaze Medium
D4RL Soure
40
40
D4RL,
20
0
40
10
20
30
Horizon
Horizon
Figure 11: D4RL score for LDCQ with varying sequence horizons H.
K Q LEARNING TD ERRORS
We found Q learning with these skill latents to be very stable. This, alongside strong empirical performance in these tasks indicates that the skill latents learnt with the 3-VAE contain useful information, and that the function is able extract this information easily. We show the training graphs for antmaze-large. kitchen-partin Each X-axis tick corresponds 3000 updates from minibatches sampled from the PER buffer. The Y-axis is the average TD error for those 3000 minibatches.
halfcheetah halfchectab-mecium.. and Adroit door-human below.
21

---

Published as a conference paper at ICLR 2024
0.50
0.00
(a) antmaze-large
(b) kitchen-parte
(c) srficheetah-mediine
(d) door-human
Figure 12: Q learning curves for different tasks. X axis is number of epochs of size 3000 from PER buffer, Y axis is average TD error over last 3000 minibatches.
22


Reviews:
[{"review_id": "cneBghShTP", "reviewer": "Reviewer_Agic", "summary": "This paper proposes a novel offline RL algorithm that leverages a diffusion model to plan over the learned temporally abstract latent space for action representation. The empirical results show that temporal abstraction can help distinguish latent skills and the proposed method shares competitive performance with other state-of-the-art baselines.", "questions": "There are some questions and concerns, which I have outlined in the previous section.\n\nUPDATE: I updated my score from 5 to 6 as the authors addressed most of my concerns in the response.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Exploring how to leverage expressive models such as diffusion models and Transformers for policy learning is an important direction in RL, especially for the offline setting. To the best of my knowledge, the idea of combining high-level diffusion planning and low-level primitive learning is novel. The paper is well-organized and clearly written.", "weaknesses": "While well-motivated, I have some questions about this work:\n\n1. My major concern is that the results on offline RL benchmarks may be insufficient to show the advantage of planning with diffusion models on latent action space. While I appreciate that the authors have covered the most popular state-of-the-art baselines in Table 1, I think it is necessary to compare LDCQ with some literature that similarly performs planning on the learned action representation space learned by VAE [1, 2, 3, 4] or Flow [5, 6]. Also, I am curious about the ablation of LDCQ on diffusion models (Section 5.2 has shown some support, however, it compares the diffusion prior with the VAE conditional prior rather than a learned latent policy).\n2. The paper has discussed some limitations of prior works depending on VAE representations. However, it is unclear to me why introducing latent diffusion can solve these limitations, as LDCQ also requires a learned VAE embedding. Is it because the diffusion model is more expressive so you can use less KL penalty or there are some other reasons? And if we learn the latent action space with a flow model, which is a lossless representation, does latent diffusion planning still have some advantages?\n\nOverall, while this paper presents some interesting ideas, I am unable to recommend acceptance at this stage given the questions mentioned above. However, I would be happy to raise my score if the authors could address my concerns.\n\n[1] Pertsch, Karl, Youngwoon Lee, and Joseph Lim. \"Accelerating reinforcement learning with learned skill priors.\" Conference on robot learning. PMLR, 2021.\n\n[2] Ajay, Anurag, et al. \"Opal: Offline primitive discovery for accelerating offline reinforcement learning.\" arXiv preprint arXiv:2010.13611 (2020).\n\n[3] Zhou, Wenxuan, Sujay Bajracharya, and David Held. \"Plas: Latent action space for offline reinforcement learning.\"  Conference on Robot Learning. PMLR, 2021.\n\n[4] Chen, Xi, et al. \"Lapo: Latent-variable advantage-weighted policy optimization for offline reinforcement learning.\" Advances in Neural Information Processing Systems 35 (2022): 36902-36913.\n\n[5] Singh, Avi, et al. \"Parrot: Data-driven behavioral priors for reinforcement learning.\" arXiv preprint arXiv:2011.10024 (2020).\n\n[6] Yang, Yiqin, et al. \"Flow to control: Offline reinforcement learning with lossless primitive discovery.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 9. 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "43Gdsj7ld2", "reviewer": "Reviewer_hYz3", "summary": "# Summary\n\nThis work brings together the ideas from Batch Constrained Q-Learning (BCQ) and Latent Diffusion Models (LDMs), in order to succeed in offline RL. As in LDMs, the proposed method (called LDCQ - Latent Diffusion Constrained Q-learning) has a two stage training process. (i) The first stage trains a variational autoencoder (VAE) for trajectories (sequences of state-actions), where the decoder is the policy network. This latent space is a compressed representation of trajectories, serving as a form of temporal abstraction. (ii) The second stage of the training learns an LDM on this latent space and a latent-conditioned Q-network. At inference time, multiple latents are sampled from the LDM, and the Q-function is used to score each latent based on estimated Q-value. This scoring mechanism can be replaced with any kind of goal-conditioning (e.g.: L2 distance to goal in a maze environment), to obtain a goal-conditioned variant of the proposed method.\n\nThe use of latent diffusion makes this work distinct from recent diffusion-based offline RL methods such as Diffuser and Decision Diffuser. The analysis presented in this work highlights that (i) diffusion models are superior to VAEs when dealing with multi-modal latent spaces, (ii) the proposed LDCQ method achieves state-of-the-art or similar performance when compared with prior works.", "questions": "# Questions\n\nSome of my questions have been merged into the weaknesses section. Here are the remaining questions.\n\n- The policy execution first samples a latent z-i (representing a skill), and then fixes that z sample throughout the policy execution for H steps. What would happen if z-i is re-sampled after every action is taken by the policy (like model predictive control - MPC). Do you expect it to improve performance and if so, will it be much more computationally expensive given the cost of sampling z?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "# Strengths\n\nThe paper proposes a novel idea: Combining BCQ and latent diffusion models can lead to an offline RL algorithm that has the best of both worlds. Overall, this idea and the execution and analysis presented in this work is significant and of high quality.\n\n1. The proposed method aims to fill a gap in existing methods such as Diffuser and Decision Diffuser by taking inspiration from latent diffusion models (LDMs): shifting diffusion into the latent space and separating the training process into policy and trajectory auto-encoding on one side and Q-learning and diffusion on the other side. With this choice, **the method inherits the strengths of both LDMs and batch-constrained Q-learning (BCQ)**.\n\n1. The analysis of how the diffusion model of the proposed method (LDCQ) deviates from a baseline (that uses VAEs instead), clearly supports the hypothesis that **diffusion models are better at dealing with multi-modality** in the distribution of latents.\n\n1. The paper has shown **similar or state-of-the-art performance** on a challenging selection of environments from the D4RL benchmark, the most notable of which is the performance on Franka kitchen-partial-v0 and kitchen-mixed-v0 environments, and the Ant-Maze large-diverse and medium-diverse environments. Specifically, two of these environments -- antmaze-large-diverse-v2 and kitchen-partial-v0 show significant improvements in the state-of-the-art.\n\n1. The source code was provided.\n\n**Clarity**: The paper has some issues with ordering of the introduced concepts and terminology, but by the end of the paper, the reader is left with a clear understanding of the presented work. The figures and algorithms are clear and greatly help in the understanding of the proposed method and analysis.", "weaknesses": "# Weaknesses\n\nOverall, the paper has some statements/claims that are stretched, the empirical evidence seems like a mixed bag with certain D4RL environments/tasks silently omitted, and some weaknesses inherited from choosing latent diffusion models.\n\n## Stretched claims\n1. The phrases “improves credit assignment”, “faster reward propagation” describing the proposed work should be avoided, or backed by empirical evidence. I don’t see how either of these quantities can be measured. I understand how the proposed method will intuitively achieve both of these -- however, they are just intuitions, not proven.\n\n1. The paper describes the decoder or policy as a “powerful decoder” that can handle “high-frequency details”. This is true for LDMs used for computer vision dealing with high resolution images. However, I don’t see how this translates well into RL. Training the VAE encoder-decoder model in a separate first stage seems like it will produce a rigid, fixed policy (i.e. decoder). The paper itself acknowledges that the “limited decoder capacity” is probably what is leading to the saturation of performance in Figure 4. How can the decoder be “powerful” and “limited” at the same time? Is there a way to empirically measure the strength of a decoder?\n\n## Issues with empirical evidence\n\n1. In Table 1, not all environments/tasks from D4RL are included. For example, the medium and umaze versions of Maze-2D are missing. In my experience, despite being easier environments, performance on these environments do highlight subtle differences between algorithms. They are certainly more important than the latter category of environments that do not require trajectory stitching.\n\n1. In the latter category of environments that do not require trajectory stitching, more environments/tasks seem to have been silently omitted. The “-expert” suffix versions of Adroit pen, hammer and relocate tasks have been omitted. I would like to see an explanation of why certain environments in Table 1 were omitted, as otherwise it appears as cherry picking. Note that these results were not present in the Appendix either.\n\n1. Some qualitative results such as visualization of trajectories at test time would have been helpful for supporting the trajectory stitching property.\n\n1. Confidence intervals (error bars) are missing in Figure 4 and Adroit environment/tasks of Table 1.\n\n1. Not all reported performance values are state-of-the-art, this claim should be qualified.\n\n## Clarity\n\n1. Batch-constrained Q (BCQ) is mentioned several times and the acronym is used once before it is explained in Section 4.2. I recommend briefly explaining BCQ as a part of “Section 3.2 Offline Reinforcement Learning”. \n\n1. Algorithms 1 and 2 explain the second stage of training and the policy execution phase respectively. However, there is no algorithm box for the first stage of training (VAE). The policy \\pi also has a subscript \\theta, but \\theta is not used anywhere else in the paper. For completeness, I would recommend adding an Algorithm to describe the first stage of training where the policy parameters \\theta are learned.\n\n-----------------\n## Updated score after author response\n\nThe authors have addressed the major concerns raised above and I have increased my rating from 6 to 8.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p6sMiqNKFk", "reviewer": "Reviewer_LeBC", "summary": "This paper proposes using latent diffusion in offline RL as a way to compress task specific skills and information that helps learning Q-functions for improved performance in offline benchmarks. This paper provides a good instantiation for how to use latent diffusion models (LDM) in the offline RL setting. The key is to use a separate diffusion model, from which to learn a latent conditioned policy (similar to prior works using latent skills for RL). The proposed approach therefore seems quite useful for long horizon credit assignment tasks, where temporal abstraction and skills play a key role, and the work tries to justifies this through sparse reward environments for experimental evaluation.", "questions": "See my questions along with the weakness section. Primarily, my worries are : \n\n1. How does the representations learnt with beta-vae compare with other representation objectives for learning the latents? How does diffusion compare to that? \n\n2. How can we evaluate the quality of the latents? This is of big importance for this work, since the Q function is now learnt over the latents; I would like to see quantitivate results showing the effectiveness of the Q-function, maybe TD errors too, and qualitative results showing how good are the latents. Figure 2 only shows that latents are diverse and more skills are introduced across the horizons; but I worry about the use of Beta-VAE in the algorithm, and think better representation objectives should instead be used in the algorithm pipeline. Results justifying this would strengthen the paper.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper proposes an interesting idea that allows decoupling the diffusion model from the policy decoder; allowed the algorithm to be used for both continuous or discrete action environments. Experiment results are primarily for continuous action D4RL offline tasks, and shows relatively improved performance across few benchmarks. The execution of the algorithm is interesting, but I worry about the easy-ness of the approach. I agree with the authors that the reason that the algorithm is structured this way is to have better flexibility in terms of the latent skills from the diffusion prior, but see my comments in weaknesses. Other than that, the execution of the algorithm in the BCQ setting taking everything in the latent space (assuming the latents are good) is an interesting idea. Figure 2 shows how well the latents also get distributed across the horizons, suggesting that diverse skills can be learnt across horizons.", "weaknesses": "1. Experiment results are difficult to follow. The D4RL results are the primary ones, but the appendix claims to have results for CARLA and Goal Conditioning tasks too? It seems the goal conditioned tasks are not the standard ones in GCRL literature, and it is not clear what the key takeaway is other than the constrained offline RL results + qualitative results evaluating the latents across the horizon. \n\n2. The proposed algorithm is interesting, but might have difficulty with the execution and easy-ness of the approach. Primarily because the Q-function is now learnt over Q(s,z) latents, and it would be helpful to do some qualitative evaluation of the Q-values over the latents. This is because the learnt Q would now heavily depend on the quality of the latents, which is coming from a beta-vae, and prior works on representation learning has showed that VAEs are not necessarily good for learning good quality latents. \n\n3. This draws to my other worry of how good the latents are when learnt from Beta-VAE, and if there is a qualitative analysis about the structure that Beta-VAE latents recover from these tasks. Recent works have shown that with representation learning objectives, the latent structure can be well recovered - this work only shows how the latents are distributed across the horizonz, but it is not clear if these latents can itself contain useful information specific to the task?\n\n4. Maybe not so related, but I wonder if an entropy-regularized objective can also be used in these settings to induce diverse skills learnt from the diffusion prior? The paper uses BCQ in offline, which is drawn from TD3 - so I agree that the algorithm is fundamentally using a deterministic policy approach; but for the discrete action tasks - maybe entropy regularized objectives to learn diverse skills may be introduced?\n\n5. Paper claims that the algorithm can work for both continuous and discrete action spaces - but there aren’t enough results showing for discrete actions? Why is that so?\n\n6. How does the claim on learning multi-modality for the diffusion approach compare with other representation objectives, rather than just the VAE? I think more baselines should be compared here, and I doubt that well tuned representation objectives would also be able to capture the multi-modality aspect similar to the diffusion model. \n\n7. Experimental results should be better organized; appendix seems to contain more domains, but they are not referred to the main; there are figures for the environments, e.g CARLA, but I couldnt gind results for that in the paper?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)