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_cc_2024_Conference
Published: 2024-05-12 23:20:14

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.

Main paper: 
 introduction: 
 Decision Transformer (DT) (Chen et al., 2021) and some concurrent work (Janner et al., 2021) have popularized the research agenda of decision-making via generative modeling. The general idea is to consider decision-making as a generative process that takes in a representation of the task objective, e.g. the rewards or returns of a trajectory, and outputs a representation of the trajectory. Intuitively, a purposeful decision-making process should shift the trajectory distribution towards regimes with higher returns. In the classical decision-making 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-gos (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 a distribution of actions conditioned on step-wise RTGs.In this work, we are interested in the problem of planning. Among various ways to identify planning as a special class of decision-making problems, we pay particular attention to its data specification and inductive biases. As designing step-wise rewards requires significant effort and domain expertise, we focus on the problem of learning from trajectory-return pairs, where a 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 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 a 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 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 a 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.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 4 and Section 6Guidelines:• 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 be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
background: 
 A sequential decision-making problem can be formulated with a decision process ⟨S, A, H, T r, r, ρ⟩ that contains a 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. T r : S + × A + → Π(S) is the transition that returns a distribution over the next state. r : S + × A + → R specifies the real-valued reward at each step. ρ : Π(S) is the initial state distribution that is always uncontrollable to the agent. The agent's decisions follow a policy π : S + × A + → Π(A). In each episode, the agent interacts with the transition model to generate a trajectory τ = (s 1 , a 1 , s 2 , a 2 ..., s H , a H ).The objective of sequential decision-making is typically formulated as the expected trajectory re-turn y = H t=0 r t , Q = E p(τ ) [y].Conventional RL algorithms solve for a policy π(a t |s t , * ), where the conditioning * denotes the optimal expected return. DT generalizes this policy to π(a t |s ≤t , a <t , RT G ≤t ), by fitting the joint distribution p(s 1 , a 1 , RT G 1 , ...s T , a T , RT G T ) with a Transformer. RT G t is the return-to-go from step t to the horizon H, RT G t = H k=t r(s ≤k , a ≤k ). It is a useful indication of future rewards, especially when rewards are dense and informative.However, RT G 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 to 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 assignment mechanism can be formalized as, RT GQ t = K k=t r(s ≤k , a ≤k )+E[Q(s ≤K , a ≤K )].Here Q can be estimated using deep TD learning with multi-step returns. Yamagata et al. (2023) instantiate a Markovian version and demonstrate improvement in trajectory stiching.Whatever credit assignment we use, be it RT G or RT G Q , the purpose is to explicitly model the statistical association between trajectory steps and final returns. This effort is believed to be necessary because of the exponential complexity of the trajectory space. This belief, however, can be reexamined 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. Given a variable-length trajectory τ , z ∈ R d is a vector that represents τ in the latent space. y ∈ R is the return of the trajectory. The joint distribution of the trajectory and its return is defined as p(τ, y).The latent trajectory variable z, conceptualized as a plan, is posited to decouple the autoregressive policy and return estimation. From a statistical standpoint, with z given, we assume that τ and y are conditionally independent, positioning z as the information bottleneck. Under this assumption, the Latent Plan Transformer (LPT) can be defined as,EQUATIONwhere θ = (α, β, γ). LPT approximates the data distribution p data (τ, y) using the marginal distribution p θ (τ, y) = p θ (τ, y, z)dz. It also establishes a generation process,EQUATIONThe prior model p α (z) is an implicit generator, defined as a learnable neural transformation of an isotropic Gaussian, z = U α (z 0 ) and z 0 ∼ N (0, I d ). U α (•) is an expressive neural network, such as the UNet (Ronneberger et al., 2015) . This approach is 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. (Chen et al., 2021) . Specifically, the latent variable z is included in trajectory generation using crossattention, as shown in Fig. 1 and controls each step of the autoregressive trajectory generation as p β (a t |s t-K:t , a t-K:t-1 , z). The action is assumed to follow a single-mode Gaussian distribution, i.e. a t ∼ N (g β (s t-K:t , a t-K:t-1 , z), I |A| ).The return predictor is a non-linear regression on the latent trajectory variable z, modeled as p γ (y|z) = N (r γ (z), σ 2 ). It directly predicts the final return from the latent variable z. The function r γ (z) is a small multi-layer perceptron (MLP) that estimates y based on z. The variance σ 2 , is treated as the hyper-parameter in our setting.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: See Section 7Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
offline learning: 
 With a set of offline training examples {(τ i , y i )} n i=1 , we aim to learn Latent Plan Transformer (LPT) through maximum likelihood estimation (MLE). The log-likelihood function is defined as L(θ) = n i=1 log p θ (τ i , y i ). The joint probability of the trajectory and final return isEQUATIONwhere p 0 (z 0 ) = N (0, I d ). The learning gradient of log-likelihood can be calculated according toEQUATIONThe full derivation of the learning method is in Appendix A.1. Let δ α , δ β , δ γ represent the expected gradients of L(θ) with respect to the model parameters α, β, γ, respectively. The learning gradients for each component are formulated as follows.For the prior model p α (z),δ α (τ, y) = E p θ (z0|τ,y) [∇ α (log p β (τ |z = U α (z 0 )) + ∇ α log p γ (y|z = U α (z 0 ))].For the trajectory generator,δ β (τ, y) = E p θ (z0|τ,y) [∇ β log p β (τ |z = U α (z 0 ))],For the return predictor,δ γ (τ, y) = E p θ (z0|τ,y) [∇ γ log p γ (y|z = U α (z 0 ))].Estimating these expectations requires Markov Chain Monte Carlo (MCMC) sampling of the posterior distribution p θ (z 0 |τ, y). We use the Langevin dynamics (Neal, 2011) for MCMC sampling, iterating as follows for a target distribution π(z):EQUATIONwhere k indexes the time step of the Langevin dynamics, s is the step size, and ϵ k ∼ N (0, I d ) is the Gaussian white noise. Here, π(z) is instantiated as the posterior distribution p θ (z 0 |τ, y). We have p θ (z 0 |τ, y) ∝ p 0 (z 0 )p γ (y|z)p β (τ |z), where z = U α (z 0 ), such that the gradient is∇ z0 log p θ (z 0 |τ, y) = ∇ z0 log p 0 (z 0 ) prior +∇ z0 log p γ (y|z) return prediction + H t=1 ∇ z0 log p β (τ (t) |τ (t-K:t-1) , z)aggregating finite-context sub-trajectories .This demonstrates that the posterior inference of z is an explicit process of optimizing a plan given its likelihood. In the presence of a finite context, p β (τ |z) parametrized with Transformer can only account for sub-trajectories with a maximum length of K. The latent variable z serves as an abstraction that integrates information from both the final return and sub-trajectories using gradients.The sampling process starts by initializing z k=0 0 from a standard normal distribution N (0, I d ). 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 a theoretical understanding of this noiseinitialized finite-step MCMC, see Pang et al. (2020) ; Nijkamp et al. (2020) ; Xie et al. (2023) . However, for large horizons (e.g.,H=1000), this method becomes slow and memory-intensive. To mitigate this, we adopt the persistent Markov Chain (PMC) (Tieleman, 2008; Xie et al., 2016; Han et al., 2017) , which amortizes sampling across training iterations. During training, z k=0 0 is initialized from the previous iteration and the number of updates is reduced to N = 2 steps. See Appendix A.2 for training and architecture details.The MLE learning of LPT gives us an agent that can plan. During testing, we first infer the latent z 0 given the desired return y using Bayes' rule,EQUATIONThis posterior sampling is achieved using Langevin dynamics similar to the training process. Specifically, we replace the target distribution in Eq. ( 5) with p θ (z 0 |y) and run MCMC for a fixed number of steps. Sampling from p θ (z 0 |y) eliminates the need for expensive back-propagation through the trajectory generator p β (τ |z).This posterior sampling of p(z 0 |y) is an explicit process that iteratively refines the latent plan z, increasing its likelihood given the desired final return. It aligns with our intuition that planning is 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 to generate actions conditioned on the current state and recent history,a t ∼ p β (a t |s t-K:t-1 , a t-K:t-1 , z = U α (z 0 )).Once a decision is made, the environment's (possibly non-Markovian) transition s t+1 ∼ p(s t+1 |a t , s t ) emits the next state. This sequential decision-making process iterates the sampling of s t and a t until termination at the horizon.Exploitation-inclined 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)EQUATIONwhich has the score ∇ z0 log pθ (z 0 |y) = ∇ z0 log p 0 (z 0 ) + w∇ z0 log p γ (y|z). This guidance weight w controls the interpolation between exploration and exploitation. When w = 1, the sampled plans collectively represent the posterior density and account for Bayesian uncertainty, resulting in a provably efficient exploration scheme (Osband and Van Roy, 2017) . When w > 1, the sampled plans are 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.Input: Learning iterations T , initial parameters θ 0 = (α 0 , β 0 , γ 0 ), offline training samples D = {τ i , y i } n i=1 , posterior sampling step size s, the number of steps N , and the learning rate η 0 , η 1 , η 2 . Output: θ T for t = 1 to T do 1.Posterior sampling: For each (τ i , y i ), sample z 0 ∼ p θt (z 0 |τ i , y i ) using Eq. ( 5) with N steps and step-size s, where the target distribution π is p θt (z 0 |τ i , y i ).2.Learn prior model p α (z), trajectory generator p β (τ |z) and return predictor p γ (y|z):α t+1 = α t +η 0 1 n i δ α (τ i , y i ); β t+1 = β t +η 1 1 n i δ β (τ i , y i ); γ t+1 = γ t +η 2 1 n i δ γ (τ i , y i ) as in Section 3.2. end forAlgorithm 2 Planning as inference Input: Expected return y, a trained model on offline dataset θ, posterior sampling step size s and the number of steps N , Horizon H and an evaluation environment. Output: Trajectory τ if Exploitation-inclined Inference (EI) then Sample z 0 ∼ pθ (z 0 |y) as in Eq. ( 7) using Eq. ( 5) with N steps and step size s, where the target distribution π is replaced by pθ (z 0 |y) ∝ p 0 (z 0 )p γ (y|z = U α (z 0 )) w and z = U α (z 0 ). else Sample z 0 ∼ p θ (z 0 |y) as in Eq. ( 6) using Eq. ( 5) with N steps and step size s, where π is replaced by p θ (z 0 |y) ∝ p 0 (z 0 )p γ (y|z = U α (z 0 )) and z = U α (z 0 ). end if while current time step t ≤ H do Sample a t using trajectory generator asa t ∼ p β (a t |s t-K:t-1 , a t-K:t-1 , z = U α (z 0 )).Once a decision is made, the environment's transition s t+1 ∼ p(s t+1 |a t , s t ) emits the next state. end whileQuestion: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to Answer: [NA] Justification: None. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
a sequential decision-making perspective: 
 We approach the sequential decision-making problem with techniques from generative modeling. In particular, our data specification of trajectory-return 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 RT G, which naturally hints the future progress of the trajectory. How is temporal consistency enforced in our model without the assistance of the RT Gs?Without loss of generality, consider the trajectory distribution conditioned on a single return value y. The MLE objective is equivalent to minimizing the KL divergence between the data distribution and model distribution, D KL (p y D (τ )∥p y θ (τ ) 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 a stochastic environment, like generating the next text token.This temporal consistency issue could be alleviated by additionally modeling the sequence of RT G.Denote ρ = (RT G 0 , RT G 1 , ...RT G H ).Modeling the joint distribution is to minimizeEQUATIONNote that the RTG prediction term is conditioned on the entire trajectory, including the future steps. Minimizing this additional KL divergence correlates predicted RT Gs with hindsight trajectory-to-go.Our modeling of the latent trajectory variable z provides an alternative solution to the temporal consistency issue. Eq. ( 4) is minimizing the KL divergenceEQUATIONwhere p y θ (z|τ ) = p y D (τ, z)/p y D (τ ) and θ = θ highlights these distributions have the same parameterization as p y θ but are wrapped with stop_grad() operator when calculating gradients for θ (Han et al., 2017) . Comparing Eqs. ( 8) and ( 9), it is now clear that z plays a similar role as RT G in promoting temporal consistency in autoregressive models. Uniquely, p y θ (z|τ ) is the temporal abstraction intrinsic to the model, in contrast to step-wise rewards. From a sequential decision-making perspective, z is effectively a plan that the agent is persistent to. From a generative modeling perspective, z from different trajectory modes would decompose the density p y (a t |s 0:t , a 0:t-1 ), relieving the burden of learning the autoregressive policy p β (a t |s 0:t , a 0:t-1 , z).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) 2022). Among them, Yang et al. (2022) pinpoints the issue by viewing RT Gs as deterministic latent trajectory variables, closely related to what we present here. Uniquely, the latent variable z in our model is inherently multi-modal (hence very non-deterministic) and ignorant of step-wise rewards. We postulate that the overfitting issue might be mitigated. This is validated by our empirical study inspired by Paster et al. (2022) .Although RTG prediction and plan prediction both promote temporal consistency, they function very differently when mixing trajectories from multiple return-conditioned regimes. RTG prediction is a supervised learning over the joint distribution p D (τ, ρ). Simply mixing trajectories from multiple regimes can't encourage generalization to trajectories that are stitched with those in the dataset. Yamagata et al. (2023) propose to resolve this by replacing RT G with RT G Q . Intuitively, this augments the distribution p D (τ, ρ) with p D (τ ′ , ρ Q ), where τ ′ denotes trajectories covered by the offline dynamic programming, such as Q learning, andρ Q = (RT G Q 0 , RT G Q 1 , ...Q H ).It significantly improves tasks requiring trajectory stitching. Conversely, plan prediction is an unsupervised learning as it samples from p D (τ, y)pθ(z|τ, y). As z contains more trajectory-related information than step-wise RT Gs, trajectories lying outside of p D (τ, ρ) may be in-distribution for p D (τ, y)pθ(z|τ, y). The return prediction training further shapes the representation of z, 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. 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 return-conditioned 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 TDlearning'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 to further extend the effective horizon. 
experiments: 
 The data specification of trajectory-return pairs distinguishes our empirical study from most existing works in offline RL. Omitting step-wise rewards naturally increases the challenges in decision-making.Our empirical study adopts the convention from offline RL. We first train our model with the offline data and then test it as an agent in the corresponding task. More training details and ablation studies of LPT can be found in Appendices A.2 and A.4.OpenAI Gym-Mujoco The D4RL offline RL dataset (Fu et al., 2020) features densely-rewarded locomotion tasks including Halfcheetah, Hopper, and Walker2D. We test for medium and mediumreplay. It also includes Antmaze, a locomotion and goal-reaching task with extremely sparse reward.The agent will only receive a reward of 1 if hitting the target location and 0 otherwise. We use its umaze and umaze-diverse variants.Franka Kitchen Franka Kitchen is 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-directed demonstrations and partially task-directed demonstrations respectively.Maze2D Maze2D is a navigation task in which the agent reaches a fixed goal location from random starting positions. The agent is rewarded 1 point when it is around the goal. Experiments are 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 is 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, 0 for a 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) and QDT (Yamagata et al., 2023) . CQL baseline results are obtained from Kumar et al. (2020) . QDT baseline results are from Yamagata et al. (2023) . The DT results for Gym-Mujoco and Maze2D tasks are from Yamagata et al. (2023) When resolving the temporal consistency issue, our model doesn't have an explicit credit assignment mechanism that accounts for the actual contribution of each step. It is not aware of the step-wise 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 to high-dimensional actions The Gym-Mujoco environment was a 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 a planning problem, for it involves decision-making 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 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. In addition to credit assignment, the setup of offline RL further presents a challenge, trajectory stitching (Fu et al., 2020) , which articulates the problem of shifting the trajectory distribution towards sparsely covered regimes with higher returns. In 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 learn to piece together relevant sub-trajectories. A similar setting happens in Maze2D domain. Taking Maze2D-medium as an example, in the training set, the average return of all trajectories is 3.98 with a 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. To probe into the agent's understanding of trajectories' returns, we visualize the representation space of the latent variables. The left of Fig. 3 is the aggregated posterior distribution of z 0 . We can see that z 0 infered from p θ (z 0 |y) 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 z, which is transformed from z 0 with the UNet, z = U α (z 0 ). We observe that zs 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.To live in a stochastic world, contingent planning that is adaptable to unforeseen noises is desirable. Paster et al. (2022) ; Yang et al. (2022) discover that DT's performance would degrade in stochastic environments due to inevitable overfitting towards contingencies. We examine LPT and other baselines in Connect Four from Paster et al. (2022) . Connect Four is a two-player game, where the opponent will make adversarial moves to deliberately disturb an agent's plan. According to the empirical study from Paster et al. (2022) , the degradation of DT is more significant than in stochastic Gym tasks from Yang et al. (2022) . As shown in  
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) . 
summary: 
 We study generative modeling for planning in the absence of step-wise rewards. We propose LPT which generates trajectory and return from a 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 a 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. In our experiments, we encounter a curious phenomenon that LPT outperforms CQL, DT and QDT in Antmaze-umaze by a large margin but falls behind in Antmaze-large. Upon closer examination of 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.86%) with length=1, 893 trajectories (12.43%) with length=1000, 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.As a reference, Fig. 5 shows the distributions of final returns and the trajectory lengths of Maze2Dlarge, a task where LPT performs well. It is important to note that TD-learning methods, such as CQL and QDT, rely solely on (s, a, s ′ , r) tuples and are less affected by the trajectory length distribution in the dataset. Consequently, Antmaze-large in D4RL remains a fair dataset for these methods to perform offline RL.We investigate the role of the expressive prior p α (z) in our Latent Plan Transformer (LPT) model by removing the UNet component, which transforms z 0 from a non-informative Gaussian distribution.Table 9 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. 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-medium-replay task with various UNet architectures. We observe that reducing the capacity or 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. 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 is optimal among the variants tested. We appreciate the reviewer's suggestion, as it prompted us to perform a more detailed analysis of the prior's impact on LPT.We are also interested in LPT's potential in finetuning or even continual learning. Inspired by ODT (Zheng et al., 2022) , we employ a trajectory replay butter (Mnih et al., 2015) to store samples from online interaction in a first-in-first-out 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; Osband and Van Roy, 2017) . We report the results in Table 11 . Despite the significance in a few tasks, the improvement is within 1 standard deviation of the mean for the majority. We observe a similar pattern in ODT (Zheng et al., 2022) . We introduce a novel generative model that leverages the idea of "planning as inference", which potentially leads to more capable and efficient AI systems that make better decisions. One potential negative impact could come from the misuse of the generative modeling in designing objects or entities for harmful purposes. Developing suitable safeguards and regulations will be crucial to mitigate potential negative impacts.NeurIPS Paper Checklist 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.

- 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?
questions: 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( )?

- Is the comparison with baselines fair in the sense that all the models have similar parameter sizes?
limitations: limitations: Yes, the authors addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
questions: questions: - Line 77-78, duplicate "Consider" sentences.
- What's the difference between $\bar{\theta}$ and $\theta$ in eq. 9?
- How the expected return is chosen for each dataset?
- 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: limitations: The authors adequately discussed the limitations of the proposed method.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

The 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?

Some experimental results are omitted.
questions: questions: 1) What are results of DT/QDT if the same code is used for their training?

2) 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?

3) 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?

4) 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.

5) 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: limitations: The main limitation is the lack of intuition of what is actually happening and lack of latent variable analysis.

Comments: title: Thank you for your insightful feedback!
Thank you for your time!
Thank you!
Reviewer Response Requested
Review Discussion Requested
Reviewer Discussion Needed
Rebuttal Acknowledgement
comment: Dear Reviewer,

Thank you for your deeply insightful feedback. We appreciate your insight that the benefits of our latent variable modeling could stem from its advantages in state compression/aggregation. We will follow your advice to conduct theoretical analysis to strengthen our claims. We'll also highlight our ablation study and comparison results in the main text, as suggested.

We're grateful for your engagement and continued support.

Best,

Authors.
Thank you for the reply. Indeed most existing MDP representation learning focuses on learning some step-by-step models, while the proposed method learns a latent variable for the entire trajectory. But from existing theories, learning such a whole trajectory variable might not be useful from the decision-making point of view, as the state itself should have included all necessary information for making optimal decisions. The benefits of the latent variable could be due to its advantages in helping state compression/aggregation. This may be beyond the scope of this paper, but having some theoretical analysis on a simple problem showing such benefits might further strengthen the claims of the paper.

As of the ablation study and fair comparison to the baseline, thank you for the additional stats and results. I think they are important and could be highlighted a bit more in the main text.
Thank you for your positive feedback and thorough review. We’re pleased that your concerns have been addressed. Have a great day!
Thank you for your positive feedback and thorough review. We’re glad all your concerns have been addressed. Have a nice day!
Dear Reviewer,

The discussion time is coming to an end soon. Please engage in the discussion process which is important to ensure a smooth and fruitful review process. Give notes on what concerns have not been addressed and why.
Dear Reviewer,

The discussion time is coming to an end soon. Please engage in the discussion process which is important to ensure a smooth and fruitful review process. Give notes on what concerns have not been addressed.
Dear Reviewer,

The discussion time is coming to an end soon. Please engage in the discussion process which is important to ensure a smooth and fruitful review process. Give notes on what parts of the reviewers responses that have and have not addressed your concerns.
I thank authors for their detailed response. Most of my concerns were resolved and I'm adjusting the score.
I would like to thank the authors for their detailed responses and additional experiments. It's very helpful to know that the expected return used for evaluation is normalized score 100. Most of my concerns and questions have been addressed, I would like to maintain my positive rating.