Efficient Planning with Latent Diffusion
Authors: Wenhao Li
Venue: ICLR 2024 poster (2024)
Keywords: Offline Decision-Making
Offline Reinforcement Learning
Generative Model
Diffusion Model
Primary area: reinforcement learning

Abstract: Temporal abstraction and efficient planning pose significant challenges in offline reinforcement learning, mainly when dealing with domains that involve temporally extended tasks and delayed sparse rewards. Existing methods typically plan in the raw action space and can be inefficient and inflexible. Latent action spaces offer a more flexible approach, capturing only possible actions within the behavior policy support and decoupling the temporal structure between planning and modeling. However, current latent-action-based methods are limited to discrete spaces and require expensive planning steps. This paper presents a unified framework for continuous latent action space representation learning and planning by leveraging latent, score-based diffusion models. We establish the theoretical equivalence between planning in the latent action space and energy-guided sampling with a pretrained diffusion model and introduce a novel sequence-level exact sampling method. Our proposed method, $\texttt{LatentDiffuser}$, demonstrates competitive performance on low-dimensional locomotion control tasks and surpasses existing methods in higher-dimensional tasks.

Published as a conference paper at ICLR 2024
EFFICIENT PLANNING WITH LATENT DIFFUSION
Wenhao Li
School of Software Engineering, Tongji University Shanghai, 201804, China
iwenhacecuhk. edu.
ABSTRACT
Temporal abstraction and efficient planning pose significant challenges in offline reinforcement learning, mainly when dealing with domains that involve temporally extended tasks and delayed sparse rewards. Existing methods typically plan in the raw action space and can inefficient and inflexible. Latent action spaces offer a more flexible paradigm, capturing only possible actions within the behavior policy support and decoupling the temporal structure between planning and modeling. However, current require expensive planning steps. This paper presents a unified framework for continuous latent action space representation learning and planning by leveraging latent, score-based diffusion models. We establish the theoretical equivalence between planning the latent action space and energy- -eedgy sampling with a pretrained diffusion model and incorporate novel sequence-level exact sampling method. Our proposed method, Lat datentDiffi fuseer, demonstrates competitive performance on low-dimensional locomotion control tasks and surpasses existing methods in higher-dimensional tasks.
actiin-bassee methods are limited discrete spaces and
INTRODUCTION
A considerable volume of samples gathered by operational systems gives rise to the issue of offline reinforcement learning (RL), specifically, the recovery of high-performing policies without additional environmental exploration (Wu et al., 2019; Kumar et al., 2020; Kostrikov et al., 2021; 2022; Ghosh al., 2022). However, domains that encompass temporally extended tasks and severely delayed sparse rewards can present a formidable challenge for standard offline approaches (Li et al., 2015; Ren et al., 2021; Li et al., 2023). Analogous to the online setting, an emergent objective in offline RL involves the development of efficacious hierarchy methodologies that can obtain temporally extended lower-level primitives, subsequently facilitating the construction a higher-level policy operating at more abstract temporal scale (Ajay et al., 2021; Pertsch et al., 2021; Villecroze et al., 2022; Rosete-Beat Beas et al., 2022; Rao et al., 2022; Yang et al., 2023).
Within the hierarchical framework, current offline RI. approaches be broadly categorized into model-free and model-based. The former conceptualizes the higher-level policy optimization as a auxilary offline RL issue (Liu al., 2020; Liu & Sun, 2022: Ma et al., 2022; Kipf al., 2019; Ajay et al., 2021; Rosete-Beaa et al., 2022). In contrast, the latter encompasses planning in the higher- level policy space by generating future trajectories through a dynamics model of the environment, either predefined or learned (Li et al., 2022; Co-Reyes et al., 2018; Lynch et al., 2020; Lee et al., 2022; Venkatraman, 2023). Concerning lower-level primitive learning, these two methods exhibit similaritiess and are typically modeled as goal-conditioned or skill-based imitation learning or offline RL problems. Conversely, the instabilities arising from offline hierarchical RL methodologies due to the "deadly triad (Sutton & Barto, 2018; Van Hasselt al., 2018)," restricted data access (Fujimoto et al., 2019; Kumar et al., 2020), and sparse rewards (Andrychowicz et al., 2017: Ma et al., 2022) remain unaddressed. This spawns another subset of model- based approaches along with more effective hierarchical variants that endeavor to resolve problems from a sequence modeling viewpoint Chen et al. (2021); Janner et al. (2021; 2022); Ajay et al. (2023).
Irrespective of whether a method is model-free or model--asedd it adheres to the traditional settings, wherein planning occurs in the raw action space of the Markov Decision Process (MDP). Although

---

Published as a conference paper at ICLR 2024
seemingly intuitive, planning in raw action space can be inefficient and inflexible (Wang et al., 2020; Yang et al., 2021: Jiang et al., 2023). Challenges include ensuring model accuracy across the entire space and the constraint of being tied to the environment's 's temporal structure. Conversely, human planning offers enhanced flexibility through temporal abstractions, high-level actions, backward planning. and incremental refinement.
Drawing motivation from TAP (Jiang et al., 2023), we put forth the notion of the latent action. Planning within the domain of latent actions delivers a twofold advantage compared to planning with raw actions. Primarily, it encompasses only plausible actions under behavior policy support, yielding reduced space despite the raw action space's dimensionality and preventing the exploitation of model frailties. Secondarily, it permits the separation of the temporal structure between planning and modeling, thus enabling more adaptable and efficient planning process unconstrained by specific transitions. These dual benefits render latent-action- baaodd approaches naturally superior to extant methodologies when handling temporally extended offline tasks.
Nevertheless, two shortcomings of TAP inhibit its ability to serve as general and practical framework. Initially, TAP confined to discrete latent action spaces. In real-world contexts, agents are likely to сarry out narrow, discrete assortment of tasks and a broader spectrum of behaviors (Co-Reyes et al., 2018). This introduces a predicament - should a minor skill modification be necessary, such as opening a drawer by seizing the handle from top to bottom instead of bottom top, completely novel set of demonstrations or reward functions might be mandated for behavior acquisition. Subsequently, once the latent action space has been ascertained, TAP necessitates distinct, isstucceinteesive planning phase for generating ceward-maximizg restricts latent actions discrete domains.
policies. The price of planning consequently
To tackle these limitations, this paper proposes a novel framework, LatentDi f fuser, by concur- rently modeling continuous latent action space representation learning and latent action-based plan- ning as a conditional generative problem within the latent domain. Specifically, Latent i f fuser employs unsupervised techniques to discern the latent action space by utilizing score-based diffusion models (SDMs) (Song et al., 2021; Nichol & Dhariwal, 2021; Ho & Salimans, 2022) within the latent sphere in conjunction with a variational autoencoder (VAE) framework (Kingma & Welling, 2014; Rezende et al., 2014; Vahdat et al., 2021). We first segment the input trajectories, map each slice to latent action space (which needs to be learned), and apply the SDM to the latent sequence. Subsequently, the SDM is entrusted with approximating the distribution over the offline trajectory embeddings, conditioned on the related return values. Planning or reward maximizing trajectory synthesis-is realized by initially producing latent actions through sampling from simple base distribution, followed by iterative, conditional denoising, and eventually translating latent actions into the trajectory space using decoder. In other words, Latent.i f ffuser be regarded as VAE equipped with an SDM prior (Vahdat et al., 2021).
Theoretically, we demonstrate that planning in the domain of latent actions is tantamount to energy- guided sampling using a pre-trained diffusion behavior model. Exact energy-guidee sampling is essential to carry out high-qualits and efficient planning. To achieve this objective, we modify QGPO (Lu et al., 2023) to realize exact sampling at the sequence level. Comprehensive numerical results on low-dimensional locomotion control tasks reveal that LatentDi f fuser exhibits com- petitive performance against robust baselines and outperforms them on tasks of greater dimensionality. Our main contributions encompass: 1) Developing a unified framework for continuous latent ac- tion space representation learning and planning that delivers flexibility and efficiency in temporally extended offline decision-making, 2) Our theoretical derivation confirms the equivalence between planning in the latent action space and energy -ueide sampling with a pretrained diffusion model. It introduces an innovative sequence-leve exact sampling technique. 3) Numerical experiments exhibit the competitive performance of La tentDi f fusee and its applicability across range of low- and high-dimensional continuous control tasks.
RELATED WORK
Owing to spatial constraints, this section will briefly present the most pertinent domain of LatentDi ffuser: offline RL or imitation learning (IL) based on hierarchical structure. In terms of algorithmic specificity, existing techniques can be broadly classified into goal- ball- and skill-based methods (Pateria et al., 2021). For further related literature, including but not limited to
2

---

Published as a conference paper at ICLR 2024
model-based RL, action representation learning, offline RL, and RL as sequence modeling, kindly refer to Appendix C and the appropriate citations within the papers.
Goal-based approaches primarily concentrate attaining a designated state. The vital aspect of such techniques concerns the selection or creation of subgoals, which reside in the raw state space. Once the higher-level subgoal is ascertained, the lower-level policy is generally acquired through standard IL methods or offline RL based on bulgoallaugmmented/conddiiooonno
policy, a universal value function (Schaul et al., 2015), or their combination. In extant methods, the subgoal is either predefined (Zhou et al., 2019; Xie et al., 2021; Ma et al., 2021), chosen based on heuristics (Ding et al., 2014; Guo & Zhai, 2016; Pateria et al., 2020; Mandlekar et al., 2020), or generated via planning or an additional offline RL technique (Liu et al., 2020; Liu & Sun, 2022; Li et al., 2022; Ma et al., 2022). Moreover, some methods (Eysenbach et al., 2019; Paul et al., 2019; Lai et al., 2020; Kujanpää et al., 2023) are solely offline during the subgoal selection or generation process. This paper also pertains to the options framework (Sutton et al., 1999; Stolle & Precup, 2002; Bacon et al., 2017; Wulfmeier et al., 2021; Salter et al., 2022; Villecroze et al., 2022)), as both the (continuous) latent actions of LatentDi f fuser and (discrete) options introduce a mechanism for temporal abstraction.
Skill-based methods embody higher-level skills low-dimensional latent codes. In this context, a skill signifies a subtask's 's policy, semantically representing "the capability to perform something adeptly" (Pateria et al., 2021). Analogous to goal-based approaches, once the higher-level skill is identified, the lower-level skill-conditioned policy is generally acquired through standard IL or offline RL methods. More precisely, few works utilize predefined skills (Nasiriany et al., 2022; Fatemi et al., 2022). The majority of studies employ a two- or multi-phase training framework: initially, state sequences are projected into continuous latent variables (i.e., skills) via unsupervised learning; next, optimal skills are generated based on offline RL (Kipf et al., 2019; Pertsch et al., 2021; Ajay et al., 2021: Rosete- Beae et al., 2022: Lee al., 2022; Venkatraman, 2023) or planning' (Co-Reyes et al., 2018; Lynch et al., 2020; Lee et al., 2022; Venkatraman, 2023) in the skill space.
Goal-candricned polioy
Skil
ae
Latent action
G
Figure 1: The physical meaning of the goal-conditioned policy, skill and latent action (corresponding timesteps in the raw MDP). The red diamond represents particular (goal) state, the gray, dotted diamond is placeholder, and the red circle denotes any state.
In contrast with the aforementioned hierarchical methodologies, LatentDi fuser initially learns a more compact latent action space and subsequently employs the latent actions to make decisions. As demonstrated in Figure 1, latent action not only differs from the conditiond policy, which pertains to the trajectory of reaching a particular state, but also from the skill, which relates to the trajectory of completing a specific (multi-step) state transition. The latent action also corresponds to the agent's 's received reward and the subsequent expected return. The unique physical implications of latent action and the methodology utilized by Latent D f ffuser render the proposed method advantageous in several ways. 1) The future information in the latent action allows the algorithm to execute more efficient planning. 2) Unlike existing works wherein multiple optimization objectives and the fully coupling or seperating of representation learning and decision making (RL or planning) lead to intricate training processes and reduced training efficiency, LatentDi fuser exhibits end-to-end training and unifies representation learning, sampling, and planning.
3 PROBLEM FORMULATION
In this paper, we approach the offline RL problem as a sequence modeling task, in alignment with previous work (Janner et al., 2022; Ajay et al., 2023; Li et al., 2023). The following subsection delin- eates the specificities of sequence modeling, or more accurately, the conditional generative modeling a
paradigm. We examine a trajectory, r, of length T, which sampled from a MDP that features a
It is important to note that planning is only feasible when the environment model is known or can be sampled from the environment model. Consequently, some of these works focus online RL. tasks, while others first. learn additional environment model from the offline dataset and then plan in the skill space.
3

---

Published as a conference paper at ICLR 2024
fixed stochastic behavior policy. This trajectory comprises (refer to Appendix for more modeling selections of 7) a series of states, actions, rewards, and reward-to-go values, Gt := Eizi '" r.s. proxies for future cumulative rewards: = ($1 a1. 11, G 82, 02. 22 G2..... ST, TT. is crucial to note that the definition of r diverges from that in prior studies (Janner et al., 2022: Ajay et al., 2023; Li al., 2023), as each timestep values. This modification has been specifically engineered to facilitate the subsequent learning of latent action spaces. Sequential peeiiionnmaainng is subsequently formulated the standard problem of conditional generative modeling:
rT, GT).
contains both the reward and reward-to-go
max [log (To y)(o)))
(1)
where 7h := т. The objective is to estimate the conditional trajectory distribution using pe to enable planning or generating the desired trajectory To based on the information y(7x)- Existing instances of may encompass the return (Janner et al., 2022; Li et al., 2023), the constraints met by the trajectory (Ajay et al., 2023; Li et al., 2023), or the skill demonstrated in the trajectory (Ajay et al., 2023). The generative model is constructed in accordance with the conditional diffusion process:
q (Tk++1 Tx); Po (Tk-1 Tke (To)))
As per standard convention, q signifies the forward noising process while pe represents the reverse denoising process (Ajay et al., 2023).
Latent Actions We introduce the concept of latent action (Figure 1) proposed in TAP (Jiang et al., 2023). TAP specifically models the optimal conditional trajectory distribution p' (r $1, using series of latent variables, 2 := (21...., 2MM Assuming that the state and latent variables (s), can be deterministically mapped to trajectory т, (" 81, z) p( (r +()1,=))==* =))x²(z 81) obtained. The terms z and x² 81) are subsequently referred to as the latent actions and the optimal latent policy, respectively. In a deterministic MDP, the trajectory corresponding to an arbitrary function ((s1,, z) with T (z 51) > will constitute an optimal executable plan, implying that the optimal trajectory be recovered by following the latent actions 2, beginning from the initial state $1. Consequently, planning within the latent action space Z facilitates the discovery of an desired, optimal trajectory. TAP, however, remain restricted discrete latent action spaces and necessitate indepentdent, resource-intensive planning. Motivated by these limitations, we present a unified framework that integrates representation learning and planning for continuous latent action via latent, score-based based diffusion models.
4 ALGORITHM FRAMEWORK
This section provides comprehensive elaboration of the model components and design choices, such as the network architecture, loss functions, as well as the details of training and planning. By unifying the representation learning and planning of latent action through the incorporation of a latent diffusion model and the exact energy-guided sampling technique, LatentDi ffuser achieves effective -ecisinn-making capabilities for temporally-extender explore the representation learning for latent action in Section 4, followed by detailed discussion on planning using energy guiddd sampling in Section 4.2, and provide a algorithm summary in Section 4.3 to close this section.
sparse reward tasks. Specifically, we first
Trajectory Decoder
State
Decoder
Encoder (Causal Transformer with Position Embedding)
with Positiin Embbdding
Action
Ras Aattra ft ta
Feedback Decoder
Reward Decoder
Diffusion Model
Return Decoder
x
Figure 2: Representation learning for latent action with the latent score-based diffusion model.
4

---

Published as a conference paper at ICLR 2024
4.1 REPRESENTATION LEARNING FOR LATENT ACTION
The latent action space allows for a more compact, efficient, and adaptable method by effectively capturing behavior policy support and detaching the temporal structure, thus providing innate ben- efits in handling temporally extended offline tasks. As indicated in Section 3, before proceeding to planning, we must first learn continuous latent action space. For this purpose, we propose the LatentDi f fuser based on latent diffusion model (LDM) (Vahdat et al., 2021), as de- picted in Figure 2. LatentDi filiiiffuss fuser constituted by an encoder qo (zo 81, +), a score-based prior Pe (zo $1), and decoder Pv (r $1 z0). In accordance with Vahdat et al. (2021), we train LatentDi i f fuser by minimizing the variational upper bound on the negative trajectory log- likelihood log (T] 81), meaning that the information (T) in Equation (1) is instantiated as the initial state $1:
(8, 7, (p) ( |- log Peb (T 81, [o]] + KT. (16 (zo 81, 7) IIpo (za 81))
(2olaa))))
| log Peb (T [1,20)] + E (zollrrr)))) [log lsb (zo [1;
(3)
+ (oo (2) log po (zo [1)]
utilizing a VAE approach (Kingma & Welling, 2014; Rezende et al., 2014), wherein the qo (zo 18, approximates the true posterior (Ro 80, T).
This paper employs Equation (3) with a decomposed KI. divergence into entropy and cross entropy terms. The reconstruction and entropy terms are easily estimated for any explicit encoder as long as the reparameterization trick is applicable (Kingma & Welling, 2014). The challenging aspect of training LatentDi f ffuser pertains to training the cross entropy term, which involves the score- based prior. Unlike Vahdat et al. (2021), which addresses this challenge by simultaneously learning an encoder/dccoder architecture alongside approach (Rombach et al., 2022) by training a VAE (go, Pw) } and a score-based diffusion model {go) consecutively based the offline dataset D,. This does not necessitate delicate balancing of reconstruction and generative capabilities.
basse-based prior, adopt simpler yet efficacious
Encoder do and Decoder Pv We use the almost consistent encoder design with TAP (Jiang et al., 2023). Specifically, we handle 2; := (st, at, Tr, G.) as a single token. The encoder & processes token 2; using GPT-2 style Transformer" yielding T feature vectors, where T is the episode horizon. Subsequently, we apply 1-dimensional max pooling with kernel size and stride of L, followed by a linear layer, and generate T/L L latent actions. Moreover, different from the TAP Decoder architecture, we use modular design idea. More concretely, each latent action is tiled J. times to match the number of input/output tokens T. We then concatenate the initial state 81 and the latent action, and apply a linear projection provide state information to the decoder. After adding positional embedding, the decoder reconstructs the trajectory f := (
Z7), with 2; == (&t, â₂, fz. G.). To enhance the decoder's s representation ability, we design the decoder modularly for different elements in æt, as shown in Figure 2. Noting that the action decoder is designed based on the inverse dynamics model (Agrawal et al., 2015; Pathak et al., 2017) in manner similar to (Ajay et al., 2023; Li et al., 2023), with the aim of generating raw action sequences founded the state sequences. The training of the encoder and decoders finally entails the use of reconstruction loss computed as the mean squared error between input trajectories {r} and reconstructed trajectories (+). coupled with low-weighted (=< 10-9) Kullback-Lein actions, akin to VAE approaches (Kingma & Welling, 2014; Rezende et al., 2014). This prevents the arbitrary scaling of latent action space.
penalty towards a standard normal on the learned latent
Score-based Prior 0
Having trained the VAE Pe.. we now have access to a compact latent action space. Distinct from VAE's adoption of a uniform prior TAP's utilization of an autoregressive, parameterized prior over latent actions, LatentDi fuser employs a score-based one. Thus, by harnessing the pattttsiionnaampiinaa---llaaaaae
framework, we seamlessly transform planning into conditional diffusion sampling, ultimately circumventing the need for an independent. costly planning stage. Concretely, the score-based prior is modeled as a conditional, reore-based diffusion probabilistic model, which is parameterized using temporal U-Net architecture (Janner et al., 2022; Ajay et al., 2023). This architecture effectively treats a sequence of noised latent action (2) as an image, where the height represents a single latent action's dimension and the width signifies
Differeet from the casual Transformer used in TAP, see Appendix for discussion.
5

---

Published as a conference paper at ICLR 2024
the number of the latent actions. Conditioning information yzz) := multi-layer perceptron (MLP). The training of the score-based prior is formulated as standard score-matching problem detailed in Appendix B.2.
$1 is then projected using
4.2 PLANNING WITH ENERGY- GUIDED SAMPLING
Upon acquiring the latent action space, are able to effectively address temporally-extendee offline tasks using planning. Intriguingly, when examined from a probabilistic standpoint, the optimal latent action sequence sampling coincides with guided diffusion sampling problem (Lu et al., 2023), wherein the guidance is shaped by an (unnormalized) energy function. By adopting a "diffusion- camplin--as-planiigg
framework (Janner al., 2022), we can perform planning through conditional sampling using the pretrained LatentDiffuser without necessitating further costly planning steps (Janner et al., 2021; Jiang et al., 2023). This renders LatentDi fuser a holistic framework that seamlessly consolidates representation learning and planning within the latent action space. In subsequent sections, the equivalence between optimal latent actions sampling and anergy-guided diffusion sampling is demonstrated, followed by the introduction of practical sampling algorithm to facilitate efficient planning.
Planning Energy- Gurdgd Diffusion Sampling Considering a deterministic mapping from T to 2, achieved by the learned encoder 96. the following theorem (refer to Appendix I.1 for the proof) is derived for the optimal latent policy defined in Section 3:
Theorem 1 (Optimal latent policy). Given an initial state s1, the optimal latent policy satisfies: n" (z 81) x u( 81 s1)e@ Σi= Qc(+),a))) wherein p(z $1) represents the behavior latent policy and Qc(:;) :) refers to the estimated Q-value function. B 2 0 signifies the inverse temperature controlling the energy strength.
By rewriting Po n² 9o l and 2p z, we can reformulate the optimal planning into the following diffusion sampling problem:
Po(zo $1) & 40(2u $1) exp (- s1))),
(4)
where (l((z. == E (c ax) and h(z0, 81) denotes the pretrained decoder Pe- The behavior latent policy qu(zu 81) is modeled by the pretrained Lat ffuser. We then adopt the "diffusion "diffusion-sampling as planning" to generate desired (e.g., reward-maximinn maximizing) latent actions diffusion timestep A 0. Then forward diffusion
Zo. Concretely, employ 40 =- 4. Po process is constructed to simultaneously diffuse do and Po into an identical noise distribution, where Pzo ko. s1) == qeo qso(z |20, $1) (N (zz/au=z0 I21). Based on (Lu al., 2023, Theorem 3.1), the marginal distribution ql and pæ of the noised latent actions at the diffusion timestep A: adhere to:
(z 81) x q 81) exp (E& (M(Zz,
(5)
where Ex (f(2,, $1)) is (h( 31)) when k = 0 and - log Equxez, (n)) (exp(=BE(h(m))) 81 s1)))
when & > O. We then need to estimate the score function of (* 81). Quoting the derivation of Lu (2023), the score function satisfies: Vz, log Pk (Zk 81) Vz, log qu (Z& (1) +
Consequently, the optimal planning has been formulated as energy-guided sampling within the latent action space, with V
Vz, &
81)) as the desired guidance.
Practical Sampling Method Estimating the target score function Vz, log (1)) is non- aniiiiivial because of the intractable energy guidance V., (h((z,,))) We borrow the energy energy-guided guided sampling method proposd in (Luu et al., 2023) and propose a sequence-level. exact sampling methods by training total of three neural networks: (1) diffusion model to model the behavior latent policy go (=o s1);
(2) state-action value function Qc(s, a) to define the intermediate energy function (h((z,, $1 s1)); and
(3) time-dependent energy model (z 81, k) estimate & (h(zk, 81)) and guide the diffusion sampling process.
Recall that already have (1) a diffusion model, i.e., the socre- based prior po(zo (1) and (2) state-action value function (s, a), i.e., the return decoder. According to Lu et al. (2023, Theorem 3.2), the only remained dime-dependent energy model, fy (2k; 81, k), be trained by minizing the

---

Published as a conference paper at ICLR 2024
following contrastive loss:
BEolh(a (1) .e1)) Σ j== e (Eo(h(z))) .s)))
log e (= ,
min min 1 I1:11 (==6))11) p()(e()))))
Σ M j==1 e (= .oakk)
(6)
where &o ~ (1, 2 = 2g + oge, and ~ (0, To estimate true latent actions distribution ((20 | 81) in Equation 6, we utilize the pretrained encoder q and score-based prior po to generate M support latent actions { ( for each initial state $1 by diffusion sampling. The contrastive loss in Equation (6) then estimated by:
M
BEOIMAA (81))
e (n (EL",.1.,)) .1.,))
min Eg
log
e (E)()(3))) ,1)) ,.))
Σn
where 20 (1) 20 (i) correspond the support latent actions for each initial state 81.
4.3 ALGORITIM SUMMARY
In general, the training phase of LatentDiffure is composed of three parts, corresponding to the training of encoder and decoders {qp. Pv)), ccsre-based prior pø. and intermediated energy model fn- as shown in Algorithm 1. Throughout the training process, it is imperative to employ two distinct datasets: the first being standard offline RI, dataset, D, which encompasses trajectories sampled from behavior policies, whereas the second dataset consists of support latent actions for each initial state $1 E D, generated by the pre-trained VAE, i.c., the encoder, score-based prior and decoders.
Algorithm Latent fuser: Efficient Planning with Latent Diffusion
Initialize the latent diffusion model, i.e., the encoder go. the score-based prior pe and the decoder Py; the intermediate energy model fn
for each gradient step do
Training the encoder and decoders
Sample B₁ trajectories T from offline dataset D
Generate reconstructed trajectories with the encoder qs and decoder P
Update {Φ, } based on the standard VAE loss
end for
for each gradient step do
>> Training the score-based prior
Sample B2 trajectories r from offline dataset DD
Sample B2 Gaussian noises from (0, I) and B, time k from (/( K) Generate latent actions 2g with the pretrained encoder 96 and decoder P
Perturb 2o according to == a4 + oge
Update (0) with the standard score-matching loss in Appendix B.2
end for
Generating the support latent actions Sample M support latent actions (zi (f) ()M M from the pretrained score-based prior pe end for
for each initial state LT in offline dataset D do
for each gradient step do
b Training the intermediate energy model
Sample B3 initial state sı from offline dataset D
Sample B3 Gaussian noises from NO, I) and B3 time k from L/(0, K)
Retrieve support latent actions { (x) ar for each
Perturb 20 (4) according to 2 (6) = ag 20 (8) OgE
Update {n} based the contrastive loss in Equation (7)
end for
Moreover, the optimal planning is tantamount to conducting conditional diffusion sampling based on the score-based prior and the intermediate energy model. Formally, the generation employs reverse denoising process at each diffusion timestep by utilizing the score function Vz, log Ps: (zx $1) based the score function of the based prior V z1 log q (zx |1) and intermediate energy model V V2,& & (h(*,, s1)), along with the state and action decoder p& (T $1. =n) to map the sam- pled latent actions 2a back to the original trajectory space. Explicitly, the generative process is
7

---

Published as a conference paper at ICLR 2024
p ($1, Zo. 7) Pu (zul s1) Ps (T) $1, 11. ). To avoid the accumulation of errors during sampling, adopt the receding horizon control used in the existing methods (Ajay et al., 2023; Li et al., 2023).
5 EXPERIMENTS
This section aims to assess the efficacy of the LatentDi fuser for extended temporal offline tasks in comparison current SOTA offline RI methods, which integrate hierarchical structures, and conditional generation models. The empirical evaluation encompasses three task categories derived from D4RL (Fu al., 2020): namely, Gym locomotion control, Adroit, and AntMaze. Gym locomotion tasks function as to ascertain whether LatentDiffi fuseer capable of accurately reconstructing trajectories for decision-making and control purposes. Subsequently, Latent fuser is evaluated on Adroit-a task with significant state and action dimensionality-42 well as on LatentDi ffuser within the AntMaze environment, which represents sparse-reward continuous-contro challenge in a series of extensive
in the lower-dimensional realm, in order
maps (Li et al., 2023). The subsequent sections will describe and examine the performance of these tasks and their respective baselines individually. Scores within 5% of the maximum per task will be emphasized in bold (Kostrikov et al., 2022).
5.1 SRDOr-OF-CONCERRE GYM LOCOMOTION CONTROL
Baselines Initially, an outline of the baselines is provided: CQL. (Kumar al., 2020), IQL (Kostrikov et al., 2022), D-QL (Wang al., 2023), and QGPO (Lu et al., 2023) are all model-free offline RL methods. MoReL (Kidambi et al., 2020) a model-based offline RL method. DT (Chen et al., 2021), TT (Janner al., 2021), Diffuser (Janner et al., 2022), and DD (Ajay et al., 2023) address offline RL tasks via conditional generative modeling. Finally, TAP (Jiang et al., 2023) and HDMI (Li et al., 2023) employ hierarchical framework grounded in generative modeling. Due to spatial constraints, only algorithms with the highest performance rankings are displayed herein; for a comprehensive comparison, please refer to the appendix.
Table The performance in Gym locomotion control in terms of normalized average returns. Results correspond to the mean and standard error over 5 planning seeds.
Dataset
Environment
CQL
D-QL
QGPO
HDMI
Med-Expert
HaliCheeetah
91.6
95
90.6:
96.8:10.3
91.8 + 0.8
93.5
92.1=1144
95.2-t0.2
Med- Med-Expert
Hopper
105.4
110.0
111.8:18
111.1:113
105.5 + 1.7
108.0:
113.5:0.9
112.9:10.3
Med-Expert
Walker2d
108.8
101.9
108.8=1.7
110.1:0.3
107.4 + 0.9
110.7 0.6
107,9=1.2
111.3+-022
Medium
HalfCheetah
44.0
46.9
49.1=1.0
51.1=20.5
45.0 0.1
54.1 0.4
48.0=0.9
53.6:10.4
Medium
Hopper
58.5
79.3:33.6
90.5:14.6
63.4
98.0
76.4::266
9..5
Medium
Walker2d
72.5
79
12.5=114
87.0:10.9
64.9 2.1
86.0 0.7
7.
86.3:10.9
Med Mep-ae
HallCheetah
45.5
41.9
39.3+4.11
4..
40.8 + 0.6
47.6 1.4
44.9++2.0
4..3++1.2
Med-Replay
Hopper
95
91.5
100=0.77
101.3:10.6
87.3 2.3
96,9 +: 2.6
99.6111.5
100.4:10.5
Med-Replan
Walker2d
77.2
82.6
75-44.3
95,5++1.5
66.8 + 3.1
84.4 + 4.1
80.7+2,11
82.6 2.1
Averuge
77.6
78.9
81.8
88.0
82.5
86.6
82.6
87.5
Table shows that LatentDiffuser surpasses specifically designed offline RL methods in the majority of tasks. Furthermore, the performance discrepancy between LatentDi ffuser and two-stage algorithms, such as TAP and HDMI, underscores the benefits provided by the the proposed framework, which unifies learning of latent action space representation and planning.
5.2 HIGH-DIMENSIONAL MDP: ADROIT
Baselines Taking into account the large dimensions characterizing the Adroit task actions, only baselines that perform well in the previous task are evaluated. Additionally, D-QL necessitates 50 repeated samplings by default for action generation (Wang et al., 2023). This requirement would result a substantial training overhead for high- high-dimensional action tasks. Consequently, to ensure a fair comparison, D-QL is configured allow only 1 sampling, akin QGPO (Lu et al., 2023).
Table 2 demonstrates the advantages of LntentDifi ffuser become even more pronounced in high- dimensional tasks. Furthermore, marked decrease in sequence modeling method performance is
8

---

Published as a conference paper at ICLR 2024
Table 2: Adroit results. These tasks have high action dimensionality (21 degrees of freedom)
Dataset
Environment
CQL
TT
DD
D-QL@1
TAP
QGPO
HDMI
LD
Human
Pen
37.5
36.4
64.1 = 9.0
66.0 = 8.3
76.5 8.5
73.9 1 8.6
66.2 + 8.8
79.0 8.1
Human
Hammer
1.3
1.1
1.1
1.2
Human
Door
6,9 = 1.2
8.0
8.
8.5 1.2
7.1
Human
Relocate
0.0
0.2 :: 0.1
0.2 = 0.1
0,2 + 0.1
0.1 0.1
0.2 0.1
Cloned
Pen
39.2
11.4
47.7 =: 9.2
49.3 8.0
$7.4 8.7
54.2 9.0
60.7 9.1
Cloned
Hammer
2.1
0.5
0.9 0.1
1.1
1.1
1.1
1.0
4.2 1 0.1
Cloned
Door
-0.1
9.0 1.6
10.6 10.6+117
11.7+11.5
11.2+144
9.3
12.0
Cloned Relocate
-0.1
-0.1
0.0
-.02
0.0
0.0
-0.1 0.0
-0.1
Expert
Pen
107.0
72.0
107.6
112.6 8.1
127.4
119.1
109.5
131.2
Expert
Hammer
86.7
15.5
106.7 = 1.8
114.8 = 1.7
127.6 1.7
123.2
111.8 1.7
132.5 1.8
Expert
Door
101.5
94.1
87.0 +: 0.8
93.7 : 0.8
104.8 1 0.8
98.8 1: 0.8
85.9 - 0.9
111.9 + 0.8
Expert
Relocate
95.0
10.3
87.5 +: 2.8
95.2 2.8
105.8 2.7
102.5 +: 2.8
91.3 +: 2.6
109.5 2.8
Average expert)
6.1
16.2
17.1
18.79
16.6
21.3
Average (w/ expert)
40.3
20.1
43,2
46.1
51.9
49.5
44.3
54.6
observed. Two primary factors are identified: first, larger action dimensions necessitate tokenization- and nutoregression-basee techniques (such TT) to process increasingly lengthy sequences; second, DD and HDMI employ an inverse dynamic model to generate actions independently. while the expansion in action dimension renders the model fitting process challenging.
5.3 LONG-HORIZION CONTINUOUS CONTROL: ANTMAZE
Baselines To validate the benefits of latent actions longer-horizon tasks, an additional com- parison is made with hierarchical offline RL methods designed explicitly for long-horizon tasks: CompILE (Kipf et al., 2019), GoFAR (Ma et al., 2022), and HiGoC (Li et al., 2022). Concurrently, CQL and TT are removed due their inability to perform well high-dimeensioal Adroit.
Table 3: AntMaze performance correspond to the mean and standard error over planning seeds.
Eavironment
CompILE
GoFAR
HBGoC
DD
D-QL@1
TAP
QGPO
HDMI
LD
AntMaze- Play
U-Maze-3
41.2+
38.5 155::222 2.2
31.2 3.2
73.1 + 2.5
52.9 4,1
82.2+ 211
30.3 1.3
86.1 2.4
85.4+ 1.9
AmtMaze- Divers U-Maze-3
23.5 1.8
25.1 3.1
25.5 1.6
49.2 3.1
32.5
5.9
69.81.0.5
38.5
73.71:111
75.6:1211
AmtMaze-
Large-2
69.21.32
71.51:35
75812.0
Single-task Average
32.4
31.8
28.4
56.4
30.0
73.7
45.4
77.1
78.9
MulliAnt-Divese Larpe-2
71.6
73.6
Averuge
Table 3 highlights that sequence modeling based hierarchical methods significantly surpass RL-based approaches. Moreover, LatentDi fuser demonstrates performance comparable to two-stage techniques such as TAP and HDMI through end-to-end training.
6 CONCLUSIONS
In this work, we present a novel approach, LatentDiffuser, for tackling temporal-extended offline tasks, addressing the limitations of previous state-of-the-art offline reinforcement learn- ing methods and conditional generation models in handling high-dimensional, long-horizon tasks. LatentDi ffuser is capable of end-to-end learning for both representation of and planning with latent action, delivering a unified, comprehensive solution for offline decision-makin and control. Numerical results on Gym locomotion control, Adroit, and AntMaze, demonstrate the effectiveness of LatentDii fuser in comparison with existing hierarchical- and planning based offline meth- ods. The performance gains are particularly noticeable in nigh-dimensional and long- -orizon tasks, illustrating the advantages of LatentDi ffuser in addressing these challenging scenarios.
ACKNOWLEDGMENTSS
We extend our heartfelt gratitude to Professor Hongyuan Zha for his enlightening discussions. This work was supported in part by Postdoctoral Science Foundation of China (2022M723039).

---

Published as a conference paper at ICLR 2024
REFERENCES
Pulkit Agrawal, Joao Carreira, and Jitendra Malik. Learning to see by moving. In ICCV, 2015.
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. Opal: Offline primitive discovery for accelerating offline reinforcement learning. In ICLR, 2021.
Anurag Ajay. Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? In ICLR, 2023.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong. Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. In NeurlPS, 2017.
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In AAAI, 2017.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In NeurIPS, 2021.
John Co-Reyes, YuXuan Liu, Abhishek Gupta, Benjamin Eysenbach, Pieter Abbeel, and Sergey Levine. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajectory embeddings. ICML, 2018.
XIAO Ding, Yi-tong LI, and SHI Chuan. Autonomic discovery of subgoals hierarchical reinforce- ment learning. The Journal of China Universities Posts and Telecommunications, 21(5))94-100 2014.
Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Search on the replay buffer: Bridging planning and reinforcement learning. In NeurIPS, 2019.
Mehdi Fatemi, Mary Wu, Jeremy Petch, Walter Nelson, Stuart Connolly, Alexander Benz, Anthony Carnicelli, and Marzyeh Ghassemi. Semi-markov offline reinforcement learning for healthcare. In CHIL, 2022.
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, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In /CML, 2019.
Dibya Ghosh, Anurag Ajay, Pulkit Agrawal, and Sergey Levine. Offline rl policies should be trained to adaptive. In /CML, 2022.
Xiaobo Guo and Yan Zhai. K-means clustering based reinforcement learning algorithm for automatic control in robots. International Journal of Simulation Systems, Science & Technology, 17(24): 6.1-6.6, 2016.
Jonathan Ho and Tim Salimans. Classifier-fie diffusion guidance. arXiv preprint arXiv:2207.12590 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In NeurIPS, 2021.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In ICML, 2022.
Zhengyao Jiang, Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefenstette, and Yuandong Tian. Efficient planning in a compact latent action space. In ICLR, 2023.
Rahul Kidambi, Aravind Rajeswaran, Pranceth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In NeurlPS, 2020.
10

---

Published as a conference paper at ICLR 2024
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. Stat, 1050:10, 2014.
Thomas Kipf. Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Sanchez-Gonzalez, Edward Grefen- stette, Pushmeet Kohli, and Peter Battaglia. CompILE: Compositional imitation learning and execution. ICML, 2019.
Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. ICML, 2021.
lyy Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In ICLR, 2022.
Kalle Kujanpää, Joni Pajarinen, and Alexander Ilin. Hierarchical imitation learning with vector quantized models. In ICML, 2023.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In NeurlPS, 2020.
Yaqing Lai, Wufan Wang, Yunjic Yang, Jihong Zhu, and Minchi Kuang. Hindsight planner. In AAMAS, 2020.
Taeyoon Lee, Donghyun Sung, Kyoungyeon Choi, Choongin Lee, Changwoo Park, and Keunjun Choi. Learning dynamic manipulation skills from haptic-play, arXiv preprint 2207. 2207.10777,, 2022.
Jinning Li, Chen Tang, Masayoshi Tomizuka, and Wei Zhan. Hierarchical planning through goal-
conditioned offline reinforcement learning. arXiv preprint arXiv:2205.11799 2022.
Lihong Li, Rémi Munos, and Csaba Szepesvári. Toward minimax off- off-policy value estimation. In AISTATS, 2015.
Wenhao Li, Xiangfeng Wang, Bo Jin, and Hongyuan Zha. Hicrarchical diffusion for offline decision making. ICML, 2023.
Jianfeng Liu, Feiyang Pan, and Ling Luo. Gochat: Goal-oriented chathots with hierarchical reinforce- ment learning. In SIGIR, 2020.
Qiang Liu, Jason Lee, and Michael Jordan. A kernelized stein discrepancy for goodness-of-fe tests. In ICML,, 2016.
Shaofan Liu and Shiliang Sun. Safe offline reinforcement learning through hierarchical policies. In PAKDD, 2022.
Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In ICMI., 2023.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In CoRI,, 2020.
Yecheng Jason Ma, Jason Yan, Dinesh Jayaraman, and Osbert Bastani. Offline coal-conditioned reinforcement learning via f-advantage regression. In NeurIPS, 2022.
Yi Ma, Xiaotian Hao, Jianye Hao, Jiawen Lu, Xing Liu, Tong Xialiang, Mingxuan Yuan, Zhigang Li, Jie Tang, and Zhaopeng Meng. A hierarchical reinforcement learning based optimization framework for large-scale dynamic pickup and delivery problems. In NeurlPS, 2021.
Ajay Mandlekar, Fabio Ramos, Byron Boots, Silvio Savarese, Li Fei-Fei, Animesh Garg, and Dieter Fox. Iris: Implicit reinforcement without interaction at scale for learning control from offline robot manipulation data. In ICRA, 2020.
Soroush Nasiriany, Huihan Liu, and Yuke Zhu. Augmenting reinforcement learning with behavior primitives for diverse manipulation tasks. In ICRA, 2022.
11

---

Published as a conference paper at ICLR 2024
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In ICMMCL, 2021.
Shubham Pateria, Budhitama Subagdja, and Ah Hwee Tan. Hierarchical reinforcement learning with integrated discovery of salient subgoals. In AAMAS, 2020.
Shubham Pateria, Budhitama Subagdja, Ah-hwee Tan, and Chai Quek. Hierarchical reinforcement learning: A comprehensive survey. ACM Computing Surveys, 54(5):1-35, 2021. 2021.
Deepak Pathak, Pulkit Agrawal, Alexei Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In ICML, 2017.
Sujoy Paul, Jeroen Vanbaar, and Amit Roy-Chowdhury. Learning from trajectories via subgoal discovery. In NeurlPS, 2019.
Karl Pertsch, Youngwoon Lee, Yue Wu, and Joseph J. Lim. Demonstration-guideed reinforcement learning with learned skills. In CoRI., 2021.
Dushyant Rao, Fereshteh Sadeghi, Leonard Hasenclever, Markus Wulfmeier, Martina Zambelli, Giulia Vezzani, Dhruva Tirumala, Yusuf Aytar, Josh Merel, Nicolas Heess, al. Learning transferable motor skills with hierarchical latent mixture policies. In ICLR, 2022.
Tongzheng Ren, Jialian Li, Bo Dai, Simon S Du, and Sujay Sanghavi. Nearly horizon-free offline reinforcement learning. In NeurlPS, 2021.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In JCML, 2014.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, 2022.
Erick Rosete-Beas, Oier Mees, Gabriel Kalweit, Joschka Boedecker, and Wolfram Burgard. Latent plans for task-agnostic offline reinforcement learning. In CoRL, 2022.
Sasha Salter, Markus Wulfmeier, Dhruva Tirumala, Nicolas Heess, Martin Riedmiller, Raia Hadsell, and Dushyant Rao. Mo2: Model-based offline options. In CoI.LAs, 2022.
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In ICMI., 2015.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In (CM1,,2015
Yang Song. Generative modeling by estimating gradients of the data distribution. yang-song.nel, May 2021. URL
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021.
Martin Stolle and Doina Precup. Learning options in reinforcement learning. In Abstraction, Reformulation, and Approximation: 5th International Symposium, SARA 2002 Kananaskis, Alberta, Canada August 2-4, 2002 Proceedings 5, pp. 212 22-32.3.222 Springer, 2002.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning, Artificial intelligence, 1 112(1-2)1811-111, 1999.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In NeurlPS, 2021.
Hado Van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil. Deep reinforcement learning and the deadly triad. arXiv preprint 1rXiv:1812. 02648, 2018.
12

---

Published as a conference paper at ICLR 2024
Siddarth Venkatraman. Latent skill models for offline reinforcement learning. Master's thesis, Carnegie Mellon University Pittsburgh, PA, 2023.
Valentin Villecroze, Harry Braviner, Panteha Naderian, Chris Maddison, and Gabriel Loaiza-Ganem. Bayesian nonparametrics for offline skill discovery. In ICML, 2022.
Linnan Wang, Rodrigo Fonseca, and Yuandong Tian. Learning search space partition for black-box optimization using monte carlo tree search. In NeurlPS, 2020.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou. Diffusson policies as expressive policy class for offline reinforcement learning. In ICLR, 2023.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprini arXiv: /911. 1361, 2019.
Markus Wulfmeeer, Dushyant Rao, Roland Hafner, Thomas Lampe, Abbas Abdolmaleki, Tim Hertweck, Michael Neunert, Dhruva Tirumala, Noah Siegel, Nicolas Heess, et al. Data-efficient hindsight off-policy option learning. In ICML, 2021.
Ruobing Xie, Shaoliang Zhang, Rui Wang, Feng Xia, and Leyu Lin. Hierarchical reinforcement learning for integrated recommendation. In AAAI, 2021.
Kevin Yang, Tianjun Zhang, Chris Cummins, Brandon Cui, Benoit Steiner, Linnan Wang, Joseph E Gonzalez, Dan Klein, and Yuandong Tian. Learning space partitions for path planning. In NeurIPS, 2021.
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. In AAAI, 2023.
Guojing Zhou, Hamoon Azizsoltani, Markel Sanz Ausin. Tiffany Barnes, and Min Chi. Hierarchical reinforcement learning for pedagogical policy induction. In IJCAI, 2019.
13

---

Published as a conference paper at ICLR 2024
Supplementary Material
Table of Contents
A Limitations and Socictal Impact
14
B Preliminaries
15
Offline Reinforcement Learning
B.2 Diffusion Probabilistic Models
15
C Missing Related Work
16
C.I. Model-based Reinforcement Learning.
16
C.2 Action Representation Learning.
16
C,3 Offline Reinforcement Learning.
Reinforcement learning as Sequence Modeling.
C.5 Controllable Sampling with Generative Models
17
D Missing Results and Analyses
18
D.1 Proof-of-Concept Example: maze-2d-open
18
D.2 Main Results
Implementation and Training Details
E.1 Baseline Details
20
E.2 Implementation Details
22
F Modeling Selection
23
F11 Planning in the Raw Action Space
wtw.y...........................
23
F2 Planning the Skill Space
25
G Ablation Studies
26
H Latent Action Visualization
27
I Missing Derivations
28
1.1 Proof Theorem
28
1.2 Proof Theorem and Theorem
29
A LIMITATIONS AND SOCIETAL IMPACT
Limitations LatentDi fuser, analogous to other diffusion-based methods for offline decision- making, exhibits a protracted inference time owing to the iterative nature of the sampling process. This challenge could be alleviated through the adoption of approaches that enable accelcrated sampling (Lu al., 2022abb) or by distilling these diffusion models into alternative methods necessitating fewer sampling iterations (Song et al., 2023). Additionally, similar with TAP (Jiang et al., 2023), empirical findings from continuous control featuring deterministic dynamics indicate that LatentDi f fuser manage epistemic uncertainty. However, the efficacy of LatentDi fuuser in addressing tasks characterized by stochastic dynamics without modifications remains unascertained. Furthermore, a deficiency in methodology is the requirement for both the latent steps L and planning horizon H for the latent action to remain constant. We hope that facilitating adaptive variation of these hyperparameters may enhance performance.
14

---

Published as a conference paper at ICLR 2024
Societal Impact Similar with other deep generative modeling techniques. the energy-guine guided diffusion sampling employed in this paper possesses the potential to generate harmful content and may perpetuate and exacerbate pre-existing undesirable biases present in the offline dataset.
B PRELIMINARIES
B.1 OFFLINE REINFORCEMENT LEARNING
In general, reinforcement learning (RL) represents the problem of sequential decision decision-making through a Markov Decision Process M (S, A, P, r, 7), encompassing state space and an action space A. Given states 8, & E S and an action a É A. the transition probability function is expressed as P (s') 8, 8,a): :
S > [0, 1] and the reward function is defined by r(s,a, S A x + TR. The discount factor is denoted
(0, 1/.
The policy represented as S state as (a s). For the timestep & E [1,7', the cumulative discounted reward, also referred to as reward-to-go. is identified by Rg -v= -Tp. The principal objective of online RL is to tr 8t, Gt, 8t++1) )] via
+ [0, 1], indicating the probability of taking action a in
determine a policy π that maximizes J = (aaraa(-lee)) (P PPll: learning from transitions (s, (s,a,r, rrr' s') during environment interaction (Sutton & Barto, 2018).
|Et=>
Conversely, in offline RL, static dataset D is employed, which has been collected through a behavior policy Ty. for acquiring a policy π that optimizes J for subsequent application in the interactive environment. The behavior policy Jy. can either constitute single policy or an amalgamation of various policies; however, it remains inaccessible. The acquisition of data is presumed to be trajectory-wise, as represented by D = {m))rr where (8;, Gg 8
B.2 DIFFUSION PROBABILISTIC MODELS
This section will provide an introduction to the diffusion probabilistic model within the context of the LatentDi ffuser. Diffusion probabilistic models (Sohl-Dickstee et al., 2015; Ho et al., 2020), constitute a likelihood-bon
generative framework that facilitates learning data distributions g(=) from the offline datasets expressed as D == {='). wherein the index i denotes specific sample within the dataset (Song, 2021), and 2 is the latent actions encoded by the pre-trained encoder qs- A core concept within diffusion probabilistic models lies in the representation of the (Stein) score function (Liu et al., 2016), which does not necessitate a tractable normalizing constant (also referred to as the partition function).
The siccrete-time generation procedure encompasses a designed forward noising (or diffusion) process 0( = (( z (1 ox) at (forward) diffusion timestep k. The for- ward process coupled with a learnable, reverse denoising (or diffusion) process po(=*-1 N(z (2) k), Σx) (backward) diffusion timestep k. NN(p, Σ) signifies Gaussian distribu- tion characterized by mean jt and variance Σ, ag E establishes the variance schedule. In order to ensure consistency with the main text notation, we denote ag := Vāx and σ := vI cg. 20 == corresponds to sample in D, 21, Z2. =K- signifies the latent variables or the noised latent actions, and žк ~ A(0 I) for judiciously selected âg values and sufficiently extensive K.
Commencing with Gaussian noise, samples undergo iterative generation via sequence of denoising steps. An optimizable and tractable variational lower-bound on log po serves to train the denoising operator, with a simplified surrogate loss proposed in (Ho et al., 2020):
Ldenoise (0) = J.te:1 [||e ca (z 13]
(8)
The predicted noise cø (2 k), parameterized through a deep neural network, emulates the noise I), integrated with the dataset sample zu yielding noisy 2 in the noising process.
Conditional Diffusion Probabilistic Models Intriguingly, the conditional distribution ((x()((2) facilitates sample generation under the condition y²x). Within the context of this paper, y(z) is instantiated as the initial state s1. The equivalence between diffusion probabilistic models and score- matching (Song al., 2021) reveals that en (2k,, k) x Vz, log p( 2k), giving rise two categorically
15

---

Published as a conference paper at ICLR 2024
equivalent methodologies for conditional sampling with diffusion probabilistic models: classifier- guided (Nichol & Dhariwal, 2021), and classifier-free (Ho & Salimans, 2022) techniques employed in work. The latter method modifies the preliminary training configuration, learning both a conditional c0(2), 81,k) k) and an unconditional eo (zx,A) model for noise. Unconditional noise manifests conditional noise co(z,, 2,, k), with a placeholder Φ replacing 81. When y(=) Ø the entries of e are zeroed out. The perturbed noise éN == eø Ee(*,. k) + (=) 81, - c k)) subsequently employed to generate samples. Additionally, we adopt low-temperture sampling in the denoising process to ensure higher quality latent actions (Ajay et al., 2023). Concretely, we compute P 1 and 1 from the previous noised latent actions zk-1 and perturbed noise č and subsequently sample 1 1 ~ N (μR-1 1, RR-- with the variance scaled by ar E [0, 1).
C MISSING RELATED WORK
C.I MODEL-BASED REINFORCEMENT LEARNING
LatentDi ffuser is incorporated into a research trajectory focused on model-based reinforcement leamnngg (RL) (Sutton, 1990; Janner et al., 2019; Schrittwieser et al., 2020; Lu et al., 2021; Eysenbach et al., 2022; Suh et al., 2023), as it makes decisions by forecasting future outcomes. These approaches frequently employ predictions in the Markov Decision Process (MDP), which entails that models accept the current raw state and action as input, outputting probability distributions encompassing subsequent states and rewards. Hafner al. (2019), Ozair al. (2021), Hafner al. (2021), Hafner et al. (2023), and Chitnis et al. (2023) proposed to acquiring a latent state space in conjunction with a dynamics function. Contrarily, in their cases, the action space accessible the planner remains identical to that of the raw MDP, and the execution of the plan maintains its connection to the original temporal structure of the environment.
C.2 ACTION REPRESENTATION LEARNING.
The concept of learning a representation for actions and conducting RL within a latent action space has been investigated in the context of model-free RL (Merel et al., 2019; Allshire al., 2021; Zhou et al., 2021; Chen et al., 2022; Peng et al., 2022; Dadashi et al., 2022). In contrast to LatentDi f fuser, where the latent action space utilized to promote efficacy and robustness in planning, the motivations for obtaining latent action space in model-free approaches vary, yet the underlying objective centers on providing policy constraints. For instance, Merel et al. (2019) and Peng et al. (2022) implement this concept for humanoid control to ensure the derived policies resemble low-level human demonstration behavior, thus being classified as natural. Zhou et al. (2021) and Chen et al. (2022) employ latent actions prevent -itt---diittibttion (OOD) actions within the offline RL framework. Dadashi et al. (2022) proposes adopting a discrete latent action space to facilitate the application of methods designed explicitly for discrete action spaces to continuous cases. In teleoperation literature, Karamcheti et al. (2021) and Losey et al. (2022) embed high- nigh-dimensional robotic actions into lower-dimensional, human-controllabe latent actions.
Additionally, several works focus on learning action representations for improved planning efficiency, Wang et al. (2020) and Yang et al. (2021) learn action representations for on-the-fly learning, appli- cable to black-box optimization and path planning scenarios. Despite high-level similarities, these papers assume prior knowledge of environment dynamics. TAP (Jiang et al., 2023) extends this framework into the offline RI. domain, where the actual environmental dynamics remain undeter- mined, necessitating joint learning of the dynamics model and the representation of latent action. Nevertheless, TАР is constrained to discrete latent action space and demands costly additional planning. LatentDi LafentDiffuser ffuser can achieve representation learning and planning for continuous latent actions by leveraging the latent diffusion model in end-to-end manner.
C.3 OFFLINE REINFORCEMENT LEARNING.
LatentDi fuser is devised for the offline RL (Ernst et al., 2005; Levine et al., 2020), precluding the utilization of online experiences for policy improvement. A principal hurdle in offline RL involves preventing outtof--iittibution (OOD) actions selection by the learned policy to circumvent value function and model inaccuracies exploitation. Conservatism (Kumar al., 2020; Kidambi et al.,
16

---

Published as a conference paper at ICLR 2024
2020b; Fujimoto & Gu, 2021; Lu et al., 2022c; Kostrikov et al., 2022) is proposed standard solution for this challenge. LatentDiffuse inherently prevents OOD actions via planning in a learned latent action space.
Pursuing adherence to a potentially diverse behavior policy, recent works have identified diffusion models as powerful generative tools, which generally surpass preceding generative approaches such as Gaussian distribution (Peng et al., 2019; Wang et al., 2020b; Nair et al., 2020) and Variational Autoencoders (VAEs) (Fujimoto et al., 2019; Wang al., 2021) concerning behavior modeling. Different methods adopt distinct strategies for action generation, maximizing the learned Q-functions. Diffusion- QL (Wang et al., 2023) monitors gradients from behavior diffusion volicy-derived actions to guide generated actions towards higher Q-value regions. StBC (Chen et al., 2023b) and Diffusion- QL employ a similar idea, whereby resampling actions from multiple behavioral action candidates occur, with predicted Q-values serving as sampling weights. Ada et al. (2023) incorporates a state reconstruction loss-based regularization term within the diffusion -based policy training, consequently bolstering generalization capabilities for OOD states. Alternative works (Goo & Niekum, 2022; Pearce et al., 2023; Block et al., 2023; Suh et al., 2023) solely deploy diffusion models for behavior cloning or planning. rendering Q- value maximization unnecessary.
Contrary to the works above aligned with the RL paradigm. LatentDiffuser addresses offline RL challenges through sequence modeling (refer to the subsequent section). Compared to RL-based offline methodologies, sequence modeling offers benefits regarding temporally extended and sparse or delayed reward tasks.
C.4 REINFORCEMENT LEARNING AS SEQUENCE MODELING
LatentDi ffuser stems from an emerging body of research that conceptualizes RL, as a sequential modeling problem (Bhargava et al., 2023). Depending on the model skeleton, this literature may be classified into two primary categories. The first category comprises models that leverage GPT- 2 (Radford et al., 2019) style Transformer architecture (Vaswani et al., 2017), also referred to causal transformers, for the autoregressive modeling of states, actions, rewards, and returns, ultimately converting predictive capabilities into policy. Examples include Decision Transformer (DT, Chen ct al. 2021) and Zheng et al. (2022), which apply an Upside Down RL technique (Schmidhuber, 2019) under both offline and online RL settings, and Trajectory Transformer (TT, Janner et al. 2021), which employs planning to obtain optimal trajectories maximizing return. Chen et al. (2023a) introduced a non-autoregressive planning algorithm based on energy minimization, while Jia et al. (2023) enhanced generalization ability for unseen tasks through refined in-context example design. Lastly, Wu et al. (2023) addresses trajectory stitching challenges by adjusting the history length employed in DT.
The second category features models based on score- diffusion process for non -oo--uutreerrssvvee modeling of state and action trajectories. Different methods select various conditional samplers to generate actions that maximize the return. Diffuser (Janner et al., 2022) emulates the classifier- guidance methodology (Nichol & Dhariwal, 2021) and employing guidance methods as delineated in $ F.1. Alternatively. Decision Diffuser (Ajay et al., 2023) and its derivatives (Li et al., 2023; Hu et al., 2023) explore classifier-free guidance (Ho & Salimans, 2022). Extensions of this concept to multi-task settings are presented by He et al. (2023) and Ni et al. (2023), while Liang et al. (2023) utilizes the diffusion model as a sample generator for unseen tasks, thus improving generalization capabilities. LatentDi f fuser offers more efficient planning solution to enable these sequential modeling algorithms to navigate complex action spaces effectively.
C.5 CONTROLLABLE SAMPLING WITH GENERATIVE MODELS
LatentDi E ffuser produces trajectories corresponding to optimal policies by employing control- lable sampling within diffusion models. Current methods for facilitating controllable generation in diffusion models primarily emphasize conditional guidance. Such approaches leverage a pretrained diffusion model for the definition of the prior distribution g(x) and strive to obtain samples from ((r) (xp(-)8(xx)). Graikos et al. (2022) introduces a training-fiee free sampling technique, which finds application in approximating solutions for traveling salesman problems. Poole et al. (2023) capitalizes on pretrained 2D diffusion model and optimizes 3D parameters for generating 3D shapes. Kawar et al. (2022) and Chung et al. (2023) exploit pretrained diffusion models for addressing linear and specific non-linear inverse challenges, such as image restoration, deblurring, and denoising. Dif-
17

---

Published as a conference paper at ICLR 2024
fuser (Janner et al., 2022) and Decision Diffuser (Ajay et al., 2023) use pretrained diffusion models to solve the offline RL problem. Zhao et al. (2022) and Bao et al. (2023) employ human- caman-crafted intermediate energy guidance for tasks including image-to-imane design. In a more recent development, Lu et al. (2023) presents a comprehensive framework for incorporating human control within the sampling process of diffusion models.
translation and inverse molecular
Remark Recently, Venkatraman (2023) proposed a novel algorithm called LDCQ, which is very similar to LatentDiffuser Specifically, LDCQ also introduces latent diffusson model to learn a latent action space. Unlike the latent action used by LatentDiffuseer the learned latent action space in LDCQ belongs the skill space, similar to what is described in Appendix F.2. Additionally, LDCQ does not perform planning within the learned skill space but instead uses model-free TD- learning methods to choose the optimal skill at each timestep and obtains the final action using decoder. In summary, it is quite coincidence that LDCQ and Latent fuuer belong to two orthogonal approaches to utilizing latent diffusion models in offline RL. The former still adopts the RL framework to model the offline RL problems, while LatentDiffuser approaches the problem from conditional generative perspective.
D MISSING RESULTS AND ANALYSES
D.1 PRONE-OF-CONCEPT EXAMPLE: MAZE-2D-OPEN
(a) Suboptimal trajectories.
(b) Stitched near optimal trajectories.
Figure Proof-of Proof-of-Concept example. We demonstrate the importance of planning through an experiment designed by Decison Diffuser (Ajay et al., 2023, DD; Appendix A.1). The diffusion model achieves trajectory stitching, process essential for handling large number of suboptimal trajectories, through implicit planning.
To demonstrate the importance of planning, designed an experiment same Ajay al. (2023, Appendix A. 1). Specifically, most tasks' offline datasets contain a large number of suboptimal trajectories. To learn better policies rather than just simple behavior cloning, trajectory stitching one of the essential abilities algorithms must possess. To validate whether LatentDi f fuser can achieve trajectory stitching through implicit planning, we adopted the experimental setup same as Ajay et al. (2023, Appendix A. In the Slaze-2D-oge towards the target area situated on the right side, with the reward being the negative distance this target area. The training dataset is composed of 500 trajectories originating from the left side and terminating at the bottom side, as well as 500 trajectories starting from the bottom side and ending at the right side. Each trajectory is constrained to maximum length of 50. At test time, the agent begins on the left side and aims to reach the right side as efficiently as possible. As demonstrated in Figure and consistent with the findings of Ajay et al. (2023, Appendix A. the Latent Di i fuser can effectively stitch trajectories from the training dataset to produce trajectories that traverse from the left side to the right side in (near) straight lines.
environment, the objective is to navigate
18

---

Published as a conference paper at ICLR 2024
D.2 MAIN RESULTS
The performance in Gym locomotion control in terms of normalized average returns of all baselines shown in Table 4.
Table 4: The performance in Gym locomotion control in terms of normalized average returns of all baselines. Results correspond to the mean and standard error over planning seeds.
Dataset
Environment
CQL
IQL
DT
MoReL
Diffuser
Med-Expert
HalrCheetah
91.6
86.7
86.8
95
53.3
79.8
Med-Expert
Hopper
105.4
91.5
107.6
110.0
108.7
107.2
Med-Expert
Walker2d
108.8
109.6
108.1
101.9
95.6
108.4
Medium
HallCheetah
44.0
47.4
42.6
Medium
Hopper
58.5
66.3
67.6
61.1
95.4
58.5
Medium
Walker2d
72.5
78.3
74.0
77.8
Med-Replay
HallCheetah
45.5
44.2
36,6
41.9
40.2
42.2
Med-Replay
Hopper
94.7
82.7
93.6
Med-Replay
Walker2d
77.2
73.9
66.6
82.6
49.8
61.2
Average
77.6
77
74.7
78.9
72.9
75.3
Dataset
Environment
D-QL
TAP
QGPO
LatentDiffuser
Med-Expert
HaltCheetah
1..6+113
96,8=0.3
91.8 = 0.8
93.5==0.3
92.1+1.4
95.2-0,2
Med-Expert
Hopper
11188-888
111.1=1.3
105.5
108.0=22.5
113.5++0.9
112.9=0.3
Med-Expert
Walker2d
108.8
110.1
107.4 1 0,9
110.7 .L 0.6
107.9
111.3_0. 0.2
Medium
HallCheetah
49.111..0
51.1.0.5
54.1
48.0.10.9
53.6 0.4
Medium
Hopper
79.3::3.6
9..5
63.4 1: 1.4
98.0 1: 2.6
76.4:1..66
98.5 0.7
Medium
Walker2d
82.5+1.4
87.0=0.9
6.19
86.0 0.7
79.9#1.
86.3=0.9
Med-Replay
HalfChectah
39.31:.11
47.8=0.3
40.8 00.8:000
47.6 + 1.4
44.942.0
47.3=1.2
Med-Replay
Hopper
100:10.7
101.3:0.6
87.3
96.9 2.6
99.6111.5
100.4=0.5
Med- Mee-RRppaay
Walker2d
75+44.3
95.5 1.5
66.8 3.1
84.4 + 4.1
80.7+-2.1
82.6 + 2.1
Average
81.8
88.0
82.5
86.6
82.6
87.5
This section will then delve into more detailed analysis of the performance differences among different baselines across various tasks. To provide an intuitive comparison of different algorithms, we classify them from three perspectives the classification method shown in Table Firstly, the Gym locomotion task has long horizon, dense rewards, and low action dimensions, making it baseline test task for offline RL. The results from Table 4 show that generative methods based on diffusion models generally perform better. The community currently attributes this to diffusion models' more powerful representation capabilities in modeling more complex policies or environmental models. However, Latent D ffuser does not demonstrate its advantages well in the low-dime now-dimensional action space. Although LatentDi fuser approaches the SOTA performance on this task, it is mainly due to a better diffusion sampling method, which is supported by the solid performance of the QGPO method. Due dense rewards, planning and hierarchy-barr based methods, such as TAP and HDMI, have not achieved the best results.
planning. hierarchy, and generative according to
Secondly, the Adroit task characterized by a high- dimeemiinnii action space. This leads the best performance for TAP and LatentDiffuser (see Table 2), two methods based on latent action, which experimentally verify the effectiveness of latent action. Additionally, generative methods based on diffusion models generally exhibit better performance. However, due to the shorter horizon of the Adroit task, the HDMI method, which is based on planning and hierarchy, does not achieve the best performance.
Lastly, the AntMaze task has longer horizon and very sparse rewards. This allows latent action ample room for improvement (see Table 3). Moreover, methods based planning and hicrarchy also achieve good results, such as HDML. In this task, non non-generative methods based on planning and hierarchy, such as ComPILE and GoFAR, approach the performance of generative methods without planning and hierarchy (D-QL).
The Performance Gap Between TAP and LatentDiffuser For TAP and LafentDiffuser, the performance gap between them on the expert dataset is smaller than on other datasets Adroit and Gym locomotion tasks. We analyzed that the primary source of this performance gap comes from
19

---

Published as a conference paper at ICLR 2024
the proportion of suboptimal trajectories in the dataset. In non-expert datasets, the proportion of suboptimal trajectories is more significant. To learn the optimal policy from the dataset, the algorithm needs to have the "trajectory stitch" ability, i.e., to splice segments of suboptimal trajectories to form an optimal trajectory.
On the one hand, most of the current offline RL methods are based on dynamic programming framework to learn a Q function. However, these methods require the Q function to have Bellman completeness to achieve good performance. Designing function class with Bellman completeness is very challenging (Zhou et al., 2023). On the other hand, Ajay et al. (2023, Appendix A.I) has found that generative methods based on diffusion models possess implicit dynamic programming capabilities. These methods use the powerful representation ability of diffusion models to bypass Bellman completeness and achieve the "trajectory stitch" ability. This allows them to perform well in datasets with more suboptimal trajectories.
LatentDi ffuser a generative method based on a diffusion model, while TAP is not. This leads to more significant performance gap between the two on non-expert datasets. In expert datasets, however, Latent Di fuser'' advantage cannot demonstrated.
E IMPLEMENTATION AND TRAINING DETAILS
In the following subsection, we delineate hyperparameter configurations and training methodologies employed in numeric experiments for both baseline models and the proposed LatentDi ffuser. Additionally, we supply references for performance metrics of prior evaluations conducted on standardized tasks concerning baseline models.
Each task undergoes assessment with total of 5 distinct training seeds, evaluated over span of 20 episodes. Adhering to the established evaluation protocols of TT (Janner et al., 2021) and IQL (Kostrikov et al., 2022), the dataset versions employed for locomotion control experiments are defined as v2, whereas v0 versions are utilized for remaining tasks.
E.1 BASELINE DETAILS
Before discussing the specific baseline implementation details, we first made simple comparison of all baselines in the 3 tasks from 3 perspectives: whether planning is introduced, whether it contains hierarchical structure, and whether generative learning is introduced, as shown in Table
Table 5:
Comparison of different baselines at three levels. means inclusive, O means exclusive,
and ! means a cheaper approximation.
ComPILE
CQL
IQL
D-QL
D-QL@1
QGPO
Diffuser
DD
Planning
O
O
O
Hierarchy
0
0
0
O
O
Generative
O
O
O
DT
TT
MoReL
HiGoC
GoFAR
TAP
HDMI
LafentDiffuser
Planning
0
Hierarchy
0
Generative
E.1.1
GYM LOCOMOTION CONTROT.
The results of CQL in Table and Table is reported in (Kostrikov al., 2022, Table 1);
The results of IQL in Table is reported in (Kostrikov al., 2022, Table 1);
The results of DT in Table is reported in (Chen et al., 2021, Table 2);
The results of TT Table and Table 4 reported in (Janner et al., 2021, Table 1);
The results of MoRel in Table 4 is reported in (Kidambi et al., 2020, Table 2);
The results of Diffuser in Table 4 is reported in the (Janner al., 2022, Table 2);
20

---

Published as a conference paper at ICLR 2024
The results of DD in Table 1 and Table is reported in the (Ajay et al., 2023, Table 1).
The results of D-QL in Table I and Table is reported in the (Wang et al., 2023, Table I).
The results of TAP in Table and Table reported in the (Jiang et al., 2023, Table 1).
The results of QGPO in Table and Table reported in the (Lu al., 2023, Table 2).
The results of HDMI in Table and Table is reported in the (Li al., 2023, Table 3),
E.1.2 ADROIT ENVIRONMENT
The results of CQL. in Table reported in (Kostrikov et al., 2022, Table 1);
The results of DD in Table 2 generated by using the offical repository" from the original paper (Ajay et al., 2023) with default hyperparameters.
The results of D-QL@1 in Table is generated by using the official repository from the original paper (Wang al., 2023) with default hyperparameters.
The results of QGPO in Table 2 is generated by using the offcial repository from the original paper (Lu al., 2023) with default hyperparameters.
The results of HDMI in Table 2 is generated by re-implementing the algorithm from the original paper (Li et al., 2023) with default hyperparameters.
The results of TT in Table 2 is reported in (Janner et al., 2021, Table 1);
The results of TAP in Table 2 reported in the (Jiang et al., 2023, Table 1).
is essential to mention that the D-QL employs resampling procedure for assessment purposes. To be more precise, during evaluation, the acquired policy initially produces 50 distinct action candidates, subsequently selecting a single action possessing the highest Q-value for execution. We empirically find that that this strategy is critical for achieving satisfactory performance in Adroit tasks. Nonetheless, the technique poses challenges in accurately representing the quality of initially sampled actions prior to the resampling procedure. Additionally, due to the high dimensionality of action, it incurs considerable computational overhead. As result, the resampling process has been climinated from the evaluation, utilizing single action candidate (referred to as D-QL @ 1) akin to QGPO (Lu et al., 2023).
E.
ANTMAZE ENVIRONMENT
The results of ComPILE, GoFAR, DD and HDMI in Table is reported in (Li et al., 2023, Table 2);
The results of HiGoC in Table 3 is generated by re-implementing HiGoC (Li et al., 2022) based on CQL and cVAE", and tune over the two hyparameters, learning rate E [3e 4, le 3) and the contribution of KL regularization [0.05, 0.2).
The results of D-QL@) @1 1 in Table 3 is generated by using the official repository from the original paper (Wang et al., 2023) with default hyperparameters.
The results of TAP in Table is generated by using the official repository from the original paper (Wang et al., 2023) with default hyperparameters.
The results of QGPO in Table 3 is generated by using the official repository from the original paper (Wang al., 2023) with default hyperparameters.
htttee///aaaaooccomm/aaaaa/aaaa. hhtt:://ithhb...
21

---

Published as a conference paper at ICLR 2024
E.2 IMPLEMENTATION DETAILS
The forthcoming release of the complete source code will be subject to the Creative Commons Attribution 4.0 License (CC BY), with the exception of the gym locomotion control, Adroit, and AntMaze datasets, which will retain their respective licensing arrangements. The computational infrastructure consists of dual servers, each possessing 256 GB of system memory, as well as a pair of NVIDIA GeForce RTX 3090 graphics processing units equipped with 24 GB of video memory.
E.2. REPRESENTATION LEARNING FOR LATENT ACTION
Encoder and Decoder. The score-based prior model is significantly inftuenced by the architec- ture of the bottleneck, encompassing both the encoder and decoder, and subsequently impacts its application for planning. Adhering the Decision Transformer Chen et al. (2021) and Trajectory Transformer (Janner et al., 2021), 1), TAP (Jiang et al., 2023) employs a GPT-2-style transformer in- corporating causal masking within its encoder and decoder. Consequently, information from future tokens does not propagate backward to their preceding counterparts. However, this conventional design remains prevalent in sequence modeling without guaranteeing optimality. For instance, one could invert the masking order in the decoder, thus rendering the planning gose-based.
detailed examination of the autoregressive (causal) versus simultaneous generation of optimal action sequences can be found in (Janner et al., 2022, $3.1). The discourse presented in Janner et al. (2022) remains germane to the context of latent action space. In particular, it reasonable to assume that latent action generation adheres to causality, whereby subsequent latent actions are contingent cecision-making or optimal control may exhibit
upon previous and current latent actions. However, anti-causality, the subsequent latent action may rely on future information, such future rewards. In general RL scenarios, the dependence on future information originates from the presumption of future optimality, intending to develop dynamic programming recursion. This notion is reflected by the future optimality variables Oz.T present in the action distribution log p (a, 8t, Our) (Levine, 2018). The aforementioned analysis lends support to the "diffusion- sawiiiio--smmppina-a-lallaa As result, causal masking is eliminated from the GPT-2 style encoder and state decoder design within the LatentDiffuse
framework.
Additionally, the action decoder is represented through the implementation of a 2-layered MLP, encompassing 512 hidden units and ReLU activation functions, effectively constituting an inverse dynamics model. Concurrently, a 3- 3-layered MLP, containing 1024 hidden units and ReLU activation functions, represents the reward and return decoder. The action decoder trained employing the Adam optimizer, featuring a learning rate of 2e 4 and batch size of 32 across 2e6 training steps. The reward and return decoder are also trained utilizing the Adam optimizer, however, with a learning rate of 2e- and batch size of 64 spanning 1e6 training steps.
Score-based Prior. Consistent with DD (Ajay et al., 2023), the score-based prior model is charac- terized by a temporal U-Net" (Janner et al., 2022) architecture encompassing a series of 6 recurrent residual blocks. Within each block, two sequential temporal convolutions are implemented, succeeded by group normalization (Wu & He, 2018), culminating in the application of the Mish activation function (Misra, 2019). Distinct 2-layer MLPs, each possessing 256 hidden units and the Mish acti- vation function, yield 128-dimelsional timestep and condition embeddings, which are concatenated and added to the first temporal convolution's 's activations within each block. We employ the Adam optimization algorithm, utilizing a learning rate of 2 x 10-4, a batch size of 32, and performing 2 x 10 training iterations. The probability, denoted by p. of excluding conditioning information 8 is set to 0.25, and K- = 100 diffusion steps are executed.
E.2.2 PLANNING WITH ENERGY-GUIDED SAMPLING
The energy guidance model is formulated as a 4-layer MIP containing 256 hidden units and leverages SiLU activation functions (Hendrycks & Gimpel, 2016). It undergoes training for 1 x 10 gradient- based steps, implementing the Adam optimizer with a learning rate of 3 x 10 and a batch size of 256. For gym locomotion control and Adroit tasks, the dimension of the latent actions set, M, is
Phtpps: :/ thub. hltbb..cmm/aanemeee fuuer..
22

---

Published as a conference paper at ICLR 2024
established 16, whereas for AntMaze tasks, it is fixed at 32. In all tasks, however, , is set to 3. The planning horizon, H, is set to 40 for gym locomotion control and Adroit tasks, and 100 for AntMaze tasks. The guidance scale, w, is selected from the range {1.2, 1.4, 1.6, 1.8), although the specific value depends on the task. A low temperature sampling parameter, a, is designated as 0.5, and the context length, C, is assigned value of 20.
F MODELING SELECTION
In this section, we shall explore the alterations elicited by diverse problem modeling approaches upon the Difteeeiffuse framework. Specifically, by examining various representations of latent action, we investigate the performance of the LatentDi fusee within both the raw action space and skill space. Subsequently, by comparing it against existing works, we distill some key observations of distinct diffusion sampling techniques on the efficacy of planning.
LatentDiffuser
Raw Action Space
Skill Space
Return Guided
E-MSE Loss
100%
Normalized Score
75%
50%
25%
0%
Gym HalfCheetah
Adroit Pen
AntMaze Play
Figure 4: Results of different modeling selections, where the height of the bar is the mean normalised scores on different tasks.
F.1 PLANNING IN THE RAW ACTION SPACE
The LatentDi fuser framework can be equally implemented within the raw action space; in this instance, merely substitute the latent diffusion model with any diffusion model for estimating raw trajectory distributions, such as those employed within the Diffuser (Janner et al., 2022) and Decision Diffuser (Ajay et al., 2023). To be precise, we can deduce the following theorem, akin to Theorem 1, within the raw action space:
Theorem (Optimal policy). Given an initial state $1, the optimal policy satisfies: r²(r | 81) xx p(( ) Σim Qc(er,oe), wherein 7 = (81, "1...... 8T, ar)), μ([ 81) represents the behavior policy and Qc(,--) refers the estimated Q-value function. signifies the inverse temperature controlling the energy strength.
By rewriting Po x² 4o / and To т, reformulate the optimal planning into the following diffusion sampling problem:
Po(to s1) (T $1) exp (-38(To,
(9)
where E(To, .1) ;= - ΣT (s az). Similarlly, the time-dependent energy model, fn (Tx. 81, 81, can then be trained by minizing the following contrastive loss:
(Ea(h(T())..) .1))
In .1.k))
min Ep(k.s1) 11. (1=0%)s))(((t)))
log M tn ( (TL2) .91
M 1=1
-Bo(((rr)) (.)))
Ej=1
(10)
23

---

Published as a conference paper at ICLR 2024
where & ~ (0, K), T = T oge, and c~ (0,, I). To estimate true distribution q $1) in Equation 10, we can utilize the diffusion model adopted in Diffuser or Decision Diffuser to generate M support trajectories {5o ( JMr for each initial state by diffusion sampling. The contrastive loss in Equation (10) then estimated by:
M
e -EEolb(f) .o:))
In (th .1.k))
min R.
log
(11)
M j= 1
.e1))
Σラ 5n((t s).))
t=1
where To (8) ()) correspond to the support trajectories for each initial state 8t. The training procedure shown in Algorithm
Algorithm 2 Efficient Planning in the Raw Action Space
Initialize the diffusion model Pe and the intermediate energy model f. for each gradient step do
b Training the diffusion model
Sample B1 trajectories T from offline dataset D
Sample B, Gaussian noises from (0, I) and B1 time k from U(0, K) Perturb to according ты = 7o o&e
Update {0) with the standard score-matching loss in Appendix B.2
end for
for each initial state in offline dataset 2 do
> Generating the support trajectories
Sample M support trajectories {f(+) yr from the pretrained diffusion model pe end for
for each gradient step do
Sample Bg initial state s1 from offline dataset D b Training the intermediate energy model
Sample B2 Gaussian noises from N I) and B2 time & from U1(0, K)
Retrieve support trajectories {ri (6) Jas for each
Perturb ri according ri (6) == og (4) + a&e
Update {n} based the contrastive loss in Equation (11)
end for
As shown in Figure 4, it becomes evident that pronounced performance degradation manifests in the raw action space when planning compared the latent action space, This phenomenon is even more pronounced in longer-horizon tasks, such as AntMaze.
Connection with Diffuser In Diffuser (Janner et al., 2022a), E(ro, $1) is defined as the return of To- Additionally, Diffuser uses a ween-sqqareeerron (MSE) objective to train the energy model (T, $1, t) and use its gradient for energy guidance (Lu al., 2023b). The training objective is:
min Bu [llf, (Te, $1, - (To. s1]1
(12)
"
Given the unlimited model capacity, the optimal fa satisfies:
f
MSE (Te, S1, E (rolmrn)))) [ (To, s1)]
(13)
However, according to Lu et al. (2023, $4.1), ) the true energy function satifies
E (To, $1) log Eu ETTo.11)]
(14)
(zo. [8 (To, 81)] f MMS (Tr, 81,
and the equality only holds when t. t 0. Therefore, the MSE energy function fMSE is inexact for all t > 0. Moreover, Lu et al. (2023) also shows that the gradient of f. fMSE is also inexact against the true gradience &, s1).
We replace the definition of E(2g, $1) in Lat ententDiff ffuseer with the return (or cumulative rewards) employed in Diffuser, culminating in the numerical results depicted in Figure 4. It is imperative to note that incorporating return into
contravenes Theorem 1. As discernible in the figure, conspicuous performance degradation ensues from the replacement. While return
24

---

Published as a conference paper at ICLR 2024
bears resemblance to cumulative state-action value, the latter demonstrates diminished variance. This advantage becomes increasingly pronounced in the longer-horizee task.
In addition, can implement alteration to the sampling method employed by the Diffuser, thereby rendering it to an exact sampling technique. More concretely, we add an exponential activation the original MSE-based loss in Equation (13), which is named E-MSE in (Lu et al., 2023b):
min S0, ||exp (fn (z,, t)) - exp (SE (zo))
In the LatentDiffuss we substitute contrastive loss with E-MSE, depicted in Figure 4. Although E-MSE belongs to the realm exact sampling methods, its inherent exponential terms precipitate significant numerical instability during training. Evidently, from Figure 4, the employment of E-MSE has culminated in a conspicuous decline in performance- finding that resonates with the conclusions drawn in the Lu et al. (2023b, SH).
F.2 PLANNING IN THE SKILL SPACE
The LatentDi ffuser framework can be equally implemented within other variants of the latent action space. Considering the following simplified trajectory Tsim of length T, sampled from an MDP with fixed stochastic behavior policy, consisting of sequence of states, and actions:
Tsim == ($1, a1 82, a2 ,,,...,,r)). ST)
(15)
Under this setting, the concept of latent action aligns perfectly with the definition of skill as delincated in prevailing works, although we persist in utilizing the LatentDi ffuser framework for efficient planning within the skill space.
Trajectory Decoder
State
Decoder
Encoder (Causal Transformer Position Embedding)
with Position Embolding
Action
+
Decocier
Fad
Fsu
Feedback Decoder
Decoder
Diffusian Model
Return
Decocier
Figure 5: Skill modeling with the score-based diffusion probabilistic model.
More concretely, can merely substitute the latent diffusion model with an variant of the latent diffusion model, as shown in Figure 5, for estimating simplified trajectory distributions. Similarilly, we also can deduce a theorem, akin to Theorem 1, within the skill space:
Theorem (Optimal skill-based policy). Given an initial state s1, the optimal skill-based policy satisfies: x (Tsim $1) x p(Tsim $1)e Et= Qc(sr,a), wherein p( 81) represents
the behavior policy and Qc(,,)) -) refers the estimated Q-value function. 2 0 signifies the inverse temperature controlling the energy strength.
Subsequently, we can employ the algorithm nearly identical to Algorithm for both the model training and the sampling of optimal trajectories.
To ascertain the efficacy of LatentDif fffieer fuser in skill space planning, we exchanged the latent diffusion model depicted in Figure 2 with the one shown in Figure 5. The experimental results can be observed in Figure 4. Evident from the illustration, a lack of encoding for future information (i.e., the reward and return) precipitates significant decline in skill space planning performance compared to that within latent action space. Similarly, this circumstance becomes increasingly pronounced in horieon tasks accompanied by sparse rewards.
25

---

Published as a conference paper at ICLR 2024
LatentDiffuser (L=3, H=40, beta=3, K=100, Causal Transformer) Latent Steps (L=1)
Latent Steps (L=4) Planning Horizon (H=10) Planning Horizon (H=50)
Inverse Temperature (beta Inverse Temperature (beta =0)
Diffusion Steps (K=30)
Diffusion Steps (K=200)
Non-causal Transformer
100%
Normalized Score
80%
60%
40%
20%
0%
Gvm Locomotion Control - Med Expert
Figure 6: Results of ablation studies, where the height of the bar is the mean normalised scores on different tasks.
G ABLATION STUDIES
In this section, ablation studies are performed on crucial hyperparameters within the LatentDi ffuser, specifically focusing the latent steps (i.e., the timesteps corresponding to the original trajectory the latent action), planning horizon (referring to the latent rather than the raw action), the inverse lemperature (impacting energy guidance intensity), and the diffusion steps (contributing to the reconstructed trajectory quality). Numerical experiments were conducted (refer Figure 6) on the Med-En Expeert dataset within the Gym locomotion control tasks, yielding the subsequent significant findings:
planning occurs in a latent action space featuring temporal abstraction. When single latent action is sampled, L transitional steps extending the raw trajectory be decoded. This design enhances planning efficiency as it reduces the number of unrolling steps to 1. hence, the search space size is exponentially diminished. Nonetheless, the repercussions of this design the decision-making remain uncertain. Therefore, we evaluated the LatentDi f fuser employing varying latent steps L. The red bars in Figure 6 demonstrate that the reduction in latent steps T. to leads to substantial performance degradation. We conjecture that this performance decline is attributed to VAE overfitting, as a higher prediction eтror has been observed with the reduced latent step, similar with TAP (Jiang et al., 2023)
Latent Steps The
Inverse Temperature The energy guidance effect is regulated by the inverse temperature; de- creasing values yield sampled trajectories more aligned with the behavior policy, while elevated values amplify the influence of energy guidance. As displayed by the pink bars in Figure 6, the LatentDi ffuterrffuurrs 's performance noticeably deteriorates when the inverse temperature is compa- rably low. Alternatively. a trivial decline occurs as the value substantially increases. We propose two plausible explanations: firstly, overwhelming energy guidance may generate discrepancies between trajectory distributions, guided by energy and induced by behavior policy, negatively impacting generated quality; secondly, the energy guidance originates from an estimated intermediate energy model, which is inherently prone to overfitting throughout training, leading to inaccuracies in the estimated energy and ultimately degrading the sampling quality.
Planning Horizon and Diffusion Steps
As demonstrated by the blue and yellow bars in Figure 6, the SetentDiffuse exhibits low sensitivity to variations the planning horizon and diffusion steps. Moreover, the conclusion regarding the planning horizon may be task-specific since dense- reward locomotion control may necessitate shorter-horizon reasoning than more intricate decision- making problems. The ablations of MuZero (Hamrick et al.. 2021) and TAP (Jiang et al., 2023) further reveal that real-time planning is not as beneficial in more reactive tasks, such as Atari and locomotion control.
26

---

Published as a conference paper at ICLR 2024
Transformers In the design of LatentDiffuser: we follow the settings of most existing generative methods (such as DT, TT. DD, TAP, etc.) for the encoder part, using GPT-2 style casual transformers for parameterization. Of course, in addition to this reason, another part of the reason is due to the modeling of latent actions. In the default setting, latent action consists of multiple timesteps of state, actions, rewards, and reward-to-go We helieve casual transformers will make the learned latent action representations more predictive. For example, predicting actions based on the state, predicting rewards and reward-to-go based on the state and action. This predictive ability, similar to model-based methods, will make the learned latent action representations more conducive high-quality planning. To verify this point, we conduct comparative experiments using non-causal transformers. Specifically, we remove the mask part of the causal transformer. This means that during the encoding and decoding process, we allow the model to use the information of the entire subtrajectory to reconstruct any element within that subtrajectory, such as the state using state information. As shown in Figure 6, the experimental results show that Lat Dattiiiffuu fusser has a significant performance degradation. Furthermore, is close to the performance when the latent step equals These results are consistent with our previous analysis, and when using non-casual transformer, the model is also prone overfitting, causing the learned latent action representations to contain less information, losing a certain degree of "predictability."
found that using non-casual transformer
TAP
DD
HDMI
D-QL
QGPO
LatentDiffuser
1.5
1
0.5
Runtime Spent on Single Decision (sl
Figure 7: The average runtime spent on a single decision of baselines based on the generative model.
Runtime To eliminate the influence of different algorithm implementation logic on runtime and focus on the model itself, we follow the settings of previous work and record the average time taken for different baselines to make the final action from the current input state 50 times. In the interest of fairness, we have only compared the runtime of generative methods. The final results are shown in Figure 7. As can be seen from the figure, the runtime of Latent fuser is at the average level, and the time required for making one decision is about 0.5 seconds, which is similar to DD. Although the sampling efficiency of the diffusion model has always been its weakness, we adopted the warm-up technique proposed by Diffuser, which can significantly shorten the sampling time without affecting performance. D-QI. has the most extended runtime, requiring multiple samplings (50 times) to select the best result. HDMI significantly increases runtime because it is two-layer method requiring two diffiusion samplings for making one decision. TT method has a longer runtime due to its tokenized data processing, which requires longer autoregressive sequence generation before generating an action. TAP and D-QL have the shortest runtimes, with the former using beam search for planning, which can be completed quickly with predetermined budget, but planning effectiveness is also constrained by the budget; the latter only needs to generate one-timestep action rather than sequence, so its runtime also shorter. However, final performance significantly lower due to the lack of planning step.
H LATENT ACTION VISUALIZATION
In order to gain a more intuitive understanding of the latent action space learned by LatentDi ffuser, this section presents a visualization of the latent actions and the corresponding trajectories obtained by decoding them. Specifically, we use the fully trained LatentDa D ffuser in the Hopper task to sample 5 trajectories and apply the t-SNE method to reduce the dimensionality of the latent actions associated with these trajectories for visualization, as shown in Figure 8. In the visualization of the trajectories, random trajectory is selected. To facilitate presentation (due to
27

---

Published as a conference paper at ICLR 2024
tinestops
timestapa
140 timesteps
200 Simandepe
900 timestops
150 Smostops
Figure 8: Visualization of latent actions and decoded trajectories.
the similarity in the robot shape at adjacent timesteps), we downsampled this trajectory by taking one latent action every 5 latent actions and overlaid the trajectory images obtained by decoding the adjacent 3 latent actions with different opacities achieve the trajectory displayed Figure 8. can be seen from the figure that the latent action space learned by Latentor feriffuss ffuse is a more compact action space, which some extent has learned a certain type of macro-action or "skill."
I MISSING DERIVATIONS
I.1 PROOF OF THEOREM
Proof. Previous works (Peters et al., 2010; Peng et al., 2019) formulate offline RL as constrained
policy optimization:
max EANDI
A((o,a)) Drc, ())la( 18))
where A is the action evaluation model which indicates the quality decision (s, a) by estimating the advantange function (=,s, a) Q (s. (s) of the current policy is inverse temperature coefficient. The first term intends to perform policy optimization, while the second term stands for policy constraint. It is shown that the optimal policy =" satisfies (Peng al., 2019):
x <a s) x p(a BAc(aa))
Since (*( has nothing do with the action a, the above formula can be further simplified to:
T (a s) x pla s) BAc(o.a) x p(a s) (Q(ra))
where Qç is the action evaluation model which indicates the quality decision (s, a) by estimating the Q-value function (=, of the current policy To simplify notation, here reuse Ç to parameterize the estimated Q Q-value function. Furthermore, we can extend the above conclusion from raw action space, single step level to latent action space, trajectory level:
28

---

Published as a conference paper at ICLR 2024
T-1
<' (1) - p((s))
I t=
x
(az
81)p(8411
8r, (r) Be, ae, 8 (++)) 81, at)
T-1
x <((8))
II
ulae
sr) exp
Rrr | st, at 1c+1)9 8t, az)
t=0
T-1
p(s1)
II
plae
se, ae. S41))((G I st,
Il 1=0
exp (3Qc
1=0
T-1
=m(z s1)
Σ (0))
=1(z s
1.2 PROOF OF THEOREM 2 AND THEOREM 3
The proof procedure of Theorem 2 and Theorem is similar with Theorem 1.
29

---

Published as a conference paper at ICLR 2024
REFERENCES FOR SUPPLEMENTARY MATERIAL
Suzan Ece Ada, Erhan Ozoop, and Emre Ugur. Diffusion policies for out-of-distribution generalization in offline reinforcement learning. arXiv preprint vrXiv:2307.04726 2023.
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 ICLR, 2023.
Arthur Allshire, Roberto Martín-Martin, Charles Lin, Shawn Manuel, Silvio Savarese, and Animesh Garg. Laser: Learning a latent action space for efficient reinforcement learning. In ICRA, 2021.
Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu. Equivariant energy- guided SDE for inverse molecular design. In ICLR, 2023.
Prajjwal Bhargava, Rohan Chitnis, Alborz Geramifard, Shagun Sodhani, and Amy Zhang. Sequence modeling is robust contender for offline reinforcement learning. arXiv preprint arXiv:2305. 145500 2023.
Adam Block, Daniel Pfrommer, and Max Simchowitz. Imitating complex trajectories: Bridging low-level stability and high-level behavior. arXiv preprint arXiv:2307. 126.223377.11999 2023.
Hongyi Chen, Yilun Du Yiye Chen, Joshua B. Tenenbaum, and Patricio A. Vela. Planning with sequence models through iterative energy minimization. In ICLR, 2023a.
Huayu Chen, Cheng Lu, Chengyang Ying. Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In ICLR, 2023b.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In NewrIPS, 2021.
Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Jianhao Wang, Alex Yuan Gao, Wenzhe Li, Liang Bin, Chelsea Finn, and Chongjie Zhang. Lapo: varinlev advantage-weighted policy optimization for offline reinforcement learning. NeurIPS, 2022.
Rohan Chitnis, Yingchen Xu, Bobak Hashemi, Lucas Lehnert, Urun Dogan, Zheqing Zhu, and Olivier Delallcau, Iql-td-mpe: Implicit q-learning for hicrarchical model predictive control. arXiv preprint arXiv:2306.00880 2023.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In ICLR, 2023.
Robert Dadashi, Léonard Hussenot, Damien Vincent, Sertan Girgin, Anton Raichuk, Matthieu Geist, and Olivier Pietquin. Continuous control with action quantization from demonstrations. In ICML, 2022.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6, 2005.
Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Russ R Salakhutdinov. Mismatched no more: Joint model-policy optimization for model-based In NeurIPS, 2022.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforvement learning. In NeurlPS, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. ICML, 2019.
Wonjoon Goo and Scout Niekum. Know your boundaries: The necessity of explicit behavioral cloning in offline rl. arXiv preprint 20X6v:2206.00695, 2022.
Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. In NeurlPS, 2022.
30

---

Published as a conference paper at ICLR 2024
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML, 2019.
Danijar Hafner, Timothy P Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In ICLR, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 04104, 2023.
Jessica B Hamrick, Abram L. Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Wither- spoon, Thomas Anthony, Lars Holger Buesing, Petar Veličković, and Theophane Weber. On the role of planning in model-based deep reinforcement learning. In ICLR, 2021.
Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xue- long Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. arXiv preprint arXiv:2305.1844 18459, 2023.
Dan Hendrycks and Kevin Gimpel. Gaussian enror linear units (gelus). arXiv preprint 1606:1660.084111 2016.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 2022.
Jifeng Hu, Yanchao Sun, Sili Huang. SiYu Guo, Hechang Chen, Li Shen, Lichao Sun. Yi Chang, and Dacheng Tao. Instructed diffuser with temporal condition guidance for offline reinforcement learning. arXiv preprint arXiv:23306. 04875, 2023.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In NeurIPS, 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In NeurlPS, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In ICML, 2022a.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In ICML, 2022b.
Zhiwei Jia, Fangchen Liu, Vineet Thumuluri, Linghao Chen. Zhiao Huang, and Hao Su. Chain-of- thought predictive control. arXiv preprint arXiv:2304. 00776, 2023.
Zhengyao Jiang. Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefenstette, and Yuandong Tian. Efficient planning in compact latent action space. In ICLR, 2023.
Siddharth Karamcheti, Albert Zhai, Dylan P Losey, and Dorsa Sadigh. Learning visually guided latent actions for assistive teleoperation. In L4DC, 2021.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In NeurlPS, 2022.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In NeurIPS, 2020a.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In NeurlPS, 2020b.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In ICLR, 2022.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In NeurlPS, 2020.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: 1805.00909 2018
31

---

Published as a conference paper at ICLR 2024
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint (iXiv:0005.0164333 2020.
Wenhao Li, Xiangfeng Wang, Bo Jin, and Hongyuan Zha. Hierarchical diffusion for offline decision making. In ICML, 2023.
Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdiffuser: Diffusion models as adaptive self-evolving planners. In ICML,, 2023.
Dylan P Losey, Hong Jun Jeon, Mengxi Li, Krishnan Srinivasan, Ajay Mandlekar, Animesh Garg, Jeannette Bohg, and Dorsa Sadigh. Learning latent actions control assistive robots. Autonomous robots, 46(1):115-147, 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. In NeurIPS, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.0009, 2022b.
Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact -hergy-guided diffusion sampling in offline reinforcement learning. In ICML, 2023a.
Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy-guidee diffusion sampling in offline reinforcement learning. In ICML, 2023b.
Cong Lu, Philip Ball. Jack Holder- Michael Osborne, and Stephen Roberts. Revisiting design choices in offline model based reinforcement learning. In ICLR, 2022c.
Kevin Lu, Aditya Grover, Pieter Abbeel, and Igor Mordatch. Reset-free lifelong learning with skill-space planning. In ICLR, 2021.
J Merel, I. Hasenclever, A Galashov, Ahuja, V Pham, G Wayne, Teh, and N Heess. Neural probabilistic motor primitives for humanoid control. In International Conference Learning Representations, 2019.
Diganta Misra. Mish: A self regularized non-monotonic neural activation function. arXiv preprint arXiv:190..08681 2019.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09335, 2020.
Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang. Metadiffuser: Diffusion model as conditional planner for offline meta-rl. In ICML, 2023.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In ICML, 2021.
Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals. Vector quantized models for planning. In ICMI., 2021.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In ICLR, 2023.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. wdvintage-weighted regression: Simple and scalable off- ooficolicy reinforcement learning. arXiv preprint arXiv: 170.0179, 2019.
Xue Bin Peng, Yunrong Guo, Lina Halper, Sergey Levine, and Sanja Fidler. Ase: Large-scale reusable adversarial skill embeddings for physically simulated characters. ACM Transactions On Graphics (TOG), 41(4):1-17, 2022.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In AAAI, 2010.
32

---

Published as a conference paper at ICLR 2024
Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. ICLR, 2023.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards-jus, map them to actions. arXiv preprint arXiv: /9/.00287. 2019.
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 leameed model. Nature, 588(7839):604-666 609, 2020.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In ICML, 2023.
HJ Suh, Glen Chou, Hongkai Dai, Lujie Yang, Abhishek Gupta, and Russ Tedrake. Fighting uncertainty with gradients: Offline reinforcement learning via diffusion score matching. arXiv preprint arXiv:2306. 14079, 2023.
Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In ICML, 1990.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS, 2017.
Siddarth Venkatraman. Latent skill models for offline reinforcement learning. Master's thesis, Carnegie Mellon University Pittsburgh, PA, 2023.
Jianhao Wang, Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-based imagination. In NeurIPS, 2021.
Linnan Wang, Rodrigo Fonseca, and Yuandong Tian. Learning search space partition for black-box optimization using monte carlo tree search. In NeurIPS, 2020a.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In ICLR, 2023.
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenherg, Scott Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, al. Critic regularized regression. In NeurlPS, 2020b.
Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. arXiv preprint arXiv:2307.02284, 2023.
Yuxin Wu and Kaiming He. Group normalization. ECCV, 2018.
Kevin Yang, Tianjun Zhang, Chris Cummins, Brandon Cui, Benoit Steiner, Linnan Wang. Joseph Gonzalez, Dan Klein, and Yuandong Tian. Learning space partitions for path planning. In NeurIPS, 2021.
Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. Egsde: Unpaired image-to-image translation via energy-guides stochastic differential equations. In NeurlPS, 2022.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In ICML, 2022.
Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforce- ment learning. In CoRL, 2021.
Zhaoyi Zhou, Chuning Zhu, Runlong Zhou, Qiwen Cui, Abhishek Gupta, and Simon Shaolei Du. Free from bellman completeness: Trajectory stitching via moodel-based eeturn-conditions supervised learning. arXiv preprint 200.:.3300.19308, 2023.
33


Reviews:
[{"review_id": "LzKSGbjOS5", "reviewer": "Reviewer_kDwN", "summary": "This paper presents a latent-action-based methods for reinforcement learning. This paper shows a unified framework for\ncontinuous latent action space representation learning and planning by leveraging latent, score-based diffusion models, brining theoretical connections. The experiments from the offline RL settings are pretty solid.", "questions": "I think in general the paper is quite self-contained", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well presented. The algorithm is straightforward and easy to understand. \n2. The paper includes some connection to the theoretical analysis.", "weaknesses": "1. The paper could be stronger if the author could provide some visualization/analysis what the latent diffusers learned", "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": "sVu3zCkWIq", "reviewer": "Reviewer_qvER", "summary": "The authors present a method that addresses offline RL using diffusion modelling. The authors combine several elements from recent work, namely using a latent representation of the path, while sampling latent variables using a diffusion model. The authors moreover make use of recent work that demonstrates effective sampling from energy based distributions. The authors demonstrate the utility of their approach on standard RL baselines and compare against a collection of other methods. They demonstrate improvement in particular for high dimensional tasks.", "questions": "Overall I enjoyed seeing the elements from previous work come together in this work. I think overall the authors did quite a good job writing this paper, although the overall notation, in particular where x and y get used, together with tau, and state tuples, can be a bit confusing. While the paper is a combination of previous work, in my opinion the algorithm the authors created is novel enough in itself to warrant publication. The authors demonstrate performance in their experiments and, importantly, provide a useful interpretation with their experiments to demonstrate specifically where they expect to make a difference. The current review system precludes 7 as a score, and until I read the rebuttal I am rounding down.\nSpecific points:\n - Can the authors comment on how the sequence length is determined? Is the sequence terminated externally?\nSmall stuff\n - A reference for the proof for theorem 1 is missing.\n - Is the text inconsistent with Fig2? It seems like the initial state in Fig 2 is labelled s0, while in the text there is only ever reference to s1 \n - On page 5 typo \"emploies\".", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The authors demonstrate improvement in the area where they expect their model to do well\n - The paper is well written \n - Elements seem to fit well together", "weaknesses": "- While I think this is not a very big deal, the paper relies mostly on existing methodologies", "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": "TwBbc8noWs", "reviewer": "Reviewer_A5bD", "summary": "This paper proposes the use of latent diffusion models for planning, especially in tasks with sparse rewards and long horizon. The key proposal is to use continuous latent action space representations for planning, through the use of latent score based diffusion models. The novel framework takes conditional generative models for planning in the latent domain, where the authors learn a latent action representation and use the latents for planning. The authors propose a multi-stage framework to address the issues of current hierarchical offline reinforcement learning algorithms. The framework utilizes diffusion models to learn latent actions from offline datasets. This paper argues the equivalence between planning with latent action representation and energy-guided sampling.", "questions": "Why does the proposed method work on par with the TAP baseline in the Expert dataset of the Adroit task?\n\nWhy was the GPT-2 style Transformer chosen over other choices of transformers as the encoder?\n\nWhat is the total runtime of the proposed method in comparison to the baselines?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper tackles offline RL tasks that require temporal abstraction, and proposes the use of latent action representations and latent actions for planning based on a diffusion based model. The task is formulated as a conditional diffusion problem, conditioning on returns. The paper seems to be well derived from an algorithmic viewpoint, and is one of the early works that seem to do planning based on a latent diffusion based approach. However, in its current form the paper is quite difficult to understand (see my points on weaknesses)\n\nThe proposed method scores better or on par with the baselines on multiple benchmarks with higher dimension action space and a long horizon. The paper presents the effectiveness of adding reward-to-go values in the formulation of a trajectory to learn sequence modeling.", "weaknesses": "The paper repeatedly uses terms like the latent action representation and latent action planning, without a carefully derived definition of it. For self consistency, it would be helpful to define these terms more concretely; otherwise, in its current format, the contributions of the paper can be hard to follow. \n\nI can see that the paper draws inspiration from the TAP paper (Jiang et al., 2023) and integrates a diffuser based latent step within this framework. The paper in its current form can be quite confusing to follow and difficult to understand. \n\nThe model seems overly complicated to understand and it is not clear what the contributions of the paper really are. I understand the relations with the TAP paper and how the work is trying to build from it, but I believe the planning approach based on diffusion doesn’t really have any novel approach? Am I right?\n\nOverall, equation 4 is the objective that needs to be optimzied, with the first term that generates the planning trajectory based on the latents from the diffuser? Am I correct? \n\nExperimental results are compared with few baselines using the D4RL setup. I believe all the results from 5.1 - 5.3 are using the same setup, although studying different aspects of the proposed algorithm? It probably would have been useful if there were much simpler proof of concept experiments, especially focusing on the planning part and how the diffuser plays a role there?\n\nThe experimental results need an explanation of performance in comparison to the baselines.\n\nThe paper requires more proofreading as it contains typos (such as “emploies”).", "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": "cjCrfcKuwz", "reviewer": "Reviewer_z1ps", "summary": "The authors propose an algorithm to leverage latent diffusion for planning skills in offline RL. Similar to previous works using latent diffusion, this work first trains a transformer encoder/decoder with weak KL regularization to model trajectory sequences, and then trains a diffusion prior to sample these latents. Then, an energy guided sampling method is proposed to plan in this latent skill space through guided diffusion sampling. The approach is evaluated on the D4RL benchmark, where it achieves strong results.", "questions": "1) Could the authors clarify why they expect a weakly regularized latent space to be useful for planning as compared to the raw actions? I see in the appendix that the comparison is made, and raw actions indeed are shown to perform worse. But I would like some intuition for why this should be the case, and if some better hyperparameter tuning could have produced similar results with raw actions.\n\n2) The authors use the return to go predictions from the decoder as Q values. Since TD learning is not done to learn optimal Q functions, these Q values correspond to those of the behavior policy. Is there a reason some form of offline Q learning technique such as CQL/BCQ/IQL was not done? Without doing Q learning, should we expect skill stitching to occur?\n\nI would also encourage the authors to cite a recent paper from Venkatraman et al. [1], which proposes a very similar algorithm. I am happy to recommend this paper for acceptance, primarily due to its strong empirical results and relevance to the current trends in offline RL.\n\n[1] Reasoning with Latent Diffusion in Offline Reinforcement Learning, arxiv.org/abs/2309.06599", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1) Leveraging latent diffusion for offline RL is a very natural extension to previous raw trajectory diffusion works such as Diffuser and Decision Diffuser. Diffusion planning latent space leverages the hierarchical structure of the planning problem very effectively, leaving the low level action prediction to the skill conditioned decoder instead.\n\n2) The experimental results on D4RL are strong and convincing.\n\n3) The work is highly relevant to the growing interest in viewing reinforcement learning through the lens of conditional generative modeling.\n\n4) The paper is written very clearly, with good motivation presented.", "weaknesses": "While I acknowledged the potential strength of hierarchical planning above, since the latent space is weakly regularized, we should not expect there to be too much benefit in leveraging this space for planning as compared to the raw trajectory space. Unlike the case of latent diffusion for image generative modeling where we simply require high quality samples with improved efficiency, here the primary motivation for leveraging latent planning is to plan with higher level skills.", "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)