QueST: Self-Supervised Skill Abstractions for Learning Continuous Control

Summary

The paper introduces an approach utilizing latent variable generative models in conjunction with FSQ quantization techniques to learn good representations of action sequences. It also proposes a prior network for autoregressive modeling at the level of behavioral representations. The superiority of the learned representations is validated across three distinct robotic environments. Additionally, the paper presents intriguing observations, such as the impact of causality on representation learning, and supports these findings with ablation studies. This paper is not too different from many previous relevant literature, and the innovation is a little weak.

Strengths

- The overall paper is easy to read, with the description of the methodology being well-written. - The experimental content of the paper is robust, and the demonstrations clearly indicate a significant performance enhancement of the proposed method over the baseline.

Weaknesses

- The overall pipeline that first learns the abstract representations of behaviors and then models the behavior with the autoregressive model is common. I believe the main difference of this paper contains two points: (1) using FSQ to learn discrete latent space; (2) implement the encoder with a causal transformer instead of a non-causal transformer. So the innovations seem not enough. - A lot of ablation studies are conducted. However, the related analysis is not sufficient.

Questions

- Why the causality of the encoder transformer is so critical?

Rating

4

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes

Summary

- This work introduces a new method called QueST that uses a latent variable model to learn a set of discrete temporal action abstractions / motion primitives / skills for imitation learning. This is done by training an auto-encoder to encode a sequence of actions using a causality preserving encoder, discretising the embedding using Finite Scalar Quantization (FSQ) to obtain a sequence of discrete latents, and then decoding the latents using a causality preserving decoder to recover the sequence of actions. - A policy is learned by predicting the latent actions from a history of observations and a language instruction while keeping the latent decoder frozen. An exception to the latter is when fine-tuning a policy in a few-shot setting. At test time, top-k sampling for sampling an optimal sequence of latents. - The method is validated on the LIBERO and Meta-World environments and compared to several recent approaches: ResNet-T, diffusion policy, ACT, VQ-Bet, and PRISE.

Strengths

Originality: - This work is a nice combination of existing methods (latent action abstraction, causal encoders/decoders, and Finite Scalar Quantisation). - The use of causal encoders/decoders for latent abstraction is an interesting idea. Quality: - LIBERO and Meta-World are reasonable benchmarks for validating the proposed method. - The method shows good performance when compared against several recent baselines: diffusion policy, ACT, VQ-BeT, and PRISE. - The experimental section contains ablations for several key design choices: FSQ vs VQ, conditioning/not conditioning the auto-encoder on observations, mirroring the architecture of the encoder for the decoder, causal/non-causal encoders, and a sensitivity analysis for the auto-encoder downsampling factor and codebook size. - Results are reported for three random seeds for the experiments on LIBERO. - Line 541-544: The paper preempts concerns about the Meta-World results due to issues with the scripted expert policies and updated results are provided at the paper website. Clarity: - The paper is generally well-written. - The videos on the project website provide a nice indication of qualitative behaviour. Significance: - Considering the good performance and the simplicity of the method, it seems reasonable that the proposed method will be adopted or further developed by the community, in particular if the code is released as stated in the NeurIPS paper checklist.

Weaknesses

Originality: - Two of the baselines, VQ-BeT and PRISE, also learn temporal action abstractions. It would be nice to have a more detailed and clearer discussion of the differences between these three methods than is currently provided, e.g., differences and similarities could be summarised in a table. Quality: - There are no results for PRISE on Multitask LIBERO-LONG. - The paper website states: “Depending on the dataset, we also tune some key hyperparameters for the baselines”. This is vague and it is unclear why this was not done for every dataset. - It is not clear whether all the baselines are based on new experiments or whether some of the results are shown as reported in previous works. - There is no analysis of inference latency of the proposed method and the baselines, which is important for practical applications. - It might be interesting to have an ablation experiment that uses Residual VQ as used by VQ-BeT instead of FSQ to assess if this choice plays a significant role for the difference in performance. Clarity: - Some information about how the hyperparameters of the baselines were tuned are on the paper website, but important information like this should be in the main body of the paper or the appendix. - The paper website states that “To ensure fair comparison of different model architectures, we use same input modalities and same observation & task encoders for all baselines…”. This is important information that should be in the main body or the appendix of the paper. It is also unclear whether this only applies to the Meta-World experiments or also the LIBERO experiments. - It is not specified where the “provided scripted policies” for collecting demonstrations on Meta-World as mentioned on the paper website can be found. - Minor comment: Figure 3b is not described or mentioned in the text. Significance: - The proposed method feels a bit like an incremental evolution of VQ-BeT and PRISE. Even though these have only come out very recently, it would still be good to have a clearer positioning with respect to these two works as already mentioned above.

Questions

- Are the updated results on Meta-World going to replace the current results in the main body of the paper? - Where can the “provided scripted policies” for collecting demonstrations on Meta-World as mentioned on the paper website be found? - Are all results for the baselines based on new experiments or are any of the results shown as reported in previous works? In the latter case, it would be useful for this to be indicated in the result tables. - How were the hyperparameters tuned for the proposed method, for the ablations, and for baselines? - How does the proposed method compare to the baselines in terms of latency? - Is the code going to be released publicly? Updates: - Presentation: 2 -> 3 - Score: 4 -> 6

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

- It would be nice to have a discussion of latency considerations.

Summary

The paper aims to capture skill abstractions through training a latent space through encoding and decoding actions. The resulting latent space is used for training a policy that converts observations into the latent space, and uses the trained decoder to output actions. The paper conducts experiments on manipulation suites (LIBERO and MetaWorld) and demonstrates improvement over existing approaches on few-shot and multitask setting.

Strengths

- The idea makes use of the structure of robotic manipulation tasks, where the sequence of actions are often similar for different items and tasks. Notably, the action decoder is not conditioned on the states to learn a latent space that is invariant to the low-level state information. - The paper is mostly well written

Weaknesses

I am happy to increase my score if these points are addressed. **Comments** - Should probably include analysis on what the latent $z$'s ended up learning. Do they actually have some temporal abstraction going on? Can we associate them? It appears that this is in the website but I believe this should be in the main paper as one premise is that the latent representation is leveraging the structure of the actions. - This work appears to rely on the fact that the demonstrations have some form of structures---probably a limitation. Perhaps in this experimentation setting it is a fair assumption, but what if we have partially observed setting (e.g. items with dynamic inertial properties). - It is not clear to me how this approach handles multimodality (in the sense of action distributions?) - Ablation on "causality": I feel the word "causality" is misused because this is simply a mask enforcing whether or not to look at the future data. Is this not still modelling correlation only? How is this causal exactly?

Questions

**Questions** - Equation 1, perhaps write out what round_std is. - Page 5, lines 179-180: What is the intuition of using cross attention between positional embedding and skill tokens but not self-attention with positional encoding? - Page 5, lines 186-187: Do the authors think that this is because the actions can be invariant (e.g. in velocity control the sequence of action is likely similar if we are reaching downwards to the items)? How would stochasticity of the dynamics play into learning this? - Figure 2: Any intuition why few-shot is better than multi-task for QueST? **Possible typos** - Page 8, line 305: viz. means visualized?

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

- The current application is only on robotic manipulation but it will be interesting to extend this to other domains

Summary

This work develops a novel framework for learning generalizable skills from demonstration data. The author’s model uses a quantized discrete latent variable model that compresses skills into a sequence of latent variables and predicts temporal sequences of actions. Their approach decodes skill by cross-attending the sequence of latent tokens against fixed positional encodings, which differs from previous works that typically condition states and actions as inputs to the decoder models. In their experimental evaluations on several benchmarks ( LIBERO and MetaWorld), the authors demonstrate their skill learning framework performs better than alternatives, and additional ablation highlights the impacts of the Quantization parameter codebook size thresholds on how much they benefit performance for inference.

Strengths

The paper is well-written and explains the author's framework in precise detail. It was interesting to suggest that a state-less sequence combined with cross-attention could yield such good performance. Given the empirically solid performance and the author's attention to ablating relevant factors of their system, the paper addresses the problems described and justifies the proposed skill model.

Weaknesses

The only major weakness of the author’s work is the limited evaluations. If the authors can justify using just three seeds across experiments, that would build more confidence. Otherwise, as this system targets robotic learning, it would have been good to see results on a real robotic system instead of just in simulation. The authors use large transformer models, so latency concerns could be relevant if their system is computationally slow. We are also concerned with the lack of error bars in Figure 4, which appears only to use a single seed for the ablation experiments, casting doubts on any conclusion the authors make from these results. The authors should clarify these details and run additional experiments to show the robustness of their results if they have not.

Questions

- How many seeds were used in Figure 4 experiments? - If the latent variables are used as key and query values, is it fair to say the generated skills can be imagined as some interpolation between a subspace of the fixed-sized vectors constructed by the positional encodings?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The major limitation we see is the lack of real robot experiments using the author’s system. For real robotic applications, such evaluation is necessary for there to be an acceptance of such methods to be deployed in the real world.

Gentle reminder!

We kindly request all reviewers to go through the rebuttals and let us know if their concerns have been addressed. Please let us know if there are any more questions/clarifications. Thank you.