QueST: Self-Supervised Skill Abstractions for Learning Continuous Control

Thanks much for your time spent reviewing and your thoughtful comments. We’ve uploaded an updated version of the paper to the website in case you’d like to review the changes we mention.

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.

Many recent works use discrete latent variable models as a mechanism to learn shared abstractions over continuous low-level skills. QueST, VQ-BeT and PRISE all do this and perform two-staged learning. However, in this work we propose several important architecture choices leading to QueST’s strong performance. Specifically, QueST encodes actions to a sequence of $n$ encodings (skill tokens) using a novel autoencoder that captures temporal correlation within an action sequence with causal convolution and masked self-attention layers.

Like QueST, VQ-BeT performs temporal abstraction by encoding a sequence of actions into one state-independent latent vector. It encodes these actions to a single output encoding using an MLP which does not explicitly model any temporal correlation, followed by an MLP decoder and continuous offset predictor. Unlike QueST, VQ-BeT’s relatively small latent space limits its expressive capacity, and thus limits its ability to learn sharable representations between tasks, as evidenced by its worse few-shot performance, worse performance with longer action chunk sizes and reliance on continuous offset prediction (whereas QueST can achieve high success rate using only the output from its action decoder).

PRISE performs temporal abstraction by learning discrete codes for state-action pairs and using BPE to group common token sequences into higher-level skills. However, BPE is known to suffer with evolving language leading to a suboptimal character-level tokenization, and it might struggle to effectively encode new actions from unseen tasks. QueST gracefully handles this by encoding new actions as combinations of its learned latent codes, and its superiority in this regard is evidenced by its stronger few-shot performance even without finetuning the decoder, as is necessary for PRISE.

A lot of ablation studies are conducted. However, the related analysis is not sufficient.

Below we mention detailed analysis which has also been added to the paper.

• Replacing FSQ with VQ still outperforms VQ-BeT in a few-shot setting suggesting that QueST’s superior performance is not only due to a better quantization scheme but also due to its architecture that flexibly maps an input sequence to multiple embeddings and allows for efficient transfer.
• It’s tempting to ground the mapping between z-tokens and actions with observation tokens with an intuition that z-tokens will define a coarse set of actions and observation tokens will aid finer action decoding. But we observe worse performance with this. We hypothesize that the reconstruction objective forces encoder and decoder for most optimal quantization at the bottleneck layer but with extra observation information the decoder might focus more on observation tokens in turn hurting the quantization. This observation goes hand-in-hand with a closely related prior work SPiRL[57] that tried the same ablation and found that state conditioned decoder hurts downstream RL.
• We observe a poorer performance in both multitask and few-shot settings with a conventional stage 1 autoencoder. This validates the QueST’s cross-attention architecture that allows for attending to all z-tokens and maintaining causality at the same time.
• We discuss the causality ablation in response to the next question.
• Figure 1b and 1d in the global response PDF illustrates the impact of decoder finetuning in fewshot IL setting. QueST outperforms all baselines in both benchmarks even without finetuning the decoder. Finetuning decoder should not be necessary in LIBERO-LONG setting, as the tasks are combination of two tasks from LIBERO-90 (pretraining set). This highlights QueST’s effectiveness in stitching trajectories using its learned skill-space. In MetaWorld, we use the ML45 benchmark that is specifically designed to evaluate transfer learning capabilities of Meta-RL algorithms, with unseen tasks being only slightly structurally similar to pretraining tasks. QueST outperforms all baselines though by a small margin but with a frozen decoder, this demonstrates the capability of QueST to effectively combine learned tokens for unseen action distributions.

Why is the causality of the encoder transformer so critical?

We ablate causal masking of encoder and decoder and observe that a fully-causal stage-1 is most optimal and a non-causal decoder does not hurt as much as a non-causal encoder does. This can be explained with a simplistic setting where the input to stage-1 are 2D trajectories of a point agent. Consider an anti-clockwise circular trajectory and an S-shaped one where the first half of the latter overlaps with the first half (semi-circle) of the former. When both of these trajectory sequences are inputted to the stage-1, a non-causal encoder will assign distinct sequences of z-tokens for both trajectories. But a causal encoder will assign the same sequence of z-tokens for the first half of both trajectories and distinct to later parts. This allows the model to re-use the z-tokens corresponding to a semi-circle for creating other shaped-trajectories that have a semi-circle in them, for example C-shaped or infinity-shaped trajectories. Thus causal masking enables better sharing of skill tokens which leads to improved performance.

Thanks much for your time spent reviewing and your thoughtful comments. We’ve uploaded an updated version of the paper to the website in case you’d like to review the changes we mention.

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

Many recent works use discrete latent variable models as a mechanism to learn shared abstractions over continuous low-level skills. QueST, VQ-BeT and PRISE all do this and perform two-staged learning. However, in this work we propose several important architecture choices leading to QueST’s strong performance. Specifically, QueST encodes actions to a sequence of $n$ encodings (skill tokens) using a novel autoencoder that captures temporal correlation within an action sequence with causal convolution and masked self-attention layers.

Like QueST, VQ-BeT performs temporal abstraction by encoding a sequence of actions into one state-independent latent vector. It encodes these actions to a single output encoding using an MLP which does not explicitly model any temporal correlation, followed by an MLP decoder and continuous offset predictor. Unlike QueST, VQ-BeT’s relatively small latent space limits its expressive capacity, and thus limits its ability to learn sharable representations between tasks, as evidenced by its worse few-shot performance, worse performance with longer action chunk sizes and reliance on continuous offset prediction (whereas QueST can achieve high success rate using only the output from its action decoder).

PRISE performs temporal abstraction by learning discrete codes for state-action pairs and using BPE to group common token sequences into higher-level skills. However, BPE is known to suffer with evolving language leading to a suboptimal character-level tokenization, and it might struggle to effectively encode new actions from unseen tasks. QueST gracefully handles this by encoding new actions as combinations of its learned latent codes, and its superiority in this regard is evidenced by its stronger few-shot performance even without finetuning the decoder, as is necessary for PRISE.

There are no results for PRISE on Multitask LIBERO-LONG.

We were unable to reproduce the PRISE results when we ran their code so we were only able to include results presented in their paper.

The paper website states: “Depending on the dataset, we also tune some key hyperparameters for the baselines”. This is vague…

Below we summarize information about hyperparameter tuning we’ve added to the appendix

• ResNet-T: We observe optimal performance with a transformer trunk with 6 layers and a hidden dimension of 256 and use those parameters for all results. We use an observation history of 10 timesteps.
• Diffusion Policy: We tuned the U-Net hidden dimension across [256, 256, 512, 512, 1024] but did not observe any performance gains. We also tune prediction and execution horizons with 16, 32 and 8, 16 respectively.
• VQ-BeT: We tune the encoder MLP dimension across (128, 256, 512) with (2, 4) layers and observe worse reconstruction loss with increase in capacity. We use residual-VQ configuration giving a similar sized codebook to QueST. We use an observation window size of 10 and tune the action window size across (1, 5, 32).
• ACT: We tune model hidden dimension (256, 512); chunk size (16, 32); and kl weight (10,100).

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.

We were able to successfully implement Diffusion Policy, VQ-BeT, ResNet-T and ACT on both benchmarks. Thus, we ran and reported those results as shown in the global response PDF. We were unable to successfully recreate PRISE results so we report the results from their paper.

There is no analysis of inference latency of the proposed method and the baselines… How does the proposed method compare to the baselines in terms of latency?

Our pipeline runs at 33Hz, which is more than sufficient for vision-based robot control where most camera systems run at 30 fps. For comparison, our implementation of Resnet-T, VQ-BeT, ACT and Diffusion Policy run at 100Hz, 100Hz, 50Hz, and 12Hz respectively. We’ll add this detail to the paper.

It might be interesting to have an ablation experiment that uses Residual VQ as used by VQ-BeT instead of FSQ…

We performed an ablation with Vanilla-VQ of the same codebook size as VQ-BeT’s Residual-VQ and observed full codebook usage throughout the training. Hence we don't expect Residual-VQ performance to vary much from the Vanilla-VQ results as Residual-VQ is mainly used to mitigate codebook collapse which doesn’t happen in our case. Moreover, the VQ-BeT paper doesn't mention if they ensured their VQ ablation to be of the same effective codebook size which might have led to poor performance in their ablation experiments. It would still be interesting to validate this hypothesis but due to time constraints it might not be viable by the end of the discussion period.

The paper website states that “To ensure fair comparison of different model architectures, we use the same input modalities and same observation & task encoders for all baselines…”. This is important information that should be in the…paper. It is also unclear whether this only applies to the Meta-World experiments or also the LIBERO experiments.

We’ve moved this detail to the appendix and will add a line confirming that this is the case for all experiments.

It is not specified where the “provided scripted policies”…can be found.

They can be found in the official Metaworld codebase.

Figure 3b is not described or mentioned in the text.

Fixed

Are the updated results on Meta-World going to replace the current results in…the paper?

Yes

Is the code going to be released publicly?

This will happen alongside the camera ready version of the paper.

Thanks much for your time spent reviewing and your thoughtful comments. We’ve uploaded an updated version of the paper to the website in case you’d like to review the changes we mention.

Should probably include analysis on what the latent 𝑧’s ended up learning. Do they actually have some temporal abstraction going on? Can we associate them?

Thanks for the suggestion. First, because the autoencoder learns to compress a sequence of T actions into a smaller sequence of n latent codes (n<T), there is automatically some temporal abstraction. Next, the t-SNE plots on the website clearly show how similar motions are aligned with one another. We visualize an example with two tasks for lifting a bowl and placing it on a plate (on right) and on a drawer (on left) on the website. The z-embeddings align for the motion of approaching and placing the bowl downwards, while deviating for moving leftwards and rightwards, demonstrating how QueST embeds diverse motion primitives into discrete tokens in a semantically meaningful way. Also, we would like to point out that there is no explicit loss for semantic alignment, it’s the reconstruction objective along with the bottleneck layer that guides the learning to extract shared representations. We’ll move these plots and discussion to the appendix.

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

While we are making this assumption of structure as an inductive bias for modeling, we do not train the model on any special datasets or use any explicit labels for primitive skills. We train and evaluate on standard datasets that many recent behavior cloning works use and our architecture is designed to capture the structure of these datasets. Motions involved in everyday manipulation tasks naturally have commonalities that can be shared and reused across tasks. This is not an assumption but a property of the data that our model most effectively leverages.

Regarding partial observability, you are correct that we didn't have to deal with partial observability for our experimental settings. The main contribution of our work is a method to tokenize continuous actions. That being said, our method can easily be adapted to handle partial observability. For example, one popular way to do that is to simply pass in a stack of historical frames, and it would be very simple to pass in several observation vectors to our policy prior as opposed to the one we pass in now.

It is not clear to me how this approach handles multimodality (in the sense of action distributions?)

First, the policy prior outputs latent actions as a categorical distribution and the inherently multimodal categorical distribution over latent actions lends itself well to multimodal data. This phenomenon has been studied in the past in papers such as BeT [56] and VQ-BeT [34]. The Diffusion Policy paper shows how BeT suffers from mode collapse due to lack of temporal consistency which both VQ-BeT and Diffusion Policy resolves by probabilistically predicting a chunk of actions. Likewise, QueST also predicts a categorical distribution over skill-tokens (this chooses mode as per training distribution) and decodes a chunk of actions (this brings temporal consistency). This effect is studied in several prior works [34, 56, 15, 72].

Ablation on "causality": I feel the word "causality" is misused... Is this not still modeling correlation only? How is this causal exactly?

You are right that causality here means not to look at the future data. There seems to be a namespace collision problem with the word. In the machine learning literature the word causality is used to refer to precisely the type of masking we are using in our architecture. For our convolutions we use the causal convolutions from wavenet [63] and for the transformer we use the square causal mask defined in pytorch, TensorFlow and HuggingFace. However we understand how this might cause confusion so we’ll replace the word ‘causality’ with ‘causal masking’, emphasizing how this boosts performance.

Equation 1, perhaps write out what round_std is.

It's the nearest integer rounding with straight-through gradients. We’ll update the manuscript.

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?

We follow the original transformer decoder architecture which consists of alternate masked self-attention and cross-attention layers. Thus we self-attend to positional embeddings with a causal mask as attending to future positional embeddings won’t provide any extra information for the i^th embedding.

lines 186-187: Do the authors think that this is because the actions can be invariant…? How would stochasticity of the dynamics play into learning this?

Yes, state-independent abstractions are generalizable as many distinct tasks share common motion primitives, especially with velocity control which offers further invariances. This is more performant than its state-conditioned counterpart as it forces action information through the bottleneck layer, preventing the decoder from ignoring it and focusing on states, hurting the quantization. A closely related prior work, SPiRL[57], tried the same ablation with similar results.

The stochasticity of dynamics are accounted for at execution time. The model predicts actions for the next 1-2 seconds, executing them in a receding horizon fashion and intermittently replanning to account for stochastic dynamics.

Figure 2: Any intuition why few-shot is better than multi-task for QueST?

The few-shot version does better because it is also trained on the data from the LIBERO-90 suite, and it transfers action representations to the tasks in the LIBERO-LONG suite. We'll add details to the paper to make this more clear.

Thanks much for your time spent reviewing and your thoughtful comments. We’ve uploaded an updated version of the paper to the website in case you’d like to review the changes we mention.

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.

For behavior cloning settings, which tend not to show significant variation across random seeds, performing experiments with comparatively low numbers of random seeds is fairly common and accepted in the literature. For example, BAKU and VQ-BeT [34] only use 1 seed per experiment. ACT [72], Diffusion Policy [15] and 3D-Diffusion Policy use 3 seeds while PRISE [73] uses four.

For our multitask experiments, we now report 4 seeds for LIBERO and 5 seeds for MetaWorld. As evidenced by the error bars, these results have very low variance (<1%) across seeds for both benchmarks. Hence, such a small number of seeds are sufficient to convincingly demonstrate the statistical significance of our reported results.

For our fewshot experiments, which have more variance across random seeds, we’ve run further random seeds for LIBERO benchmark. Specifically, we run three fewshot training seeds for each of the three pretraining seeds, resulting in a total of 9 seeds. Please see the updated results in the global response PDF, summary: QueST outperforms best performing baselines by 14% (absolute). Since variance across pretraining seeds is very low (following multitask results) and all stage-2+decoder parameters are finetuned, the fewshot results have very low reliance on pretraining seeds and hence our seeding methodology is justified. For MetaWorld, we report results for both settings across 5 seeds and observe all methods to perform comparatively with QueST slightly better than others.

In order to more rigorously verify these assertions, we’ve performed a t-test, whose results are contained in the PDF response. In summary we see statistically significant improvements in performance across all benchmarks except Metaworld fewshot.

The latency concerns could be relevant if their system is computationally slow.

Our pipeline runs at 33Hz, which is more than sufficient for vision-based real-robot control where most camera systems run at 30 fps. For comparison, our implementation of Resnet-T, VQ-BeT, ACT and Diffusion Policy run at 100Hz, 100Hz, 50Hz, and 12Hz respectively. We’ll add this detail to the paper under a new subsection in section 5.

It would have been good to see results on a real robotic system instead of just in simulation.

We agree! Unfortunately we do not have the capacity to run real robot experiments at this time but this is an important limitation we’ll add to the limitations section of the paper.

Regarding lack of error bars in Figure 4. How many seeds were used in Figure 4 experiments?

Unfortunately this is a very computationally demanding experiment to run and since our university-provided compute is swamped with other authors working on NeurIPS rebuttals, we might not be able to show more seeds for this experiment by the end of the discussion period. That being said, we will add two further random seeds before the camera ready version is released. While this is an unfortunate circumstance, hopefully you'll agree that the lack of seeds in this experiment don't meaningfully detract from our overall claims about the effectiveness of QueST in multitask and fewshot settings.

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?

Not quite. Since decoder cross attends to latent skill tokens, their Key vectors (K) and Value vectors (V) are used with Query vectors (Q) from positional encodings. Thus the output actions (skills) lie in the space spanned by the value vectors from skill tokens. Additionally, the softmax and GeLU non-linearity makes the relationship more complex than simple interpolation. A more accurate way to describe might be that the output actions are a weighted combination of transformed vectors derived from skill tokens, where the weights are contextualized based on positional encodings (for temporal correlation).