Latent Plan Transformer for Trajectory Abstraction: Planning as Latent Space Inference
Authors: Deqian Kong
Dehong Xu
Minglu Zhao
Bo Pang
Jianwen Xie
Andrew Lizarraga
Yuhao Huang
Sirui Xie
Ying Nian Wu
Venue: NeurIPS 2024 poster (2024)
Keywords: Generative models
Reinforcement learning
Decision transformer
Primary area: reinforcement_learning
TL;DR: We present the Latent Plan Transformer (LPT), an novel generative model designed to explore planning and maintain temporal consistency through the latent space.

Abstract: In tasks aiming for long-term returns, planning becomes essential. We study generative modeling for planning with datasets repurposed from offline reinforcement learning. Specifically, we identify temporal consistency in the absence of step-wise rewards as one key technical challenge. We introduce the Latent Plan Transformer (LPT), a novel model that leverages a latent variable to connect a Transformer- based trajectory generator and the final return. LPT can be learned with maximum likelihood estimation on trajectory-return pairs. In learning, posterior sampling of the latent variable naturally integrates sub-trajectories to form a consistent abstrac- tion despite the finite context. At test time, the latent variable is inferred from an expected return before policy execution, realizing the idea of planning as inference. Our experiments demonstrate that LPT can discover improved decisions from sub- optimal trajectories, achieving competitive performance across several benchmarks, including Gym-Mujoco, Franka Kitchen, Maze2D, and Connect Four. It exhibits capabilities in nuanced credit assignments, trajectory stitching, and adaptation to environmental contingencies. These results validate that latent variable inference can be a strong alternative to step-wise reward prompting.

Latent Plan Transformer for Trajectory Abstraction: Planning as Latent Space Inference
Deqian Kongl*, Dehong Xul.,, Minglu Zhao Bo Pang', Jianwen Xie Andrew Lizarraga', Yuhao Huang', Sirui Xie Ying Nian Wu
Department of Statistics and Data Science, UCLA Salesforce Research Akool Research 'Xi'an Jiaotong University Department of Computer Science, UCLA
Abstract
In tasks aiming for long-term returns, planning becomes essential. We study gener- ative modeling for planning with datasets repurposed from offline reinforcement learning. Specifically, we identify temporal consistency in the absence of step- rewards as one key technical challenge. We introduce the Latent Plan Transformer (LPT),, novel model that leverages latent variable to connect Transformer- based trajectory generator and the final return. LPT can be learned with maximum likelihood estimation on trajectory retuum pairs. In learning, posterior sampling of the latent variable naturally integrates sub-trajectories tion despite the finite context. At test time, the latent variable inferred from an expected return before policy execution, realizing the idea of planning inference. Our experiments demonstrate that LPT can discover improved decisions from sub- optimal trajectories, achieving competitive performance across several benchmarks, including GGmmMuujoc Franka Kitchen, Maze2D, and Connect Four. It exhibits capabilities in nuanced credit assignments, trajectory stitching, and adaptation environmental contingencies. These results validate that latent variable inference can be strong alternative to step-wise reward prompting.
form consistent abstrac-
Introduction
Decision Transformer (DT) (Chen al., 2021) and some concurrent work (Janner al., 2021) have popularized the research agenda of decision-makin via generative modeling. The general idea is to consider docision-making e.g. the rewards or returns of trajectory, and outputs representation of the trajectory. Intuitively, purposeful eeciisionmmaiing process should shift the trajectory distribution towards regimes with higher returns. In the classical decision- -eaiinn literature, this is achieved by two interweaving processes, policy evaluation and policy improvement (Sutton and Barto, 2018). Policy evaluation promotes consistency in the estimated correlations between the trajectories and the returns. In DT, this is realized by the maximum likelihood estimation (MLE) of the joint distribution of sequences consisting of states, actions, and return-to-go (RTG). Policy improvement shifts the distribution to improve the status quo expectation of the returns. In DT, this is naturally entailed since the policy is distribution of actions conditioned on step-wise RTGs.
generative process that takes in representation of the task objective,
In this work, are interested in the problem of planning. Among various ways identify planning as a special class of decision -aainng problems, we pay particular attention to its data specification
"Equal Contribution
Project page: https://sites google. www.le.co...viww/ea.....o.oo www./iiwwwaatttttlaannnnnnn trenrtrt//.n.t...a.attttttteerrr Code: http:/gghhhh.ooommmmnuuuuuaa///a aar.aaeeeeeaa----Tannoeeeeee
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

and inductive biases. As designing step-wise rewards requires significant effort and domain expertise, we focus on the problem of learning from trajectory-retumm pairs, where trajectory is a sequence of states and actions, and the return is its total rewards. This design choice forces the agents to predict into the long-term future and figure out step-wise credits by themselves. A competitive Temporal Difference (TD) learning baseline, CQL (Kumar et al., 2020), was reported to be fragile under this data specification (Chen et al., 2021).
Our design of inductive biases reflects our intuition of a plan. While a policy is a factor of the trajectory distribution, a plan is an abstraction lifted from the space of trajectories. As a plan is always made in advance of receiving returns, it implies significance, persistence, and contingency. An agent should plan for more significant returns. It should be persistent in its plan even if the return is assigned in hindsight. It should also be adaptable to the environment's s changes during the execution of the plan. We formulate this hierarchy of decision-making with a top-down latent variable model. The latent variable we introduce is effectively plan, for it decouples the trajectory generation from the expected improvement of returns. The autoregressive policy always consults this temporally extended latent variable to be persistent in the plan. The top-down structure enables the agent to disentangle the variations in its plan from the environment's 's contingencies.
In this work, we introduce the Latent Plan Transformer (LPT), a novel generative model featuring a latent vector modeled by a neural transformation of Gaussian white noise, a Transformer- based policy conditioned on this latent vector and return estimation model. LPT is learned by maximum likelihood estimation (MLE). Given an expected return, posterior inference of the latent vector in LPT is an explicit process for iterative refinement of the plan. The inferred latent variable replaces RTG in the conditioning of the auto-regressive policy, providing richer information about the anticipated future. We further develop a mode-seeking sampling scheme that strongly enforce the temporal consistency for long-range planning, which is particularly effective in stitch trajectory, i.e., to compose parts of sub-optimal trajectories to reach far beyond (Fu et al., 2020), LPT demonstrates competitive performance in Gym-Mujoco locomotion, Franka kitchen, goal-reaching tasks in maze2d and antmaze, and a contingent planning task Connect Four. These empirical results support that latent variable inference can enable and improve planning in the absence of step-wise rewards.
2 Background
A sequential decision-making problem can be formulated with decision process (5, A, H, Tr, r, p) that contains set S of states and a set A of actions. Horizon H is the maximum number of steps the agent can execute before the termination of the sequence. We further employ S+ to denote the set of all non-empty state sequences within the horizon and A for action sequences likewise. Tr S x A + II(S) is the transition that returns a distribution over the next state. r : S A + E specifies the real-valued reward at each step. P : II(S) is the initial state distribution that is always uncontrollable to the agent. The agent's decisions follow a policy S+ x A+ + (A)).. In each episode, the agent interacts with the transition model to generate a trajectory ($1, 01. 82, a2..., 8H, (H).
The objective of sequential decision-making is typically formulated as the expected trajectory re- turn y = - Tt, Q = [v]- Conventional RL algorithms solve for a policy π <a, *). where the conditioning * denotes the optimal expected return. DT generalizes this policy to a (aglsgr, ac,, RTG by fitting the joint distribution p(8), 01, RTG aT, RTGr) with a Transformer. RTG, is the return-to- from step t to the horizon H, RTG; = E r(ssk, agx). It is a useful indication of future rewards, especially when rewards are dense and informative.
However, RTG becomes less reliable when rewards are sparse or have non-trivial relations with the return. Distributing the return to each step is a credit assignment problem. Consider an example of an ideal credit assignment mechanism: When students receive partial credits for their incomplete answers, it's more fair to give points equal the full marks minus the expected points for all possible ways to finish the answer, rather than assuming students have no knowledge of the remaining parts. This credit [. Here
assignment mechanism can be formalized as, RTG? = Σk= ((8g,, agx) (asx)+E((((ss)) Q can be estimated using deep TD learning with multi-step returns. Yamagata et al. (2023) instantiate Markovian version and demonstrate improvement in trajectory stiching.

---

Whatever credit assignment we use, be it RTG or RTG?, the puppse is to explicitly model the statistical association between trajectory steps and final returns. This effort is believed be necessary because of the exponential complexity of the trajectory space. This belief, however, can be re- examined given the success of sequence modeling. We explore an alternative design choice by directly associating the latent vector that generates the trajectory with the return.
Latent Plan Transformer (LPT)
3.1 Model
d1j 02
x N
Cross-attention
Un (zo)
(a())
Causal Transformer
p=(T(z)
Pr z)
== (81, 82,
811,
Figure 1: Left: Overview of Latent Plan Transformer (LPT). E I& is the latent vector. The prior distribution of 2 is a neural transformation of 20, i.c., 2 = U (2))) 20 ~ NN Given 2, T and y independent. PT(T(z) is the trajectory generator. Py (y(2)) 2) is the return predictor. Right: Illustration of trajectory generator pT(T(z).
Given variable-length trajectory т. a E Idd a vector that represents J in the latent space. y E R is the return of the trajectory. The joint distribution of the trajectory and its return is defined as p(T, v).
The latent trajectory variable z, conceptualized as a plan, is posited decouple the autoregressive policy and return estimation. From a statistical standpoint, with given, we assume that 7 and y are conditionally independent, positioning as the information bottleneck. Under this assumption, the Latent Plan Transformer (LPT) can be defined as,
Po (T. y. 2) Pa (T(())p((1a))p((lle))
(1)
where 0 (a. (, LPT approximates the data distribution Pdata y) using the marginal distribu- tion Po (T. = pa v. 2 22. It also establishes generation process,
Pa(:),
[r|=] ps(T)2),
[ul=] (vl=)>>(uul>))
The prior model p. (=) is implicit generator, defined as learnale neural transformation of an isotropic Gaussian, Ua (zn) and 20 ~ (0, Ia). U. (.) is an expressive neural network, such the UNet (Ronneberger al., 2015). This approach inspired by, yet contrasts with Pang et al. (2020), wherein the latent space prior is modeled as an Energy-based Model (EBM) (Xie et al., 2016). While EBM offers explicit unnormalized density, its sampling process is complex. Conversely, our model provides an implicit density with simpler sampling.
The trajectory generator p is conditional autoregressive model with finite context K, pa(T(z) (z) 11 (T(Tt))T(c-K).--. Tit-1): where T()) (st, It can parameter- ized by causal Transformer with parameter B, similar to Decision Transformer (Chen et al., 2021). Specifically, the latent variable is included in trajectory generation using cross- attention, shown in Fig. and controls each step of the autoregressive trajectory generation as (a St-K:tt Ot-K:t-1, 2). The action is assumed to follow a single-mode Gaussian distribution, i.c. a: N(gg(s:- --.-111. 2), I[a]).
The return predictor is non-linear regression on the latent trajectory variable 2, modeled as ps (ylz) (x), o²). ). It directly predicts the final return from the latent variable z. The function r2(2) is a small multi-layer perceptron (MIP) that estimates y based on z. The variance 2 is treated as the hyper-parameter in our setting.
3

---

3.2 Offline Learning
With of offline training examples {(T,. 1=1 we aim learn Latent Plan Transformer
(LPT) through maximum likelihood estimation (MLE). The log rigliike function is defined as
L(0) = Di= po w.). The joint probability of the trajectory and final return
Po(T, y) p=(Tl2 (z()))) (yl=
(3)
where po(zo) -V, Ia). The learning gradient of log-likelihood can be calculated according to V log (T v) (po(sol VV log log + Vo (=)(,
(4)
The full derivation of the learning method is in Appendix A.1. Let dex d3, dy represent the expected gradients L()) with respect to the model parameters at v,,3, 7, respectively. The learning gradients for each component are formulated as follows.
For the prior model (2)...
(T, v) (so(in))))) (a Ua(z))) V. logp> Un Ua(=)))]
For the trajectory generator,
53(T,1)= (aa Vogopp/(T (a(z)))).
For the return predictor,
(T, v) = E V, log = U. (20)))
Estimating these expectations requires Markov Chain Monte Carlo (MCMC) sampling of the posterior distribution pe(2oTT, T, y). We use the Langevin dynamics (Neal, 2011) for MCMC sampling, iterating as follows for target distribution <2)):
k+1 = + SV log ((*) + 22e**
(5)
where & indexes the time step of the Langevin dynamics. s the step size, and ex (O Ia) the Gaussian white noise. Here, T22) is instantiated as the posterior distribution pe (20 y). We have pe(zolT, v) x po(zo)p (T)2)) where Uz(z)). such that the gradient
Vzo log Po Po( +V lolp)(ul=) Σ aggregating finite context sub trajecturies
V log 1o
prior
return prediction
This demonstrates that the posterior inference of = is an explicit process of optimizing plan given its likelihood. the presence a finite context, Pa (T(2)) parametrized with Transformer can only account for sub-trajectee to--trajectories with maximum length of K. The latent variable serves as abstraction that integrates information from both the final return and sub-trajectories using gradients.
The sampling process starts by initializing k=0 from standard normal distribution A(. Ia). We then apply N steps of Langevin dynamics (e.g., N 15) to approximate the posterior distribution, making our learning algorithm an approximate MLE. For theoretical understanding of this noise- initialized finite-step MCMC, see Pang et al. (2020); Nijkamp et al. (2020); Xie et al. (2023). However, for large horizons (e.0.,/=1000), this method becomes slow and memory-intensive. To mitigate this, we adopt the persistent Markov Chain (PMC) (Tieleman, 2008; Xie al., 2016; Han al., 2017), which amortizes sampling across training iterations. During training, =²=0 is initialized from the previous iteration and the number updates is reduced to N = 2 steps. See Appendix A.2 for training and architecture details.
3.3 Planning Inference
The MLE learning of LPT gives us an agent that can plan. During testing, we first infer the latent 20 given the desired return using Bayes' rule,
20 (z0l)))) x po(zo)p, (=()))
(6)

---

This posterior sampling is achieved using Langevin dynamics similar to the training process. Specifi- cally, we replace the target distribution Eq. (5) with po po(2oly) ly) and run MCMC for fixed number of steps. Sampling from (= ()) eliminates the need for expensive back-propagation through the trajectory generator pa (T().
This posterior sampling of p(zoly) ly) is an explicit process that iteratively refines the latent plan increasing its likelihood given the desired final return. aligns with our intuition that planning an inference process. This inferred z, fixed ahead of the policy execution, effectively serves as a plan. At each step, the agent consults this plan generate actions conditioned on the current state and recent history, a2 pa(as) 8t-K:11: 1asss-::0---iee K 2 = (z6)).
Once a decision is made, the environment's (possibly non-Markovian) transition St+11 p(& ((8++1 (0t, 84) emits the next state. This sequential decision-making process iterates the sampling of 8t and Oe until termination the horizon.
Exploitation-incline Inference (EI) Inspired by the classifier guidance (CG) (Dhariwal and Nichol, 2021; Ho and Salimans, 2022) in conditional diffusion models, we introduce a guidance weight w to the original posterior in Eq. (6)
(z & Po (ul2)) = Ual(2))
which has the score 2., log Lo (zolv) = V log + +V w.,loogpp. (v12). This guidance weight w controls the interpolation between exploration and exploitation. When = 1, the sampled plans collectively represent the posterior density and account for Bayesian uncertainty, resulting in provably efficient exploration scheme (Osband and Van Roy, 2017). When w > 1, the sampled plans more concentrated around the modes of the posterior distribution, which are plans more likely to the agent. The larger the value of w, the more confident the agent becomes, and the stronger the inclination towards exploitation.
An overview of the algorithms for both offline learning and inference can be found in the following.
Algorithm 1 Offline learning
Input: Learning iterations T, initial parameters 0g = (ag. So: %)), offline training samples D (Tr,3: M- posterior sampling step size s, the number of steps N, and the learning rate 17o, 71. 172- Output: 0T
for / = to 7 do
yt) using Eq. (5) with N steps Posterior sampling: For each (Ti, V:), sample 2p ~ Po, (zo Ti, (zal
and step-size where the target distribution is Pe, (2o Ts, y1).
2.Lrn prior model Pa (z), trajectory generator (T(() and return predictor (y(z): 0141 0t+ + 70 E3 8 (T), Be41 = B,+m #/ E, (T 9c); 7641 = 7t # >, 8y (Ty, yi) as in Section 3.2. end for
Algorithm 2 Planning as inference
Input: Expected return y. trained model on offline dataset 6, posterior sampling step size s and the number of steps N, Horizon H and evaluation environment.
Output: Trajectory т if Exploitation-inclined Inference (EI) then
Sample 2g ~ po 20 ly) as in Eq. (7) using Eq. (5) with N steps and step size s, where the target
distribution is replaced by po(zoly) x (o(2))p (y/ Ua (20))" and U. (20).
else
Sample 2g Po Po(zaly) ly) as in Eq. (6) using Eq. (5) with N steps and step size s, where π is replaced by po (2oll) x (a(2o)P, (yl = (2)) and : = (20))
end if
while current time step t < H do Sample o, using trajectory generator as 0, pa(arlst-KKa K:t- (t:te//sK-Kk111 Once a decision is made, the environment's transition 8141 ~ p(se+1 lat, 84) emits the next state. end while
at ku-1, (zo)).
5

---

4 A Sequential Decision-Making Perspective
We approach the sequential decision-making problem with techniques from generative modeling. In particular, our data specification of trajectory-retumi pairs omits step-wise rewards, based on the belief that the step-wise reward function is only a proxy of the trajectory return. However, step-wise rewards are indispensable input to classical decision-making algorithms. Accumulating the rewards from the current step to the future gives us the RTG, which naturally hints the future progress of the trajectory. How is temporal consistency enforced in our model without the assistance of the RTGGs
Without loss of generality, consider the trajectory distribution conditioned y. The MLE objective is equivalent to minimizing the KI. divergence between the data distribution and model distribution, DKL (+% (7) [T) Here, PD denotes the data distribution and po denotes the model distribution. MLE upon autoregressive modeling imposes additional inductive biases by transforming the objective to DKL (P' AR (P), (T)), which is reduced next-token prediction for behavior cloning and transition model estimation:
single return value
H Σ
(a (1a,)))))))
behavior cloning
transition model estimation
However, behavior cloning is believed to suffer from drifting errors since it ignores covariate shifts in future steps (Ross and Bagnell, 2010). This concern is unique to sequential decision-making, as the agent cannot control the next state from stochastic environment, like generating the next text token.
This temporal consistency issue could be alleviated by additionally modeling the sequence of RTG. Denote = (RTG,, .T.G,,,..TG.... Modeling the joint distribution is to minimize
(T = DKL. (p) (T) (I)((T) + DRI (PTT))
Ep Σ H
1=1 (KI (ATGG)
(8)
RTG prediction
Note that the RTG prediction term conditioned on the entire trajectory, including the future steps. Minimizing this additional KL divergence correlates predicted RTTGs with hindsight trajectory-to-go
Our modeling of the latent trajectory variable 2 provides an alternative solution to the temporal consistency issue. Eq. (4) is minimizing the KL divergence
(K 2)p 2)) = DKI. (P) = =DKI AR (T) (R) + (g) [Dr (r (=+))))))))
+ (xs..))((())))))(((())))
plan prediction
where (=(7) = pT 2)/p" and 0 = 0 highlights these distributions have the same parameteri- zation as pÿ but wrapped with stop_groar -gradgrr operator when calculating gradients for (Han et al., 2017). Comparing Eqs. (8) and (9), it is now clear that = +plass a similar role as RTG in promoting temporal consistency in autoregressive models. Uniquely, 2xx (=|7) is the temporal abstraction intrinsic to the model, in contrast to step-wise rewards. From a sequential effectively plan that the agent is persistent to. From generative modeling perspective, from different trajectory modes would decompose the density pd 99:r,, a0:t-1,, relieving the burden of learning the autoregressive policy Pa (a, Su:4, ag.
Dacision-makinn perspective, 2
One caveat is that the transition model estimation should not be conditioned on y. Mixing up more trajectory regimes could provide additional regularization for its estimation and generalization. Actually, environment stochasticity is a more concerning issue for autoregressive behavior cloning, as highlighted by Yang et al. (2022); Paster et al. (2022); Strupl et al. (2022); Brandfonbrener et al. (2022); Villaflor et al. (2022); Eysenbach et al. (2022). Among them, Yang et al. (2022) pinpoints the issue by viewing RTGs as deterministic latent trajectory variables, closely related to what we present here. Uniquely, the latent variable in our model is inherently multi-mode (hence very non-deterministic) and ignorant of step- hise rewards. We postulate that the overfitting issue might be mitigated. This is validated by our empirical study inspired by Paster al. (2022).
Although RTG prediction and plan prediction both promote temporal consistency, they function very differently when mixing trajectories from multiple return-conditoondd regimes. RTG prediction is

---

supervised learning over the joint distribution PD( p). Simply mixing trajectories from multiple regimes can't 't encourage generalization to trajectories that are stitched with those the dataset. Yamagata et al. (2023) propose to resolve this by replacing RTG with RTGP. Intuitively, this aug- ments the distribution (T p) with pD (T',, where r' denotes trajectories covered by the offline dynamic programming, such as Q learning, and p = (RTG) .TGGG...Qu It significantly improves tasks requiring trajectory stitching. Conversely, plan prediction an unsupervised learn- ing as it samples from pD (T, (=)T,(=TT,)))) step-wise RTGs, trajectories lying outside of (D( P)) may be in-distribution for pD (T, ())pp(2lpp y)- The return prediction training further shapes the representation of 2, which can be benefited from denser coverage of y. With more return values covered, we may count on neural networks' strong interpolation capability to shift the trajectory distribution with y-conditioning.
As = contains more trajectory-relatte
information than
Related work
Decision-Making via Sequence Modeling Chen et al. (2021) propose Decision Transformer (DT), pioneering this paradigm shift. Concurrently, Janner et al. (2021) explore beam search upon the learned Transformer for model-based planning and inspired later work that searches over the latent state space (Zhang et al., 2022). Lee et al. (2022) report DT's capability in multi-task setting. Zheng al. (2022) explore the online extension of DT. Yamagata et al. (2023) augment the Monte Carlo RTG in DT with Q function and show improvement in tasks requiring trajectory stitching. Janner et al. (2022) explore diffusion models (Ho et al., 2020) as an alternative generative model family for decision-making. Our model differentiates from all above in data specification and model formulation.
Latent Trajectory Variables in Behavior Cloning Yang et al. (2022): Paster et al. (2022) investigate the DT's overfitting to environment contingencies and propose latent variable solutions. Our model closely related to theirs but unique in an EM-style algorithm for MIEE. Ajay et al. (2021); I Lynch et al. (2020) propose latent variable models to make Markovian policies temporally extended. Their models are more related VAE (Kingma and Welling, 2014).
Offline Reinforcement Learning Since the offline static datasets only partially cover the state transition spaces, efforts from a conventional RL perspective focus on imposing pessimistic biases to value iteration (Kumar et al., 2020; Kostrikov et al., 2021; Uehara and Sun, 2021; Xie et al., 2021; Cheng et al., 2022). Fujimoto and Gu (2021) show that simply augmenting value-based methods with behavior cloning achieves impressive performance. Emmons et al. (2021) report that supervised learning on eeturn-conditiones policies is competitive to value-based methods in offline RL. Our MLE objective is more related to the supervised learning methods. The latent variable inference further imposes temporal consistency, acting as a replacement of value iteration.
Hierarchical RL Methods like OPAL (Ajay et al., 2021), OPOSM (Freed et al., 2023) address TD- learning's 's limitations in long-range credit assignment using a two-stage approach: discovering skills from shorter subsequences to reduce the planning horizon, then applying skill-level CQL or online model-based planning on the reduced horizons. This paper focuses on comparing various methods for long-range credit assignment on the original horizon. Future work includes first discovering skills and then modeling them with a skill-level LPT further extend the effective horizon.
6 Experiments
The data specification of trajectory- pairs distinguishes our empirical study from most existing works in offline RL. Omitting step-wise rewards naturally increases the challenges in decision-making.
6.1 Overview
Our empirical study adopts the convention from offline RL. We first train our model with the offline data and then test an agent in the corresponding task. More training details and ablation studies of LPT can be found in Appendices A.2 and A.4.
OpenAl Gym-Mujoco The D4RL offline RI. dataset (Fu et al., 2020) features densely-rewarde locomotion tasks including Halfcheetah, Hopper, and Walker2D. We test for medium and medium- replay. It also includes Antmaze, a locomotion and goal-reaching task with extremely sparse reward.
7

---

The agent will only receive a reward of if hitting the target location and 0 otherwise. We use its umaze and umaze-diverse variants.
Franka Kitchen Franka Kitchen a multitask environment where a Franka robot with nine degrees of freedom operates within a kitchen setting, interacting with household objects to achieve specific configurations. Our experiments focus on two datasets of the environment: mixed, and partial, which consists of non-task-direttes demonstrations and partially task-directed demonstrations respectively.
Maze2D Maze2D a navigation task in which the agent reaches fixed goal location from random starting positions. The agent is rewarded 1 point when it is around the goal. Experiments conducted on three layouts: umaze, medium, and large, with increasing complexity. The training data of the Maze2D task contains only suboptimal trajectories from and to randomly selected locations.
Connect Four This a tile-based game, where the agent plays against a stochastic opponent (Paster et al., 2022), receiving at the end of an episode 1 reward for winning, O for draw, and -1 for losing.
Baselines We compare the performance of LPT with several representative baselines including CQL (Kumar et al., 2020), DT (Chen et al., 2021) 1) and QDT (Yamagata et al., 2023). CQL baseline results obtained from Kumar et al. (2020). QDT baseline results are from Yamagata et al. (2023). The DT results for al. (2022), and Kitchen implemented based on the published source code. CQL and DT results in the Connect Four experiments are from Paster et al. (2022). The mean and standard deviation of our model, shown LPT and L.TTEI,, are reported over seeds.
and Maze2D tasks are from Yamagata et al. (2023), Antmaze from Zheng
Table I: Evaluation results of offline OpenAl Gym MuJoCo tasks. We provide results for data specification with step-wise reward (left) and final return (right). Bold highlighting indicates top scores. LPT outperforms all final-retur baselines and most step- wise-reward baselines.
Step-wise Reward
Final Return
Dataset
CQL
DT
QDT
CQL
DT
QDT
LPT (Ours)
LPT-EI (Ours)
halfcheetah-mmddium
44.4
42.1
12.3
1.0
12.4
42.4
43.13 + 0.38
43.53 + 0.08
46.2
34.1
35.6
7.8
33.0
32.8
39.64 + 0.83
40.66 + 0.12
hopper-medium medium
58.0
60.3
66.5
23.3
57.3
50.7
58.52
1.92
63.83 : 1.47
hopper-men reellay
48.6
63.7
52.1
7.7
50.8
38.7
82.29 + 1.26
89.93 + 0.61
walker2d-medium
79.2
73.3
67.1
0.0
69.9
63.7
77.85 + 3.18
81.15 + 0.33
walker2d-medium- replay
26.7
60.2
58.2
3.2
51.6
29.6
72.31 + 1.92
75.68 + 0.34
51.0
22.3
17.2
61.9 1.22
64.7 0.51
kitchen-partinn
49.8
20.4
10.5
61.2 1.22+12 1.75
65.3 + 0.62 0.62
6.2 Credit assignment
When resolving the temporal consistency issue, our model doesn't 't have an explicit credit assignment mechanism that accounts for the actual contribution of each step. It not aware of the step-wirs rewards either. We are therefore curious about whether the inferred latent variable z can effectively assign fair credits to resolve compounding errors.
Distributing sparse rewards nigh-diimnsional actions The Gym-Muuoco environment was standard testbed for high-dimensional continuous control during the development of modern RL algorithms (Lillicrap et al., 2015). In this environment, step-wise rewards were believed to be critical for TD learning methods. In the setup of offline RL, Chen et al. (2021) reported the failure of the competitive CQL baseline when delaying step-wise rewards until the end of the trajectories. DT and QDT are reported to be robust to this alternation. As shown in Table 1, the proposed model, LPT,, outperforms these baselines when the data specifications are the same. Notably, LPT even excels in most of the control tasks when compared with the baselines with step-wise rewards.
Distributing delayed rewards to long-range sequences Maze navigation tasks with fully delayed rewards align with our intuition of planning problem, for it involves decision-makinn maainng at certain critical states absent of instantaneous feedback. An ideal planner would take in the expected total return and calculate the sequential decisions, automatically distributing credits from the extremely sparse and fully delayed rewards. According to Yamagata et al. (2023), DT fails in these tasks. Our
8

---

Table 2: Evaluation results of Maze2D tasks. Bold highlighting indicates top scores.
Dataset
CQL
DT
QDT
LPT
LPT-EI
Maze2D-umen
5.7
31.0
21.3
57.3
65.43 2.91
70.57 1.39
Maze2D-meddum
5.0
8.2 4.4
13.3 5.6
20.62 1.81
26.66 0.74
Maze2D-large
12.5
2.3
31.0 19.8
37.21 2.05
45.89 2.98
proposed model LPT outperforms QDT by a large margin in all three variants of the maze task. These results validate our hypothesis that the additional plan prediction KL imposes temporal consistency on autoregressive policies.
Table 3: Evaluation results of Antmaze tasks. Bold highlighting indicates top scores.
Dataset
CQL
DT
LPT
LPT-EI
Antmazz-umaze
74.0
53.3 + 5.52
80.8 + 4.83
92.4 + 0.80
Antmaze-umaze-divarse
84.0
52.5 + 5.89
78.5 + 1.66
84.4 + 1.96
Dataset
LPT-generated
Dataset
LPT-generated
TT
(a) Maze2d-Medium
(b) Maze2d-Large
Figure 2: (a) Maze2D-medium environment (b) Maze2D-large environment. Left panels show example
trajectories from the training set right panels show LPT generations. Yellow stars represent the goal states.
6.3 Trajectory stitching
In addition credit assignment, the setup of offline RL further presents challenge, trajectory stitching (Fu et al., 2020), which articulates the problem of shifting the trajectory distribution towards sparsely covered regimes with higher returns. the Franka Kitchen environment, both the mixed, and partial datasets contain undirected data where the robot executes subtasks that do not necessarily achieve the goal configuration. The "mixed" dataset contains no complete solution trajectories, necessitating that the agent Iearn to piece together relevant sub-trajectories. A similar setting happens as an example, in the training set, the average return of
in Maze2D domain. Taking Maze2D-modine all trajectories is 3.98 with standard deviation of 10.44, where the max return is 47. DT's score is only marginally above the average return. Yamagata et al. (2023) attribute DT's failure in Maze2D to its difficulty with trajectory stitching.
Fig. visualizes samples from the training data and successful trajectories in testing. The left panels show that trajectories in training are sub- optimal in terms of (1) being short in length and (2) containing very few goal-reaching in- stances. Trajectories the right are generated by random runs with LPT, where the agent successfully navigates to the end goal from ran- dom starting positions an effective manner. This indicates that the agent discover the relation between different ys to facilitate such stitching.
Training Xg
Training Testing
Testing Zo
Figure t-SNE plot of latent variables in the Maze2D-medium Left: Training 20 from aggre- gated posterior (T))().)) (po(2n r, v-] Testing 20 from po (5o y), disjoint from training population. Right: Distribution of z = (z)).
To probe into the agent's understanding of trajec- tories' returns, visualize the representation

---

space of the latent variables. The left of Fig. 3 is the aggregated posterior distribution of 20- We can see that 20 infered from Po(2)l3) yy) are distant away from the training population. The agent understands they are not very likely in the training set. The right of Fig. 3 is the distribution of which is transformed from Eg with the UNet, 2 (2)) We observe that 28 from the generated trajectories become "in-distribution" in the sense that some of them are mingled into the training population and the remaining lie inside a region coverable through linear interpolation of training samples. The agent understands what trajectories to generate even if they are unlikely among what it has seen.
6.4 Environment contingencies
To live in a stochastic world, contingent planning that is adaptable to unforeseen noises is desirable. Paster al. (2022); Yang et al. (2022) discover that DT's performance would degrade stochastic environments due to inevitable overfitting towards contingencies. We examine LPT and other baselines in Connect Four from Paster et al. (2022). Connect Four a two-player game, where the opponent will make adversarial moves to deliberately disturb an agent's 's plan. According to the empirical study from Paster al. (2022), the degradation of DT is more significant than in stochastic Gym tasks from Yang et al. (2022). As shown in Table 4, LPT achieves the highest score with minimal variance. The ESPER baseline is from Paster et al. (2022), which is very relevant to LPT as it is also latent variable model. ESPER learns the latent variable model with adversarial loss. It further adds a clustering loss in the latent space. LPT's on-par performance may justify that MLE upon more flexible prior can play an equal role.
Table 4: Evaluation results on Connect Four. Bold highlighting indicates top scores.
Dataset
CQL
DT
ESPER
LPT
Connect Four 0.61 + 0.05
0.8 + 0.07
0.99 0.03
0.99 + 0.01
7 Limitation
We omit the Antmaze-large experiment from the main text and included potential reasons for LPT's unsatisfactory performance in Appendix A.3. Another interesting direction is to study LPT's continual learning potential. During planning, LPT explores with provably efficient posterior sampling (Osband et al., 2013; Osband and Van Roy, 2017).
8 Summary
We study generative modeling for planning in the absence of step-wise rewards. We propose LPT which generates trajectory and return from latent variable. In learning, posterior sampling of the latent variable naturally gathers sub-trajectories to form an episode-wise abstraction despite finite context in training. In inference, the posterior sampling given the target final return explores the optimal regime of the latent space. It produces latent variable that guides the autoregressive policy to execute consistently. Across diverse evaluations, LPT demonstrates competitive capacities of nuanced credit assignments, trajectory stitching, and adaptation to environmental contingencies. Contemporary work extends LPT's application to online molecule design (Kong et al., 2024). Future research directions include studying online and multi-agent variants of this model, exploring its application in real-world robotics, and investigating its potential in embodied agents.
Acknowledgements
The work was partially supported by NSF DMS-2015577 NSF DMS-241522 and a gift fund from Amazon. We sincerely thank Mr. shanwei Mu and Dr. Jiajun Lu at Akool Research for their computational support, well as the anonymous reviewers for their valuable feedback.
10

---

References
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofrr Nachum. Opal: Offline primitive discovery for accelerating offline reinforcement learning. In International Conference Learning Representations (ICLR), 2021.
David Brandffobrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does retur-conditioons supervised learning work for offline reinforcement learning? In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pages 1542-1553, 2022.
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 Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 15084-15097, 2021.
Ching-An Cheng, Tengyang Xie, Nan Jiang, and Alekh Agarwal. Adversarially trained actor critic for offline reinforcement learning. In International Conference on Machine Learning (ICML), pages 3852-3878, 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 8780-8794, 2021.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:21 2021.
Benjamin Eysenbach, Soumith Udatha, Russ R Salakhutdinov, and Sergey Levine. Imitating past suc- cesses can be very suboptimal. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pages 6047-6059, 2022.
Benjamin Freed, Siddarth Venkatraman, Guillaume Adrien Sartoretti, Jeff Schneider, and Howie Choset. Learning temporally abstractworld models without online ex experimentation. In International Conference Machine Learning (ICML), pages 10338-10356, 2023.
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 and Shixiang Shane Gu. minimalist approach offline reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 20132-20145, 2021.
Tian Han, Yang Lu, Song-Chun Zhu, and Ying Nian Wu. Alternating back-propagation for generator network. In AAAI Conference Artificial Intelligence (AAAI), pages 1976 1984, 2017.
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 Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 6840-6851, 2020.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 1273-1286, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning (ICMI.,, pages 9902-9915, 2022.
Diederik P. Kingma and Max Welling, Auto-encoding variational bayes. In International Conference on Learning Representations, (ICLR), 2014.
Deqian Kong, Yuhao Huang, Jianwen Xie, Edouardo Honig, Ming Xu, Shuanghong Xue, Pei Lin, Sanping Zhou, Sheng Zhong, Nanning Zheng, and Ying Nian Wu. Molecule design by latent prompt transformer. In Advances in Neural Information Processing Systems (NeurIPS), 2024.

---

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations (ICLR), 2021.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 1179-1191, 2020.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, al. Multi-game decision transformers. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pages 27921-27936, 2022.
Timothy Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint (xX.v:1509.0220 2015.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on robot learning (CoRL), pages 1113-1132, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529- 53399553332200 2015.
Radford M Neal. MCMC using hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2, 2011.
Erik Nijkamp, Bo Pang. Tian Han, Linqi Zhou, Song- Song-Chun Zhu, and Ying Nian Wu. Learning multi-layer latent variable model via variational optimization of short run meme for approximate inference. In European Conference on Computer Vision (ECCV), pages 361-378, 2020.
Ian Osband and Benjamin Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In International Conference on Machine Leurning (ICML), pages 2701-2710, 2017.
Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. In Advances in Neural Information Processing Systems (NIPS), volume 26, 2013.
Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't 't count on luck: Why decision transformers and rvs fail in stochastic environments. In Advances in Neural Information Processing Systems (NeurIPS), volume 35. pages 38966-38979, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical Image Computing and Computer- Assisted Intervention (MICCAI), pages 234-241, 2015.
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In International Conference on Artificial Intelligence and Statisties (AISTATS), pages 661-668, 2010.
Miroslav Strupl, Francesco Faccio, Dylan Ashley, Jürgen Schmidhuber, and Rupesh Kumar Srivastava. Upside-down reinforcement learning can diverge in stochastic environments with episodic resets. arXiv preprint arXiv:2205.06599, 2022.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. International conference on Machine learning (ICML), pages 1064 1071, 2008
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. In International Conference Learning Representations (ICLR), 2021.
12

---

Adam R Villaflor, Zhe Huang, Swapnil Pande. John M Dolan, and Jeff Schneider. Addressing optimism bias in sequence modeling for reinforeement learning. In International Conference on Machine Learning (ICML), pages 12200222282 2022.
Jianwen Xie, Yang Lu, Song-Chun Zhu, and Ying Nian Wu. A theory of generative convnet. In International Conference on Machine Learning (ICML), pages 2635-2644, 2016.
Jianwen Xie, Yaxuan Zhu, Yifei Xu, Dingcheng Li, and Ping Li. A tale of two latent flows: Learning latent space normalizing flow with short-run langevin flow for approximate inference. In AAAI Conference on Artificial Intelligence (AAAI), pages 10499-10509, 2023.
Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistent pessimism for offline reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 6683-6694, 2021.
Taku Yamagata, Ahmed Khalil, and Raul Santos Rodrriuz. Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl. In International Conference on Machine Learning (ICML), pages 38989-39007, 2023.
Sherry Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Dichotomy of control: Sepa- rating what you can control from what you cannot. In International Conference on Learning Representations (ICIR), 2022.
Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefenstette, Yuandong Tian, et al. Efficient planning in a compact latent action space. In International Conference on Learning Representations (ICLR), 2022.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In International Conference on Machine Learning (ICML), pages 27042-27005 2022.
A Appendix
A.1 Details about model and learning
Given trajectory T, Rd is the latent vector to represent the variable-length trajectory, y E R is the return of the trajectory. With offline training trajectory-retum pairs {(T,, 1,..., n). The log-likelihood function is L(0) Σ lo M1). with learning gradient VoL(
Σ Vo log (o yi). We derive the form of Vo log po (T: st), proving Eq. (4) below, dropping index subscript, for simplicity.
V po(p Vopo(T.V)
po(T,
/ Vo Un (20)))
Po(T, y)
po(T, U. (z)))
V Volog v, = (z(z)))
po (o(T.
/ ppe(zol= F, Vo (T (z(z)))
V lo U. (z)))
log Vol Vo log
(
V Velogps(T(V.<<z)) Vo (glo
A.2 Training details
For Gym- -ujoco offline training, as shown in Table 5, most of the hyperparameters were shared across all tasks except context length and hidden size. However, due to the significant variations in the scale of the maze maps and the lengths of the trajectories within the Maze2D umaze, medium, and large nategories-moot sizes were adjusted accordingly accommodate these
13

---

differences, where the detailed setting can be found in Table 6. We also show the parameters for Franka Kitchen environment Table 7 and Connect Four Table
Training time for the Gym-Mujoco tasks using single Nvidia A6000 GPU is 18 hours on average. We train Maze2d tasks using single Nvidia A 100 GPU using 30 hours average. Kitchen tasks using a single Nvidia A6000 GPU takes 60 hours on average. Connect-4 on a single Nvidia A6000 GPU takes 10 hours.
Table 5: Gym-Muuoco Environments LPT Model Parameters
Parameter
HalfCheetah
Walker2D
Hopper
AntMaze
Number of layers
3
3
3
Number of attention heads
1
Embedding dimension
128
128
128
192
Context length
32
64
64
64
Learning rate
le-4
le-4
le-4
1e-3
Langevin step size
0.3
0.3
0.3
0.3
Nonlinearity function
ReLU
ReLU
ReLU
ReLU
Table 6: Maze2D Environments LPT Model Parameters
Parameter
Umaze
Medium
Large
Number of layers
4
Number of attention heads
8
4
Embedding dimension
128
192
192
Context length
32
64
64
Learning rate
le-3
le-3
2e-4
Langevin step size
0.3
0.3
0.3
Nonlinearity function
ReLU
ReLU
ReLU
Table 7: Franka Kitchen Environments LPT Model Parameters
Parameter
Mixed
Partial
Number of layers
4
3
Number of attention heads
4
16
Embedding dimension
128
128
Context length Learning rate
16
16
1e-3
1e-3
Langevin step size
0.3
0.3
Nonlinearity function
ReLU
ReLU
Table 8: Connect LPT Model Parameters
Parameter
Value
Number of layers
3
Number of attention heads
4
Embedding dimension
128
Context length
4
Learning rate
le-3
Langevin step size
0.3
Nonlinearity function
ReLU
14

---

Dataset antrnazeelarggeddddereeee Trajectory length vs return
0.040
Return
Return
0.035
0.025
Density
0.020
0.015
0.010
0.000
200
400
600
800
1000
1200
Trajectory length
Figure 4: Trajectory length and return distribution in dataset Antmaze-largeddivrrse
Trajectory returns for maze2d-large-v1
17500
15000
12300
maggancc
10000
7500
5000
2300
Trajectory lengths for maze2d-arre-v1
300
400
300
200
100
Trajectory length
Figure 5: Trajectory length and return distribution in dataset Maze2D-Large
A.3 Discussion data quality of Antmaze medium and large
In our experiments, encounter a curious phenomenon that LPT outperforms CQL, DT and QDT in Antmaze-umant umaaz by large margin but falls behind in Antmaze-large. Upon closer examination the data from D4RL, we gained valuable insights into the potential reasons behind LPT's performance on this task.
Fig. 4 plots the distributions of final returns and the trajectory lengths. Surprisingly. this dataset consists of 5448 trajectories (75.8%%) with length=1, 893 trajectories (12.43%) with enggt=11000 and only 841 trajectories (11.71%) with lengths in between. Such a biased trajectory coverage can be detrimental to sequence models like LPT,, which learn to make decisions by discovering correlations between trajectories and returns.
15

---

As reference, Fig. 5 shows the distributions of final returns and the trajectory lengths of Maze2D- large, a task where LPT performs well. It is important to note that TD-leariing methods, such as CQL and QDT, rely solely on (s, a, s,,''rr tuples and are less affected by the trajectory length distribution in the dataset. Consequently, Antmaze-large in D4RL, remains fair dataset for these methods to perform offline RI..
A.4 Ablation study
We investigate the role of the expressive prior (z) in our Latent Plan Transformer (LPT) model by removing the UNet component, which transforms 20 from non-informative Gaussian distribution. Table reports the results on three Gym-Mujoco tasks and Connect Four. We observe that the performance of LPT drops in all environments when the UNet is removed. For example, in the stochastic environment Connect Four, LPT's performance decreases from 0.99 to 0.90, while the baseline Decision Transformer (DT) without latent variables achieves 0.80. These results indicate that a more flexible prior benefits the learning and inference of LPT.
Table 9: Ablation study results on Gym-Mujoco tasks and Connect Four.
Dataset
DT
LPT
LPT w/o UNet
halfcheetab-medium---ee medium-replay
33.0 = 4.8
39.64 + 0.83
34.70 = 1.58
hopper-mediune maapprr-edium-replayy replay
50.8 + 14.3
71.17 3.01
53.41 = 6.95
walker2d-medium-ree
51.6 - 24.7
72.31 + 1.92
56.88 + 4.20
Connect Four
0.8 + 0.07
0.99 = 0.01
0.90 + 0.06
To further explore the impact of the prior, we conducted additional experiments testing the effects of different UNet configurations on LPT's performance. Table 10 shows the normalized scores on the walker2d 2ddmmdddieee -e--medduuueeee um-replay task with various UNet architectures. We observe that reducing the capacity expressiveness of the UNet (e.g., smaller dimension, fewer multipliers, smaller initial convolution, or fewer ResBlocks) consistently degrades performance, though still outperforming the model without the UNet prior. This suggests that a more expressive prior enhances LPT's ability to model complex policies.
Table 10: Effect of different UNet configurations on LPT performance.
Model Prior
Normalized Score
UNet (original)
72.31 + 1.92
UNet w/ smaller dimension
64.06 + 1.94
UNet w/ fewer multipliers
64.59 = 1.54
UNet w/ smaller initial convolution
70.49 + 2.84
UNet w/ single ResBlock
67.95 + 4.64
No UNet (Standard Normal prior)
56.88 + 4.20
Our results underscore the crucial role of the learned prior in LPT's performance. The original UNet configuration achieves the highest normalized score, indicating that our current UNet design optimal among the variants tested. We appreciate the reviewer's suggestion, as it prompted to perform a more detailed analysis of the prior's impact on LPT.
A.5 Continual learning with online data
We are also interested in LPT's potential in finetuning even continual learning. Inspired by ODT (Zheng et al., 2022), we employ trajectory replay butter (Mnih et al., 2015) to store samples from online interaction
a firstin--irst-ouu manner. After the completion of each episode, we update LPT with the same learning algorithm as with the offline data. Note that ODT introduces some techniques additional to DT. In contrast, LPT explores with the provably efficient posterior sampling (Osband et al., 2013; 30 Osband and Van Roy, 2017). We report the results in Table I.. Despite the significance in a few tasks, the improvement within standard deviation of the mean for the majority. We observe similar pattern in ODT (Zheng al., 2022).
16

---

Table 11: Evaluation results of online Open AI Gym MuloCo and Antmaze tasks. ODT baselines are sourced from Zheng et al. (2022). Our results are reported over 5 seeds.
Step-wise Reward
Final Return LPT
Dataset
ODT
A
A
halfcheetah-meddee
42.16 1.48
-0.56
43.26 0.59
0.13
halfcheetah sertcheeth--medeee
medium-replay
40.2 - 1.61
0.43
40.63 = 0.28
0.99
hopper-mediie
97.54 + 2.1
30.59
64.84 10.29
6.32
hopper-meddim-replay
88.89 + 6.33
2.25
+
1.27
walker2d-medium
76.79 2.30
4.60
79.54
1.69
walker2dd-mcdium-rrpla
76.86 + 4.04
7.94
78.99 3.84
6.68
Antmaze-umaze
88.5 + 5.88
35.4
83.5 + 3.28
22.3
Antmaze-diverse
56.00 5.69
5.8
75.6 + 2.8
8.0
A.6 Broader Impact
We introduce novel generative model that leverages the idea of "planning inference", which potentially leads to more capable and efficient Al systems that make better decisions. One potential negative impact could come from the misuse of the generative modeling designing objects or entities for harmful purposes. Developing suitable safeguards and regulations will be crucial to mitigate potential negative impacts.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: See Section 3, Section and Section 6
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See Section 7
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
18

---

Justification:
See Section 3, Section 4 and Appendix A.1
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Section Appendix A.2.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
19

---

Answer: [Yes]
Justification: Code is released.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
Sonideee/Cooeuumiisseeee oniooiccy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ //iips.cooo wwtpccccpppccceeeeccec.cccccccct haabblicggiids..... /CoodSubmbia
/ ceobbmissiie ssinnn oniioniyl) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See Section 6, Appendix A.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: See Section 6, Appendix A.4, and Appendix A.5
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
20

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix A.2
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that NeurIPS Code of Ethics https: : /ttp///eerrp/.ee/ee//ee cc ct/pcc//pp/....ccc httpppbicccccchccccceeeeeee.eeee
didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification: None.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Appendix A.6.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
21

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: None.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: None.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
wnwerswwthcodeee
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
22

---

Answer: [NA]
Justification: None.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/coddemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: None.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: None.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
23


Reviews:
[{"review_id": "BejgPaw3jj", "reviewer": "Reviewer_uuQw", "summary": "Building on the idea of decision transformer, the paper introduces a new generative model based decision-making agent called Latent Plan Transformer (LPT). Instead of directly generating the trajectories and returns as in the prior work, LPT would first generate a latent vector, and then generate the trajectory and its return conditioned on the latent vector. The idea of introducing this latent vector is to view it as a plan which provides the agent a temporal consistency guideline in the long decision-making process. Experimental results show improved performance of LPT over existing baselines in robotic and navigation tasks.", "questions": "- Are there derivations for Eq. (9)? Why is the second term in (9) are between the two parameters of the same model and one with stop_grad( )?\n\n- Is the comparison with baselines fair in the sense that all the models have similar parameter sizes?", "limitations": "Yes, the authors addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea of having a plan for decision-making is very interesting. Discussion in Section 4 on some intuitive reasons why having a plan may help in long range problems when temporal consistency is an issue.\n\n- Experimental results show strong performance in various tasks. Especially in the maze domains with long-range delayed rewards, the performance improve from DT to LPT highlights the benefits of the latent plan.", "weaknesses": "- Although there are motivations and some high-level discussions provided, it would probably be more convincing if those ideas have connection to some theoretical analysis in MDP representation learning.\n\n- The ablation study in A.5 comparing different latent prior seems important, but little discussion is provided. Intuitively, DT may be view as LPT with a trivial prior. Are there more comparisons with different priors for the latent vector? Since there is a UNet used in LPT, LPT may end up having much more parameters than DT. Are the sizes of the models adjusted to account for this extra prior network in LPT?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing the importance of our approach to modeling plans in decision-making processes! We greatly appreciate your insightful feedback and the opportunity to clarify and expand upon our work.\n\n> W1: What are the connections between LPT and some theoretical analysis in MDP representation learning?\n\nThank you for bringing up this interesting perspective! While there are indeed some conceptual similarities, we believe it is important to highlight a few key differences. Traditional MDP representation learning typically focuses on learning state representations that capture the MDP structure, including state transitions and reward functions, often operating on a step-by-step basis. In contrast, LPT doesn't explicitly learn the MDP structure. Instead, our latent variable $z$ serves as a policy representation, encoding information about the entire trajectory. This fundamental difference in focus - state representation versus policy representation - leads to distinct approaches. While MDP representation learning aims to facilitate efficient learning of optimal policies through compact state representations, latent variable $z$ in LPT directly informs action selection throughout the episode. This episode-level approach in LPT can be seen as a temporal abstraction over the trajectory, instead of over a timestep, as in MDP representations. Despite these differences, we acknowledge that both approaches share the goal of finding useful abstractions that facilitate decision-making in sequential problems. \n\n\n> W2.1: A.5 ablation study important but lacks discussion -- DT as LPT with a trivial prior. Are there more comparisons with different priors for the latent vector? \n\nWe would like to clarify that in our LPT model, the prior is not a trivial distribution but rather a learned distribution modeled by an implicit generative model, specifically a UNet. This choice allows for a more flexible and expressive prior compared to standard distributions. Our ablation studies in A.5 demonstrate the importance of this learnable prior. For example, in the Connect Four environment, LPT with the UNet prior achieves a score of $0.99 ± 0.01$, while LPT with a standard Normal prior (without UNet) achieves $0.90 ± 0.06$. Similar improvements are observed across Gym-Mujoco tasks. \n\nPer the reviewer's suggestion, in addition to comparing learnable and non-informative priors, we performed an additional set of experiments testing the effects of different UNet priors on the model performance. The results are as follows:\n\n\n| Model prior                          | Normalized score |\n| --------                             | -------- |\n|  UNet                                | $72.31 \\pm 1.92$ |\n|  UNet w/ smaller dimension           | $64.06 \\pm 1.94$ |\n|  UNet w/ fewer multipliers           | $64.59 \\pm 1.54$ |\n|  UNet w/ smaller initial conv        | $70.49 \\pm 2.84$ |\n|  UNet w/ single ResBlock             | $67.95 \\pm 4.64$ |\n|  No UNet                             | $56.88 \\pm 4.20$ |\n\nOur results underscores the crucial role of the learned prior in LPT's performance. Our ablation studies further reveal that our current UNet configuration is optimal among the variants tested. We appreciate the reviewer's suggestion, as it inspires us to perform a more detailed model analysis.\n\n\n\n> W2.2 and Q2: LPT may have much more parameters than DT due to the UNet used. Are the sizes of the models adjusted to account for this extra prior network in LPT?\n\nThank you for bringing up this important aspect in model training! Throughout our project, we have taken cautious to ensure a fair comparison between DT and LPT. We used the same implementation of the autoregressive decoder as DT, and since we don't have return-to-go conditioning in LPT, the total number of parameters end up being comparable:\n\n* DT with 3 layers and hidden size 128 has approximately 1.5 million parameters.\n* LPT with the UNet prior has around 1.6 million parameters.\n\nWhile there is a slight increase in the number of parameters for LPT, we believe this difference is not substantial enough to solely account for the performance improvements we observe. Importantly, both models use the same decoder implementation, ensuring a consistent base architecture for comparison. The additional parameters in LPT are primarily in the UNet-based prior, which allows for more expressive modeling of the trajectory distribution.\n\n\n> Q1: Are there derivations for Eq. (9)? Why is the second term in (9) are between the two parameters of the same model and one with stop_grad()?\n\n\nWe kindly refer you to our **Global Response [Explanation on Equation 9]**: on explanations of Eq. (9). To be brief, the understanding of this equation stems from the principle that maximizing likelihood is equivalent to minimizing the KL divergence between the model and data distribution. Then we can make it explicit that in this equation, we have $p^y_\\mathcal{D}(\\tau, z) = p^y_{\\bar{\\theta}}(z|\\tau) p^y_\\mathcal{D}(\\tau)$. Then we can derive,\n\n\\begin{align}\nD_\\mathrm{KL}(p^y_\\mathcal{D}(\\tau, z) || p^y_\\theta(\\tau, z)) &=\nD_\\mathrm{KL}(p^y_{\\bar{\\theta}}(z|\\tau) p^y_\\mathcal{D}(\\tau) || p^y_{\\theta}(z|\\tau) p^y_\\theta(\\tau)) \\\\\\\\&=\nD_\\mathrm{KL}(p^y_\\mathcal{D}(\\tau) || p^y_\\theta(\\tau))  + \\mathbb{E}\\_{p^y_\\mathcal{D}(\\tau)}D_\\mathrm{KL}(p^y_{\\bar{\\theta}}(z|\\tau) || p^y_\\theta(z|\\tau))\n\\end{align}\n\n\nThe gradient of this KL divergence with respect to model parameters $\\theta$ is consistent with Eq. 4, considering the stop_grad() operator.\n\nThank you again for your insightful review!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oZlm0kgip5", "reviewer": "Reviewer_NC2f", "summary": "The paper introduces the Latent Plan Transformer (LPT), a novel framework for trajectory generative modeling in the absence of step-wise reward. This framework employs a top-down latent variable model, using a temporally-extended latent variable z to represent a plan for decision-making. The framework comprises three components connected by the latent plan, a neural transformation of a Gaussian noise, a transformer-based trajectory generator, and a return estimator. LPT is optimized through maximum likelihood estimation on offline datasets composed of trajectory-return pairs. During testing, a latent plan is inferred based on an expected return, after which the trajectory generator is applied to extract actions. The framework is extensively evaluated across several environments, including gym locomotion (including antmaze), kitchen, maze2d, and connectfour. LPT addresses the challenge of temporal consistency while demonstrating competitive benchmark performance compared to baselines and excelling in various aspects, spanning credit assignments, trajectory stitching, and dealing with environment contingencies.", "questions": "- Line 77-78, duplicate \"Consider\" sentences.\n- What's the difference between $\\bar{\\theta}$ and $\\theta$ in eq. 9?\n- How the expected return is chosen for each dataset?\n- In Figure 3, are the testing $z_0$s sampled based on unseen returns? For LPT, we need to choose an expected return before the latent plan is sampled, in many cases we may want to sample a trajectory with a return even beyond the highest one in the training data. Therefore, it would be valuable for the authors to provide an analysis of how LPT performs across a range of returns to assess its robustness when interpolating and extrapolating out-of-distribution returns.", "limitations": "The authors adequately discussed the limitations of the proposed method.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- LPT creatively enables learning from trajectory-return pairs without any step-wise reward, while resolving the temporal consistency issue.\n- The paper provides a detailed analysis from a sequential decision-making perspective, identifying the significance of plan prediction in enforcing temporal consistency.\n- Exploitation-inclined Inference provides a simple and flexible way to control exploration and exploitation for latent plan sampling, which generally leads to better plans given the evaluation results.\n- The method is thoroughly evaluated, and LPT-EI exhibits superior performance compared to final-return baselines and strong stitching capabilities.\n- The paper is well-written and well-structured.", "weaknesses": "- Training and inference may suffer from inefficiency because of MCMC sampling in more complex scenarios, especially when modeling high-dimensional distribution (e.g. image-based benchmarks) or requiring high accuracy.\n- As mentioned by the authors, the LPT may fail on datasets with a skewed distribution of trajectory lengths (e.g. antmaze-large) due to its sequence modeling nature.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and constructive review of our paper! We shall address your questions point by point.\n> W1: Potential training and inference inefficiency because of MCMC sampling in more complex scenarios\n\nThank you for your insightful question! We would like to clarify that with careful design, MCMC does not become a bottleneck in our method.\n\n**Sampling in low-dim latent space**: Our MCMC sampling occurs in an extremely low-dimensional latent space compared to the entire trajectory. This latent space dimensionality can be controlled independently of input data dimensionality, mitigating the curse of dimensionality often associated with MCMC in high dimensions, even for high-dimensional inputs like images.\n\n**Training Efficiency**: During training, when sampling $p(z|\\tau,y)\\propto p(\\tau|z)p(y|z)p(z)$,  the main overhead is backpropagating through $p(\\tau|z)$ over long sequences multiple times. We mitigate this by using a fixed context window in the causal Transformer (e.g., K=32) to simplify attention computations. Additionally, we adopt a persistent Markov chain (line 153) that amortizes sampling across training iterations, reducing sampling iterations to 2. This approach has proven effective, allowing us to train on complex tasks without prohibitive computational costs.\n\n**Inference Efficiency**: During inference, we don't need to backpropagate through the trajectory generator $p(\\tau|z)$. As the latent plan decouples trajectory generation and return prediction (Sec 3.3), sampling only involves the much faster MLP for return prediction and the prior i.e. $p(z|y)\\propto p(y|z)p(z)$, which is computationally lightweight. This makes the inference process fast.\n\n> W2: As mentioned, LPT may fail on datasets with a skewed distribution of trajectory lengths due to its sequence modeling nature.\n\nThank you for raising this important point. The performance of sequence models like LPT and DT can indeed be affected by datasets with skewed distributions of trajectory lengths, as in the antmaze-large environment. This sensitivity to trajectory length distribution stems from the essence of sequence models, which learn to make decisions by discovering correlations between trajectories and their associated returns. When faced with a dataset that has an uneven distribution of trajectory lengths, these models may struggle to generalize effectively across time scales. Moving forward, we are keen to explore ways to address this challenge. We welcome any further thoughts or suggestions you might have on this matter.\n\n> Q1: Line 77-78, duplicate \"Consider\" sentences.\n\nThank you for catching that. We will correct the duplicate sentence in the revision.\n\n> Q2: Difference between $\\bar\\theta$ and $\\theta$ in eq. 9\n\nThanks for raising this question. We kindly refer you to our **Global Response [Explanation on Equation 9]**: on explanations of Eq. (9). To be brief, this equation stems from the principle that maximizing likelihood is equivalent to minimizing the KL divergence between the model and data distribution. Then we can make it explicit that we have $p^y_\\mathcal{D}(\\tau, z) = p^y_{\\bar{\\theta}}(z|\\tau) p^y_\\mathcal{D}(\\tau)$. Then we can derive,\n\\begin{align*}\nD_\\mathrm{KL}(p^y_\\mathcal{D}(\\tau, z) || p^y_\\theta(\\tau, z)) &= \nD_\\mathrm{KL}(p^y_{\\bar{\\theta}}(z|\\tau) p^y_\\mathcal{D}(\\tau) || p^y_{\\theta}(z|\\tau) p^y_\\theta(\\tau)) \\\\\\\\&=\nD_\\mathrm{KL}(p^y_\\mathcal{D}(\\tau) || p^y_\\theta(\\tau))  + \\mathbb{E}\\_{p^y_\\mathcal{D}(\\tau)}D_\\mathrm{KL}(p^y_{\\bar{\\theta}}(z|\\tau) || p^y_\\theta(z|\\tau))\n\\end{align*}\nThe gradient of this KL divergence with respect to model parameters $\\theta$ is consistent with Eq. 4, considering the stop_grad() operator.\n\n> Q3: Expected return selection\n\nWe chose the testing expected return as the expert-level performance metrics provided by the D4RL benchmark (Fu et al., 2020). The expected returns correspond to 100 on the normalized scale in D4RL for each environment. This setting aligns with experiments in Decision Transformer. \n\n> Q4-1: Plot of the testing $z_0$s in Figure 3\n\nYes you are correct. For the training and testing $z_0$, we refer to $p(z_0|\\tau,y)$ and $p(z_0|y)$ respectively. The testing $z_0$ values in Figure 3 are sampled based on unseen returns. For the Maze2D-medium dataset, the maximum trajectory normalized score in the training data is 12.8, significantly lower than our testing return (normalized score of 100).\n\n> Q4-2: LPT performance across return ranges\n\nWe thank the reviewer for this valuable suggestion. Our choice of normalized score 100 for $y$ is already out-of-distribution (OOD) for all datasets used. To illustrate this, we provide a table of D4RL dataset statistics with normalized scores:\n| Dataset| max. score|95 pct. score | 90 pct. score |\n|-|-|-|-|\n| maze2d-umaze-v1|21.1|13.2|10.3|\n| maze2d-medium-v1|12.8| 6.8|4.9|\n| maze2d-large-v1|16.9| 6.5| -2.5|\n| hopper-medium-v2 | 99.5 | 63.2|57.0|\n| hopper-medium-replay-v2| 98.6| 46.4|31.5 |\n| halfcheetah-medium-v2|45.0| 43.0|42.5|\n| halfcheetah-medium-replay-v2 |42.4|39.9|39.2|\n| walker2d-medium-v2| 92.0|83.4|82.4|\n| walker2d-medium-replay-v2|89.9|66.6|42.5|\n\nAs evident from this table, our chosen $y$ is an extrapolation beyond the dataset trajectories for all environments, particularly for maze2d datasets where the max score is notably suboptimal.\n\nTo further address the reviewer's concern about LPT's performance across different return ranges, we conducted an additional experiment on the Walker2D-medium-replay environment. We tested LPT with $y$ being the 50th and 90th percentiles of dataset returns, as well as the maximum dataset return and expert-level performance. As shown in **Fig. 1 of global response PDF**, LPT shows a linear trend in performance as we increase the target return. This analysis highlights LPT's robustness in both interpolating and extrapolating to OOD returns, a crucial capability for effective planning in various scenarios.\n\nThank you again for your valuable feedback!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ql8RDnWwUq", "reviewer": "Reviewer_W2Dk", "summary": "Authors propose a novel method Latent Plan Transformer (LPT). The key idea is to modify Decision Transformer (DT) approach by adding latent variable conditioning instead of return-to-go. This latent variable is assumed to represent a \"plan\" that the agent will follow. The motivation of replacing return-to-go is the fact that in practice return-to-go might be unavailable during the inference.", "questions": "1) What are results of DT/QDT if the same code is used for their training?\n\n2) What is the performance of the LPT using Antmaze medium/large tasks? Umaze tasks are not very representative. It is claimed that performance there is poor and appendix shows that there are some problems with those datasets. Anyway, what are the results? What if we  remove the trajectories of length 1 and train all baselines and LPT using modified datasets?\n\n3) How is the Figure 3 (left) obtained? It is said that $z_0$ is sampled from isotropic normal distribution but why is it in different t-sne space?\n\n4) Ablation: what if we remove the return-to-go conditioning during the training of the U-Net? What if we remove the trajectory conditioning during the training of the U-Net? I find these missing.\n\n5) What kind of distribution over $z$ do we obtain? I would recommend to try some toy example where we have $z$ with dim of 2. Is it  much different from gaussian?", "limitations": "The main limitation is the lack of intuition of what is actually happening and lack of latent variable analysis.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "To my prior knowledge, authors provide a novel idea of usage of latent variable for decision making with autoregressive generative model. \n\nAuthors provide strong empirical results on a diverse range of tasks which show that the proposed method has much higher performance than other types of the DT.", "weaknesses": "There is a chance I did not understand some parts of the work which lead to the lack of intuition behind the behavior that we observe. \n\nThe problem is that we sample some random variable $z$ that is conditioned on the final return and just plug this fixed into the Transformer's cross-attention (if I understood everything correctly). This provides better results into a better performance than DT even when it has access to the step-wise return to go. It seems very strange as intuitively updated return-to-go should let the model to adapt to the situation in which the agent ended up and change the behavior according to it while fixed latent variable might only provide the direction which the agent should follow and not updated. Can it be the implementation/hyperparameters differences? Baselines scores are taken from previous works and not claimed to be reproduced with LPT codebase and it is not provided. Or is it just some good latent representation of the return-to-go?\n\nSome experimental results are omitted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback! We greatly appreciate the opportunity to clarify and expand our work.\n> W1: Lack of intuition: plugging in a sampled fixed random variable conditioned on the final return performs better than DT with step-wise RTG\n\nThank you for this insightful observation. Indeed, we can view both RTG and the latent variable $z$ as functions of the whole trajectory and final return. RTG is hard-designed, while our latent $z$ with a learnable prior distribution is learned from the posterior distribution $p(z|\\tau, y)$. The learned $z$ potentially captures more information than RTGs, serving as a rich, abstract representation of the entire trajectory and expected return, providing guidance throughout the process. In Sec.4, we show that given final return, both latent $z$ and RTG can help promote temporal consistencies. During inference, although RTG updates at each time-step, the non-Markovian state representation at each time-step also consults $z$ to make decisions via the cross-attention mechanism, allowing the model to maintain and utilize information about the trajectory's progress and the remaining challenges.\n> Q1: Results of DT/QDT if the same code is used.\n\nRegarding the implementation of DT, we used the publicly available code provided by Huggingface. For QDT, since the authors did not release a public version of the code, we relied on the results reported in their original paper. We will release our codebase to facilitate reproducibility and further investigation.\n> Q2.1: LPT performance in Antmaze medium/large tasks.\n\nThanks for pointing out this important aspect of our work, which indeed presents some challenges as we noted in our paper's limitations section. For Antmaze, we attach here the results as the reviewer suggested:\n||CQL|DT|LPT-EI|\n|-|-|-|-|\n|Antmaze-diverse-large|14.9 ± 3.2|0.0 ± 0.0|3.6 ± 0.8|\n\nThese results show that LPT's performance is lower than CQL, though it outperforms DT. The difference in performance highlights an important distinction between trajectory-based models and TD-learning methods in offline reinforcement learning, as we mentioned in the **global response**. Sequence models like LPT and DT are designed to learn from correlations between sequences of actions and states and their associated returns. However, this strength can become a limitation when faced with datasets that have skewed distributions of trajectory lengths, as in the Antmaze-large environment. If the dataset is dominated by very short or very long trajectories, it can bias the model's understanding of this relationship. \n> Q2.2: Removing the trajectories of length 1 and train all baselines and LPT using modified datasets.\n\nThis is a very interesting and insightful point! We conduct another set of experiments where we remove trajectories of length 1 from the Antmaze-large-diverse dataset to perform model training. The results are as follows:\n||DT|LPT-EI|\n|-|-|-|\n|Antmaze-diverse-large|0.0 ± 0.0|3.2±0.9|\n\nResults indicate that the performance of LPT remains suboptimal -- unfortunately upon closer examination of the new dataset, we found that the distribution the trajectory lengths are still biased, as we show in **Fig. 2 in our global response PDF**. We can see that most of the trajectories with final return=1 are very short in length (median = 13). The trajectory coverage can be detrimental to sequence models like LPT, which learn to make decisions by discovering correlations between trajectories and returns. We appreciate the reviewers' constructive feedback, which has prompted us to conduct a thorough analysis. This analysis highlighted the importance of dataset quality for sequence models.\n> Q3: Figure 3: distribution of $z_0$ and t-sne plot\n\nGood catch! In Figure 3 (left), we illustrate the learned distribution in the $z_0$ space. For the blue points, we plot $\\mathbb{E}_{p(\\tau, y)}p(z_0|\\tau,y)$ where each point represents one training data point in $z_0$ space. For yellow points, we plot, in the testing or inference case, the sampled $p(z_0|y^*)$, given expected final return $y^*$. Since $y^*$ is chosen outside the training distribution, these $z_0$s may be distant from training dataset.\n> Q4: Ablation: removing the RTG conditioning during the training of the UNet\n\nThanks for raising this point! We would like to clarify that in LPT, the UNet is not conditioned on RTG or trajectories during training. Instead, it serves as a flexible prior for the latent variable $z$. It transforms an initial Gaussian noise into a more expressive distribution. The UNet is trained jointly with the rest of the model through our overall objective function. It learns to shape the prior distribution of $z$ to better match the posterior distribution inferred from trajectories and returns, but it does not directly take these as inputs. As a result, the suggested ablations of removing RTG or trajectory conditioning from the UNet are not applicable to our model.\n> Q5: Distribution over $z$: recommend to try 2D $z$\n\nGreat idea! Following your suggestion, we add a new experiment using a 2D latent space ($z$ with dim=2) and apply it to the maze2d-umaze environment. This low-dimensional setting allows for easy visualization and analysis of the latent space distribution. \n\nAs designed, $z_0$ is sampled from a 2D Gaussian distribution. After UNet, the distribution of $z$ is notably different from Gaussian, as seen in **Fig. 3 of our global response PDF**. This experiment shows that our model learns a non-Gaussian latent representation that captures task-relevant information. We would like to highlight that our method doesn't impose any specific distributional assumptions on $z$. The UNet is free to transform the initial Gaussian distribution into the form that is most informative for the task at hand, allowing the model to encode richer, more complex patterns about the trajectories and strategies.\n\nThanks again for your insightful feedback and for taking the time to review our work.", "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:
## Global Response

We would like to thank all reviewers for their helpful feedback! Here, we would like to add some clarifications based on points raised by multiple reviewers.

**[Latent variable $z$ formulation]**

We first higlight that the latent variable of our model is $z$, whose prior distribution is a learnable transformation from Gaussain white noise $z_0\sim N(0,I)$. It is parameterized by an expressive UNet. Meanwhile, this prior model can be viewed as an implicit generator with implicit densities.

**[Additional explanation on Equation 9]**

To understand Eq. 9, we can start from the well-known equivalence between Maximum Likelihood Estimate $\max\_\theta\mathbb{E}\_{p^y_\mathcal{D}(\tau, z)}[\log p^y_\theta(\tau, z)]$ (Eq. 4) and KL minimization $\min\_\theta\mathbb{E}\_{p^y_\mathcal{D}(\tau, z)}[\log p^y_\mathcal{D}(\tau, z) - \log p^y_\theta(\tau, z)]$ (Eq. 9). Here $p^y_\mathcal{D}(\tau, z) = p^y_{\bar{\theta}}(z|\tau) p^y_\mathcal{D}(\tau)$. We have a stop_grad version of $\bar\theta$ because deriving Eq. 4 from the marginal MLE $\max_\theta\mathbb{E}\_{p^y_\mathcal{D}(\tau)}[\log p^y_\theta(\tau)]$ does not involve a derivative through $\mathbb{E}\_{p^y_{{\theta}}(z|\tau)}$ (as shown in Appx. A.1). 

**[Comparison between decision-making algorithms]**

We hope it is helpful to further clarify on the implicit data specification in different decision-making algorithms. Sequence-based methods (e.g. DT and LPT) and classical TD-learning methods (e.g. CQL) have different assumptions in data distribution. The non-Markovian autoregressive nature of DT and LPT requires the training data to have reasonable coverage of sequence length. In contrast, CQL operate primarily on individual state-action-reward-next_state tuples, which can be advantageous in datasets with skewed trajectory length distributions.

Meta-review: This paper presents a new method to improve the representations learned in decision transformers. After the rebuttal and discussion time, the reviewers agree that the paper is promising and adds to our understanding of how to improve decision transformers. There are still some outstanding questions about how this representation learning method relates to other representation learning methods in RL, and we encourage the authors to continue investigating this area.
Decision: Accept (poster)