Latent Plan Transformer for Trajectory Abstraction: Planning as Latent Space Inference

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.

Strengths

To my prior knowledge, authors provide a novel idea of usage of latent variable for decision making with autoregressive generative model. Authors 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. 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

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?

Rating

6

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

The main limitation is the lack of intuition of what is actually happening and lack of latent variable analysis.

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.

Strengths

- LPT creatively enables learning from trajectory-return pairs without any step-wise reward, while resolving the temporal consistency issue. - The paper provides a detailed analysis from a sequential decision-making perspective, identifying the significance of plan prediction in enforcing temporal consistency. - 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. - The method is thoroughly evaluated, and LPT-EI exhibits superior performance compared to final-return baselines and strong stitching capabilities. - 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. - 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

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

Rating

6

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

The authors adequately discussed the limitations of the proposed method.

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.

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

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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors addressed the limitations.

Rebuttal Acknowledgement

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.