QueST: Self-Supervised Skill Abstractions for Learning Continuous Control
Authors: Atharva Mete
Haotian Xue
Albert Wilcox
Yongxin Chen
Animesh Garg
Venue: NeurIPS 2024 poster (2024)
Keywords: Behavior Clonning
Action Tokenization
Self Supervised Skill Abstraction
Few-shot Imitation Learning
Primary area: robotics
TL;DR: QueST is a multitask latent-variable behavior model that learns sharable low-level skills by representing temporal action abstractions (1-2 secs motion) with a sequence of discrete codebook entries (skill-tokens).

Abstract: Generalization capabilities, or rather a lack thereof, is one of the most important unsolved problems in the field of robot learning, and while several large scale efforts have set out to tackle this problem, unsolved it remains. In this paper, we hypothesize that learning temporal action abstractions using latent variable models (LVMs), which learn to map data to a compressed latent space and back, is a
promising direction towards low-level skills that can readily be used for new tasks. Although several works have attempted to show this, they have generally been limited by architectures that do not faithfully capture sharable representations. To address this we present Quantized Skill Transformer (QueST), which learns a larger and more flexible latent encoding that is more capable of modeling the breadth of low-level skills necessary for a variety of tasks. To make use of this extra flexibility, QueST imparts causal inductive bias from the action sequence data into the latent space, leading to more semantically useful and transferable representations. We compare to state-of-the-art imitation learning and LVM baselines and see that QueST’s architecture leads to strong performance on several multitask and few-shot learning benchmarks. Further results and videos are available at https://quest-model.github.io.

QueST: Self-Supervised Skill Abstractions for Learning Continuous Control
Atharva Metel, Haotian Xue', Albert Wilcox', Yongxin Chen Animesh Garg!²
Georgia Institute of Technology, 2NVIDIA
Abstract
Generalization capabilities, or rather lack thereof, is one of the most important unsolved problems in the field of robot learning, and while several large scale efforts have set out to tackle this problem, unsolved it remains. In this paper, we hypothesize that learning temporal action abstractions using latent variable models (LVMs), which learn to map data to compressed latent space and back, is a promising direction towards low-level skills that can readily be used for new tasks. Although several works have attempted to show this, they have generally been limited by architectures that do not faithfully capture sharable representations. To address this we present Quantized Skill Transformer (QueST), which learns a larger and more flexible latent encoding that more capable of modeling the breadth low-level skills necessary for a variety of tasks. To make use of this extra flexibility, QueST imparts causal inductive bias from the action sequence data into the latent space, leading to more semantically useful and transferable representations. We compare to state-of-theeaat imitation learning and LVM baselines and see that QueST's architecture leads to strong performance several multitask and few- shot learning benchmarks. Further results and videos are available at https: //questt model.. gi thub.
Introduction
One of the grand goals of robotic learning is general- purposs model that can learn from complex multitask demonstration data and generalize new tasks a zero-shot or few-shot manner. While such purpose models have hecome ubiquitous in natural language (NLP) [66, 72, 52, 53, 63] and computer vision (CV) [77, 33, 7]. they have eluded robotics researchers. Whereas CV and NLP can achieve positive transfer by scaling up models trained on internet scale datasets [33, 72, 56, 62, 39]. even large scale robot data collection efforts [10, 11, 47, 20, 31] have been insufficient for this approach. To that end, posit that in order to achieve positive transfer in robotics, it is important to design architectures that specifically lend themselves to efficient cross-task transfer.
There has recently been surge of work towards the goal of learning generalist policies from large, diverse datasets. Several papers have used techniques such action discretization [57, 16, 35, 10, 11, 47, 55] and implicit models [23, 15, 25] to model multimodal action distributions. In particular, the behavior transformer line of work shows that carefully discretized action space combined with GPT-style transformer leads to impressive capabilities modeling multimodal behavior distributions
[57, 16, 35]. another vein, several works have attempted scale demonstration data and achieve positive transfer of low-level skills between high-level tasks [10, 11, 47, 46, 18, 12, 55, 3]. While these works have shown some transfer, for example applying policies for known tasks to unfamiliar objects, they have generally failed to achieve transfer of low-level skills to novel tasks [4]. We
"Correspondence amete7dgatech.edu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

hypothesize that in the relatively low-data regime of robot learning, it is promising to explicitly force the model to learn sharable representations. To that end we study latent variable models (LVMs), which learn to map data to compressed latent space and back, introducing an information bottleneck which encourages the model to learn shared representations across the training data. Specifically we consider the application of LVMs to learn low-dimensional representations of action sequences. Such representations are termed temporal action abstractions or motion primitives - in this paper we refer to these abstractions as "skills'.
wide body of work has considered the application of LVMs to robotics. One line of work learns temporal action abstractions (skills) in continuous latent space with a Gaussian prior [41, 50, 59]. While this line of work showed some initial promise of learning latent plans, it has failed to scale to difficult multitask settings due to the loose nature of the latent structure and posterior collapse issues that inhibit the learning of shared representations. On the other hand, recent work in CV and NLP has shown that vector-quantized discrete latent spaces are capable of learning semantically meaningful representations from data like phoneties in speech [9, 6] or melody in music [17, 2]. This insight, along with prior work showing that discretized action spaces can help to address the multimodality problem in when learning from large datasets [57, 16, 35, 14], motivates methods learning temporal action abstractions with discrete latent spaces. Several recent works have set out to do this [35, 74, 76, 30), showing some degree of positive transfer between tasks in multi-task and few-shot settings. However, they are generally limited by architectures that do not faithfully capture transferable representations [35, 74], or depend on state prediction and state- based objective functions which are impractical for many real robot tasks [76, 30].
In this paper, we present Quantized Skill Transformer (QueST), a simple yet novel architecture for learning generalizable low-level skills within a discrete latent space. The key insight behind QueST is its ability to flexibly capture variable length motion primitives by representing them with a sequence of discrete codebook entries. We achieve this through unique encoder-decoder architecture primarily designed to impart causal inductive bias in action sequence data into the latent space. Such formulation enables to employ powerful sequence modeling approaches to plan and composably reason within the space of low-level skills. Through our experiments, we show that autoregressive modeling of these latent skills with GPT-like transformer outperforms challenging robotic manipulation benchmarks, where QueST shows an 8% improvement in multitask and 14% improvement in few-shot imitation learning over the next best baselines. We also conduct detailed ablation and sensitivity study to validate our key architectural design decisions.
baselines on
2 Related Works
The proposed framework in this paper introduces a methodology for self-supervised skill abstraction, followed by decision-making within this skill space. Several related works have explored similar sub-directions such as decision-making in the latent space and decision making with transformer:
2.1 Learning from Offline Data
Behavior cloning (BC) [51] aims to learn a policy by directly mapping observations to actions, and is typically trained end-to-end using pre-collected pairs of observation and behavior data. While this on its surface is simple supervised learning problem, there are several properties of robot demonstration data that should be considered when building BC systems. First, large BC datasets collected from a variety of human demonstrators tend to contain data sampled from multimodal distributions. To address this, some works opt to sample actions from Gaussian Mixture Models (GMM) [43], while others explore implicit models including those derived from mnergy-based models [23, 29] or diffusion models [15, 71, 14, 68, 28]. The Behavior Transformer (BcT) line of work [57, 16, 35] shows that wansformer-basee categorical policies in carefully discretized action spaces do a good job handling multimodal demonstrator distributions and QueST builds upon this by contributing more capable discrete latent skill model.
Another key property of robot demonstration data is that sequential actions are often highly correlated with one another, and exploiting this can lead to stronger performance while ignoring it can lead to policies which are susceptible to temporally correlated confounders [61]. Recently several works have set out to handle this by predicting action chunks. For example, the Action Chunking Transformer (ACT) line of work [73, 24] shows that transformer trained CVAE [60] to output chunks of

---

actions performs well for wide variety of manipulation tasks, and diffusion policy [15] shows across the board improvements when predicting action chunks. As discussed in detail in Section 2.3, QueST builds on long line of work which handles sequential correlations through action abstractions [41, 50, 59, 35, 74, 76, 30].
2.2 Multi-task and Few-shot Imitation Learning
In the past, robot learning researchers have approached multi-task decision making settings using wide variety of methods such as supervised pre-training and fine-tuning [20, 42], meta-learning
[22, 19] and action retrieval [48, 44]. There has recently been large focus on multi-task language- conditioned imitation learning for robotics with several papers attempting to address the problem by training large models on large demonstration datasets [10, 11, 47, 46, 18, 12, 1]. While these papers achieve impressive multitask results, they mostly rely sufficient data coverage and fail to generalize beyond their training distribution [4]. Thus, they lack abstractions that can readily be applied to learn new tasks, especially in low-data regime. On the other hand LVMs like QueST designed to learn sharable representations that can be applied to new tasks.
2.3 Decision Making in Learned Latent Spaces
LVMs, modeled by a paired encoder-decoder, have found extensive applications in computer vi- sion [65, 32, 8] and generative models [56, 21, 13]. Recent studies also demonstrate the utility of latent space representations in robot Decision-makinn
spanning offline RL [58, 50, 40J, imitation learning [14, 67, 37, 35], and temporal action abstraction [57, 74, 75]. Most similar to our work are those that learn temporally abstracted discrete latent skill spaces. PRISE [74] learns single-step state-action abstractions within discrete space and then does temporal abstraction by applying BPE tokenization. While this method shows promise in learning multi-task and few-shot policies, BPE is well suited for text and is known to suffer in domains with highly dynamic vocabularies, in robotics its equivalent to varying action distributions across unseen tasks. TAP [76] and H-GAP [30] utilize self-supervised auto-encoder to learn skill codes, but their functionality relies on Model Predictive Control (MPC) using state prediction with ground-trutt state-conditioned objective functions, which make it difficult to apply to real-world manipulation tasks. VQ-BeT [35] bears the strongest resem- blance to QueST, also using a pre-trained discrete latent skill space to discretize the action space for a transformer-base
policy prior. However, their quantization approach does not leverage the inherent structure in action sequences, limiting the representational capabilities of the latent space. Thus it's shown to heavily rely on an continuous offset predictor for best performance. Unlike these works, QueST's learned latent space is highly flexible yet structured and expressive, allowing it to effectively model action distribution across many distinct tasks in a meaningful shared representation.
3 Preliminary
3.1 Problem Setting
We consider dataset D () at),,,,,)..... (rr,,ar))) 1im Le M Ni
consisting of M robot inter- action trajectories where a, is continuous-valued action and O is tuple consisting of high- dimensional sensory observation. The data is collected via either human teleoperation or scripted policies for M different task each with label La. In our setting, O, consist of RGB image observa- tions from the front camera and gripper camera (if available) along with proprioceptive state of the agent. L3 is natural language description of the tt task but can also be one-hot encoding.
(O,,,.,,,,,,,,)))
3.2 Finite Scalar Quantization
We build on Finite Scalar Quantization (FSQ) [45] as a discrete bottleneck in our model. It's a drop-in replacement for Vector Quantization (VQ) layers in VAEs with simple scalar quantization scheme. Here the input representation is projected to very few dimensions (typically 3 to 5) which are then bounded and rounded, creating implicit codebook.
round_ ste(f where is the bounding function
Given a feature vector e E d to quantize, instead of learning a parameterized codebook [65] and quantizing e by matching the nearest neighbor in the codebook, FSQ quantizes by first bounding
3

---

Stage Self-supervised Skill Abstruction
Stage II: Decision Making with Skill Tokens
sampled action sequeence
Inteence
comseptrrinng F&D enbooiiing
1 2 A)
Causal Conv-1D
Autoregressive Transformer
mp(Z(,,))
Gragee
CLPjere-hot
Food treedings
MLP
ResNet
do (
Transformer Decoder
Proprioception
deel
4113
Observetion history
predicted oction sequence
Figure Overview Quantized Skill Transformer: we factorize the policy that outputs action based on task descriptions e and observations o encoding into two parts: π T(Alo, e) = the (,, e). where Z/ a sequence of skill tokens for the action sequence A. In Stage I, we learn skill ahstraction in self-supervised way with quantized autnencoder. In Stage II, we learn skill-based policy in the style of next-token prediction using multi-modal transformer.
it into certain range with (e.g. [a/2] 3 tanh(e)) and then rounding each dimension into integer numbers directly with straight-throus gradients (round ste). a E 2d defines the width of the codebook for each dimension (e.g. [8, 5, 5, 5), d 4). Finally, it is easy to that the size of the quantization space is II 04 An MI.P can be used to transform into continuous space of required dimension further.
04.
A common problem with vector-quantized codebooks (VQ) [65] is the under-utilization of the codebook. Recent works have attempted address by heuristics like reinitializing the codebook, stochastic formulations, or some regularization [34, 69]. In contrast, FSQ achieves much better codebook utilization for large codebook sizes with much fewer parameters and simplified training without any auxiliary losses or aforementioned tricks. Due to its simplicity and proven benefits, we use FSQ our main experiments, but since many prior works in this space use VQ we also perform an ablation with it (see section 5.6).
4 Method
In this section, we describe the key ideas behind Quantized Skill Transformer. Section 4.1 we present encoder-decoder architecture, which is designed to provide the flexibility to learn wide range of skills with inductive biases to ensure that the learned skills are useful. In Section 4.2, we detail skill prior, which we train to autoregressively predict codebook skills, Our full pipeline is shown in Figure
4.1 Stage I: Learning the Skill Codebook
As motivating example, consider the task of lifting pot and placing it on stove beside. This consist primitives like reaching the pot, grasping it, lifting it to a certain height, reaching the stove and finally placing it on the stove. Each of these primitives are of variable lengths, and to properly model these skills it is important to learn a latent skill space with the flexibility to model all of them. At the same time, it is important that the learned skills semantically meaningful can be reused for new tasks, for example reusing the reaching skill for an object lifting task. In order to address these desiderata, introduce the novel autoencoder architecture shown in Figure 1 consisting of an encoder Φy and decoder Va--
that they
The input to the encoder dy is action sequence T-1 sampled from the dataset, which we pass through several ID causal strided-convolution layers [64]. This step reduces the sequence length to achieve the desired temporal abstraction depending on the stride lengths and the number of layers.

---

We follow the convolutional layers with masked self eeff-attention layers for sequence modeling, With a downsampling factor of F, the encoder outputs total T/F embeddings. The embeddings are then quantized using FSQ as per the equation into n discrete latent codes {=*] termed as skill tokens:
(s',..., FFSS ".,, (, ((4T-1))
Having an input sequence of actions mapped to multiple skill tokens gives this architecture more flexibility to model complex sequences of actions. At the same time, each component of the encoder is causal, meaning that an output representation at a position / cannot depend on input from any future timesteps. We found this inductive bias encourage the model to learn semantically useful action representations by modeling the inherent causality in the action data. We validate this design choice in the ablations. (see section 5.6)
Typical autoencoder decoders are simply mirrored versions of the encoders, but this would prevent the decoder from attending to all quantized codes. This is important because individual codes do not represent anything meaningful but a sequence of codes represents a particular meaningful motion [45]. In order maintain causality while attending to all codes, the decoder the cross attends between fixed sinusoidal positional embedding inputs and the skill tokens, similarly with [73]. The architecture is a transformer decoder block consisting of alternate masked self-attention and cross-attention layers, after which the output embeddings are projected back to the original action dimension using an MLP layer. Thus, given a sequence of skill codes, tho reconstructs the original action
(āz,.... T++T-1) = v/u(2' ((,)(",,,,")
(3)
As in [35], the autoencoder is trained by minimizing the &ı reconstruction loss:
Lrecon (0) = |/ (FSQ( (( ama+T-1ll-
Unlike prior work which often conditions on the state well the actions [30, 5, 74, 37], we choose to learn state-independent abstractions that solely capture motion primitives irrespective of the current scene or task. Through our experiments show that our model learns generalizable abstractions that are shared and can be transferred across tasks.
4.2 Stage II: Learning the Skill Prior
After training the encoder de and decoder the, train skill prior Tzl to predict skills Z 2 corresponding to the demonstrator action distribution conditioned on a task embedding and a length sequence of image observations and proprioception inputs, o= (4r-k. Pt-n), (4e, Pc). We encode image observations with separate learned vision encoder for each camera view and encode proprioception using an MLP encoder, all of which are trained end-to-end with the rest of the skill prior. The observation token T,  for timestep is obtained by concatenating outputs from all the aforementioned encoders. Task embeddings are designed specifically for each environment suite, as discussed in more detail in Section 5. See Appendix B for more details about the encoders.
Because skill tokens are highly dependent on one another according to the complex nonlincar representations learned by the autoencoder, it is important that the skill prior has the modeling capacity to reason about these dependencies. To achieve this, we employ a decoder-only transformer to model the distribution of skill tokens n.P (T|TTTrna e) autoregressively as:
T (Z e) = II n <=*
2 21:11 -T e)
where 
 is learnable start token that marks the start of skill tokens. We add sinusoidal positional embeddings only to the skill tokens. To optimize the skill prior, we sample sequence of demonstrator
actions d-:+T-1 and use the trained encoder de to extract a latent skill vector Z, = 2 Imm according
to Equation 2. Then, optimize TiP using the following negative lig-likelihood loss:
Craex(p) = - log u. (2,(T) (T e).
(6)
The full skill prior pipeline is shown in Figure 1.
5

---

Few-Shot Finetuning: For few-shot finetuning on new tasks, we use model pre-trained on large set of tasks and finetune on a small number of demonstrations (5 in our experiments) from the held-out task. Although finetuning only stage-2 is enough, we empirically found that finetuning the decoder on the predicted skill tokens gives boost in the performance. Specifically, we finetune the decoder using following decoder loss:
Ldccoder (0) = (looss(2)) 0u:+T-1lla
where sg is the stop gradient operator. We present the results with and without decoder finetuning both. Additionally, we note that the encoder is still frozen in this setting.
4.3 Inference with Quantized Skill Transformer
At inference time, QueST uses the skill prior Ty alongside the decoder the to sample actions. Conditioned the encoded observation history T and task embedding e, use top-k sampling
with a temperature of r to autoregressively sample skill vector 2 ~ x <- e) from the skill prior. In practice, the decoder to map the skill vector back to the action space, producing a sequence of predicted actions - the 2). In receding horizon fashion, we execute the first T. S T actions before replanning.
find i 5 and 1 to work well across all environments. Then, we
5 Experiments
We design the experiments to empirically evaluate the performance of Quantized Skill Transformer in three practical settings: (1) Multitask IL, (2) Few-shot transfer, and (3) Long-horizon II. Lastly, we perform some ablations to empirically justify our model design choices.
5.1 Benchmarks and Baselines
We use the following benchmark suites to evaluate in the settings discussed above:
LIBERO [38] lifelong learning benchmark featuring several task suites consisting of variety of language-labeled rigid- and articulated-bool manipulation tasks. Specifically, we evaluate on the LIBERO-90 suite, which consists of manipulation tasks, and the LIBERO-LONG suite, which consists of 10 long-horizon tasks composed of two tasks from the LIBERO-90 suite. As described study the multitask IL. few-shot transfer
in more detail below, we use the LIBERO benchmark and onng-horizon II. settings. Because tasks from this benchmark are language-annotated, we use the output of a frozen CLIP [54] encoder for the task conditioning input e.
Meta World [70] features a wide range of manipulation tasks designed to test few-shot learning algorithms. We use the Meta-Learning 45 (ML45) suite which consists of 45 training tasks and difficult held-out tasks which are structurally similar to the training tasks. We use this benchmark to test multi-task and few-shot learning. Because this benchmark does not include language labels, use learned task embeddings e for task conditioning.
We compare to the following baselines, which include similar discrete LVM pipelines as well as state of the art imitation learning algorithms:
Baselines:
1. The ResNet-T model from [38], which encodes observation and task instructions using ResNet-18 with FiLM [49], applies a transformer sequence model, and uses a GMM head to predict actions.
The UNet-based diffusion policy from [15], which uses a convolutional UNet to map samples from Gaussian prior to action samples from the demonstrator distribution according to learned denoising process.
ACT [73], which trains a transformer as a CVAE [60] to predict action chunks.
VQ-BeT [35], which learns a discrete latent space using a VQ-VAE [65] and uses a transformer to predict discrete latent codes.
5. PRISE [74], which first quantizes observation-actin pairs and performs temporal abstraction
using byte pair encoding (BPE) to learn skill token vocabulary, which it uses as an action space for few-shot learning.

---

120
10
014
70
th44
20
18.2
60-
60-
$61
50-
50
40
32.3
20
70
10
10
NcsNes1 AG1
Uifisknn FRGE
V0.50T 0v55
Que51
Failyy
Pallyy
(a) Multiask
shot LIBRR-LLON.
Multiask LIBERO
Figure 2: Multitask performance on LIBERO-90 (a) and LIBERO-LONG (c), and few shot performance on LIBERO-LONG (b). For (a) and (c) we train on the datasets described in Sections 5.2 and 5.4. For (b) we finetune the model from (a) on condensed dataset described in Section 5.3, Results show the mean and error bar represents standard error across four random seeds for multitask experiments and nine random seeds for fewshot experiments. Results PRISE are taken from Zheng et al. [74], the others we reimplemented and ran ourselves.
For a detailed discussion between QueST and the baseline methods, please refer to Appendix B.5,
5.2 Performance on Multitask BC
We evaluate the goal-conditioned multi-task imitation learning capabilities of QueST and the baselines using the LIBERO-90 and ML45 benchmark suites. For LIBERO-90 the learner receives 50 expert demonstrations per task from the author-provided dataset. For ML45, we use the scripted policies provided in the official Metaworld codebase to collect 100 demonstrations per task. We evaluate the model at the end of training and for each task run 40 evaluation rollouts (50 for MetaWorld) starting from the initial states selected sequentially from a predefined set. We report the aggregated results across 4 seeds (5 seeds for Meta World).
In Figure 2a, we present the average success rate across 90 tasks in LIBERO-99 against the afore- mentioned baselines. Quantized Skill Transformer achieves sate-offthhe-art results on LIBERO-90 benchmark, outperforming the baseline VQ-BeT and Diffusion Policy by a margin of and 13% respectively. We attribute its performance its learned latent space that enables effective knowledge sharing across tasks. While VQ-BeT also shows strong performance, we see that QueST's architecture lends itself better sharing representations across tasks. Our implementation of ResNet-T achieves significantly better performance than the reported number (16.8%) in [38] but is still lower than QueST. Figure 3a shows the average success rate across 45 tasks in ML45 benchmark. Being simpler benchmark. all methods perform almost similar in Multitask- IL setting which consistent with the trend observed in [74].
We attribute the reasonably good performance of the diffusion policy to its nature as a latent variable model, which employs continuous latent variable with the same dimensionality However, the consistent outperformance of QueST over the diffusion policy provides compelling evidence for the benefits of using bottlenecked latent variable. This bottleneck encourages the model to learn shared representations, resulting enhanced performance. While both VQ-BeT and PRISE employ latent bottleneck, VQ-BeT's architecture neglects the inherent inductive biases in the action data, which we believe results in a less werl-structured latent space. PRISE incorporates this using latent forward transition model, but is bottlenecked by the use of BPE which posit is not suitable for such a dynamic latent space.
the actions.
5.3 Few-shot Transfer to Unseen Tasks
In this setting we take the pretrained model from section 5.2 and test its 5-shot performance unseen tasks from TIBER--LONG and held-out set in ML45. We sample only five demonstrations for each task, generate the skill tokens using pretrained encoder and them to finetune the skill prior and the decoder as described Section 4.2. We also present the results without finetuning the decoder (frozen via in figure 2b & 3b) to validate its generalization unseen skill tokens sequences.
Figure shows the average success rate for 5-shot II. across unseen tasks in LIBERO-LONG. QueST achieves SOTA performance, surpassing all other basclines by absolute margin of 14%.
7

---

90.8
90.3
10
87.6
917
96.1
66.7
Miah Succes Rate
60
50
40
20
0.
ResNetT
ACT
Difusion PAISE
Va- 0JeST
Policy
Palcc
(a) Multitask on ML45
(b) Few-shot on ML45
Figure 3: Multitask and few-shot success rate on the Metaworld ML45 task suite. In (a) we train on the dataset described in Section 5.2, and in (b) we finetune the model from (a) using 5 demonstrations each from set of held out tasks. Results show the mean and error bar represents standard error across five random seeds. Results for PRISE are taken from Zheng et al. [74], and the others we reimplemented and ourselves.
VQ
Obs. Cond.
Mirror Dec.
Ours
LIBERO-90
81.2 + 0.6
$
86.3 = 0.9
88.6 = 0.4
Few Shot
62.5
+ 2.0
61.3 + 2.2
45.4 2.0
68.8 = 1.7
Table I: Success after ablating design details of QueST. We present the across four random seeds error tolerances show the standard error.
Though we see marginal drop of 1.5% without decoder finetuning, it still outperforms all the baselines. These results highlight the superiority of QueST in learning transferable representations of action abstractions and effectively leveraging them for downstream decision making. For fair comparison, we also tried fine-tumi fine-tuning the decoder of VQ-BeT but did not observe any gains from it. VQ-BeTT struggles in this setting as it heavily relies on offset head to output continuous action corrections which requires more data-samples for sufficient coverage the continuous action space. Figure 3b shows the average success rate for 5-shot IL across 5 unseen tasks in World. Similar to multitask results, all methods perform comparably, with QueST showing a slight improvement over the others. QueST leverages its learned skill tokens to compositionally model their distribution for an unseen task in just demonstration examples. For few-shot evaluation protocol please refer Appendix D.I.
5.4 Long-horizon BC
In this setting, we aim to purely study and compare the performance of our model on long-horize tasks. We train the model (both stages) solely on LIBERO-LOWD complete dataset (50 demonstra- tions per task) and evaluate with the same scheme as described earlier.
Figure 2c shows the average success rate across 10 LIBERO-LONG tasks. All LVMs perform significantly worse than the ResNet-T model. We attribute this to the relatively small size of LIBERO- LONG dataset which is just not enough to learn good latent space in LVMs. QueST still outperforms all LVM baselines by large margin demonstrating its long-horizon modeling capabilities.
Overall, we see that our model outperforms baselines like VQ-BeTT in multitask settings, showing stronger modelling capacity. At the same time, it has the correct latent structure outperform baselines like diffusion in few shot settings, especially even with frozen decoder, indicating strong generalization capabilities of learned skill-space.
5.5
Latency
Our pipeline runs at 33Hz, which is more than suffcient for vision-based real-robot control where most camera systems run 30 fps. For comparison, our implementation of Resnet-T, VQ-BeT, ACT and Diffusion Policy run at 100Hz, 100Hz, 50Hz, and 12Hz respectively.
Non Causal da
Non Causal the
Fully Non Causal
Ours 88.6 + 0.4
LIBERO-90
82.0 + 1.6
85.1 + 1.8
78.5 + 0.5
Few Shot
58.8 + 3.0
61.6 + 2.5
56.1 + 1.8
68.8 + 1.7
Table Success rates after ablating the causality QueST. We present the mean across four random seeds and tolerances show the standard error.

---

LIRERO-50
Fnw Shot
1924
2048
4096
Downs Tplle fastor
Cod2bioorsde 202e
(a) Downsampling factor sensitivity.
(b) Codebook size sensitivity.
Figure 4: We conduct sensitivity experiment across downsampling factors (a) and codebook sizes (b) the LIBERO benchmark. For fix sequence length 32. Overall, we that the few-shot version more sensitive to hyperparameters and that F with 1024 codebook vectors are good choices.
5.6 Ablations
We validate the proposed architecture by ablating some of its key design decisions. All the ablations performed on LIBERO studying their effects on both multitask and fewshot IL settings.
1. Vector Quantization: We replace the FSQ layer with a Vector Quantization layer of nearly the same codebook size, shown in the VQ column of Table 1. We see that FSQ's superior codebook utilization leads to an improvement in performance.
Many prior condition the action decoder with current observation [76, 30]. We experiment with this by appending observation tokens to the skill tokens and allowing the transformer decoder to jointly cross-attend to both, shown in Table We see that conditioning on observations leads to a deterioration in performance.
Observation Conditioned Decoder:
Following a typical autoencoder design. we use decoder that mirrors the encoder, using transposed convolutions instead of strided convolutions, and with the strides in reverse order as in the encoder. This decoder directly takes skill-token embeddings as input and outputs the continuous actions, and results are shown in Table We see that this method performs worse, suggesting attending to all quantized codes in our decoder does, is important for faithfully predicting actions.
3. Mirrored Decoder:
We ablate the use causal layers in various parts of our network in Table 2. see that removing causality from any part of our architecture leads to worse performance, suggesting that causal masking imparts the model with helpful inductive bias for modeling robot action data.
Causality:
We also perform a sensitivity experiment over several hyperparameters including downsampling factor and codebook size in Figure 4. Across the board we see that the hyperparameters are more important in the difficult few-shot learning setting. In Figure 4a we see that both algorithms have the best performance with modest downsampling factor of F 4, and Figure 4b we see that QueST does well with 1024 codebook vectors. For more discussion on ablations please refer Appendix C.
5.7 Latent Skill-Space Analysis
We present a t-SNE visualization (Figure 5) illustrating the learned skill-space across multiple set of similar tasks. We consider four different combinations of similar tasks to effectively examine the z-embeddings corresponding to their trajectories. Each data point in the plot represents a vector of z-embeddings specific timestep throughout the entire episode, with decreasing transparency indicating temporal progression. We show that the QueST encoder learns a semantically meaningful skill-space that encodes shared representations of similar motion primitives across different tasks. This analysis includes the first 11 tasks from LIBERO-99 90. For better comprehension, we encourage readers review the corresponding rollouts on the website. Notably, the skill-space learning happens in the first stage training which does not make use of any task labels.
Conclusion
We present Quantized Skill Transformer, a novel LVM architecture for learning sharable skills in discrete latent space. The key idea behind QueST is to represent action sequences as series

---

betton trveet of oidect cruver calnet
caliee
Stonee
padra
arower oi caliica drowm the chese
the black boal calamat
Bcal
rigay
applactees
Figure 5: t-SNE visualization of skill-token embeddings. Here, the transparency decreases as the episode progresses. The overall patterns clearly shows how similar motion primitives like approaching, picking and placing from different tasks are aligned with one another.
of codebook vectors, and demonstrate that using causal convolutions and masked transformers provides an inductive bias that encourages the model to learn useful shared representations. We evaluate QueST across 145 robot manipulation tasks, and show that it outperforms several state-of- the-art baselines in multitask and few-shot learning settings. Our results highlight the usefulness of QueST's encoder (decoder) semaanically-sound, task-agnostic tokenizer (detokenizer) for continuous actions, and its potential to leverage Large Multi-modal Language Models in stage-2.
While the benchmarks consider encompass wide variety of tasks, the held-out tasks are still structurally similar to the pretraining set, which makes few-shot transfer feasible. In scenarios with more diverse task, current model may struggle to solve new tasks solely within the learned skill space. A promising direction is to train stage-1 on larger datasets, such as Open X- E-Ebbodiment [47], with an expanded codebook that could capture more diverse motion primitives. Additionally, our current architecture only accounts for causality. Future work should explore other inductive biases, like geometric invariance and dynamic consistency, enhance abstraction learning.
Limitations:
A statement on societal impact: This paper works towards the broader goal of automating wide range manipulation tasks. While this can have positive impacts, such as helping people with mobility impairments or performing menial tasks humans would rather not do, it can also have negative impacts such as automating peoples' jobs away and further concentrating wealth in the hands of handful of companies. It is important that we in the machine learning community advocate for equitable use of the technology we develop.
10

---

References
[1] Q-transformer: Scalable offline reinforcement learning via autoregressive q-functions. In 7th Annual Conference on Robot Learning, 2023.
[2] A. Agostinelli, T. Denk, Z. Borsos, J. Engel. M. Verzetti, A. Caillon, Q Huang. A. Jansen, A. Roberts, Tagliasacchi, M. Shariti, N. Zeghidour, and C. Frank. Musiclm: Generating music from text, 2023.
[3] M. Ahn, A. Brohan, N. Brown, Y. Chebotar, O. Cortes, B. David, C. Finn, C. Fu, K. Gopalakrishnan. K. Hausman, A. Herzog, D. Ho, Hsu, Ibarz, B. Ichter, A. Irpan, E. Jang, R. J. Ruano, K. Jeffrey, S. Jesmonth, N. J. Joshi, R. Julian, D. Kalashnikov, Y. Kuang. K.-H. Lee, S. Levine, Y. Lu, L. Luu, C. Parada, P. Pastor, J. Quiambao, K. Rao, J. Rettinghouse, D. Reyes, P. Sermanet, N. Sievers, C. Tan, A. Toshev, V. Vanhoucke, F. Xia, T. Xiao, P. Xu, S. Xu, M. Yan, and A. Zeng. Do as i can, not as i say: Grounding language in robotic affordances, 2022.
[4] M. Ahn, D. Dwibedi, C. Finn, M. G. Arenas, K. Gopalakrishnan, K. Hausman, B. Ichter, A. Irpan, N. Joshi, R. Julian, Kirmani, I. Leal, E. Loe, S. Levine, Y. Lu, I. Leal, S. Maddineni, K. Rao, Sadigh, P. Sanketi, P. Sermanet, Q. Vuong, Welker, F. Xia, Xiao, P. Xu, S. Xu, and Z. Xu, Autort: Embodied foundation models for large scale orchestration of robotic agents, 2024.
[5] A. Ajay, A. Kumar, P. Agrawal, Levine, and O. Nachum. Opal: Offline primitive discovery for accelerating offline reinforcement learning, 2021.
[6] A. Baevski, Zhou, A. Mohamed, and M. Auli, wav2vec 2.0: A framework for self-supervised learning of speech representations. CoRR, abs/2006.1 1477, 2020. URL https://arxiv org/ abtaa/bbb./0..//////////////// 11477.
[7] Y. Bai, X. Geng, K. Mangalam, A. Bar, A. Yuille, T. Darrell, J. Malik, and A. A. Efros. Sequential modeling cnables scalable learning for large vision models. arXiv preprint arXiv:331..0000 2023.
[8] H. Bao, L. Dong, S. Piao, and F. Wei. Beit: Bert pre-training of image transformers. arXiv preprint 10X.v22106.0855: 2021.
[9] Z. Borsos, R. Marinier, D. Vincent, E. Kharitonov, O. Pietquin, M. Sharifi, D. Roblek, O. Teboul, D. Grangier, M. Tagliasacchi, and N. Zeghidour. Audiolm: a language modeling approach to audio generation, 2023.
[10] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar, J. Dabis, C. Finn, K. Gopalakrishnan, K. Hausman, A. Herzog, J. Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint axXiv:2222.06817, 2022.
[11] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar, X. Chen, K. Choromanski, T. Ding, D. Driess, A. Dubey, C. Finn, al. Rt-2: Vision-language-action models transfer web knowledge to robotic control. arXiv preprint Xiv: 15818, 2023.
[12] A. Bucker, Figueredo, S. Haddadin, A. Kapoor, S. Ma, S. Vemprala, and Bonatti. Latte: Language trajectory transformer. 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 7287-7294. IEEE, 2023.
[13] H. Chang, H. Zhang, L. Jiang, C. Liu, and W. T. Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11315- 11325, 2022.
[14] I.. Chen, S. Bahl, and D. Pathak. Playfusion: Skill acquisition via diffusion from language-annotated play. In Conference on Robot Learning, pages 2012 2029990 PMI.R, 2023.
[15] C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song. Diffusion policy: Visuomotor policy learning via action difflusion. arXiv preprint urXiv:2303.041:1 2023.
[16] 7. J. Cui, Y. Wang, N. M. M. Shafiullah, and I.. Pinto. From play to policy: Conditional behavior generation from uncurated robot data, arXiv preprint arXiv:2210. 10047, 2022.
[17] P. Dhariwal, H. Jun, C. Payne, J. W. Kim, A. Radford, and I. Sutskever. Jukebox: A generative model for music, 2020.
[18] D. Driess, F. Xia, M. S. Sajjadi, C. Lynch, A. Chowdhery, B. Ichter, A. Wahid, J. Tompson, Q. Vuong, T. Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint arXhv:2303.0337 2023.
[19] Y. Duan, M. Andrychowicz. B. C. Stadie, J. Ho, Schneider, I. Sutskever, P. Abbeel, and W. Zaremba. One-shot imitation learning, 2017.
11

---

L. Y. Chen, K. Bllis, P. D. Fagan, J. Hejna, M. Itkina, M. Lepert, J. Ma, P. T. Miller, J. Wu, S. Belkhale, S. Dass, H. Ha, A. Jain, A. Lee, Y. Lee, M. Memmel, S. Park, I. Radosavovic, K. Wang, A. Zhan, K. Black, C. Chi, K. Hatch, S. Lin, J. La, J. Mercat, A. Rehman, P. R. Sanketi, A. Sharma, C. Simpson, Q. Vuong, H. R. Walke, B. Wulfe, T. Xiao, H. Yang, A. Yavary, Z. Zhao, C. Agia, R. Baijal, M. G. Castro. D. Chen, Q. Chen, T. Chung. J. Drake, E. P. Foster, J. Gao, D. A. Herrera, M. Heo, K. Hsu, J. Hu, D. Jackson, C. Le, Y. Li. K. Lin, R. Lin, Z. Ma, A. Maddukuri, S. Mirchandani, D. Morton, T. Nguyen, A. O'Neill, R. Scalise, D. Seale, V. Son, S. Tian, E. Tran, A. E. Wang, Y. Wu, A. Xie, J. Yang, P. Yin, Y. Zhang, Bastani, G. Berseth, Bohg, K. Goldberg, A. Gupta, A. Gupta, D. Jayaraman, J. Lim, [20] F. Ebert, Y. Yang, K. Schmeckpeper. B. Bucher. G. Georgakis, K. Daniilidis, C. Finn, and S. Levine. Bridge data: Boosting generalization of robotic skills with cross-domain datasets, 2021.
[21] P. Esser, R. Rombach, and B. Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pages 12873-12883, 2021.
[22] C. Finn, T. Yu, T. Zhang, P. Abbeel, and S. Levine. One-shot visual imitation learning via meta learning. 2017.
[23] P. Florence, C. Lynch, A. Zeng, O. A. Ramirez, A. Wahid, L. Downs, A. Wong, J. Lee, I. Mordatch, and J. Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168. PMLR, 2022.
[24] Z. Fu, T. Z. Zhao, and C. Finn. Mobile aloha: Learning bimanual mobile manipulation with low-cost whole-body teleoperation. arXiv preprint arXiv:2401.00 17, 2024.
[25] H. Ha, P. Florence, and Song. Scaling up and distilling down: Language-guided robot skill acquisition. In Proceedings of the 2023 Conference on Robot Learning, 2023.
[26] S. Haldar, Z. Peng, and L. Pinto. Baku: An efficient transformer for multi-task policy learning, 2024, URL https: / //arxiv...iiv////a
aaaaa/as//220/ 07539.
[27] K. He, X. Zhang, S. Ren, and Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pultern recognition, pages 770-778, 2016.
[28] M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.0999 2022.
[29] D. Jarrett, I. Bica, and M. van der Schaar. Strictly batch imitation learning by basey distribution matching. Advances in Neural Information Processing Systems, 33:7354 73654,73355770 2020.
[30] Z. Jiang, Y. Xu, N. Wagener, Y. Luo, M. Janner, E. Grefenstette, T. Rocktäschel, and Y. Tian. H-gap: Humanoid control with generalist planner. arXiv preprint 2311:22112.02682 2023.
[31] A. Khazatsky, K. Pertsch, S. Nair, A. Balakrishna, S. Dasari, S. Karamcheti, S. Nasiriany, M. K. Srirama,
J. Malik, R. Martín-Martin, S. Ramamoorthy, D. Sadigh, S. Song, J. Wu, M. C. Yip. Y. Zhu, T. Kollar, S. Levine, and C. Finn. Droid: A large-scale in-the-wild robot manipulation dataset. 2024.
[32] D. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6111 2013.
[33] A. Kirillov, F. Mintun, N. Ravi, H. Mao, C. Rolland, I., Gustafson, T. Xiao, S. Whitchead, A. C. Berg. W.-Y. Lo, al. Segment anything. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 4015-4026, 2023.
[34] D. Lee, C. Kim. S. Kim, M. Cho, and W.SS. Han. Autoregressive image generation using residual quantization, 2022.
[35] Lee, Y. Wang, H. Etukuru, H. J. Kim, N. M. M. Shafiullah, and L. Pinto. Behavior generation with latent actions. arXiv preprint wxXiv:2403.03/8// 2024.
[36] S. Levine, C. Finn, T. Darrell, and P. Abbeel. End-to-end training of deep visuomotor policies, 2016.
[37] Z.. Liang, Y. Mu, H. Ma, M. Tomizuka, M. Ding, and P. Luo. Skilldiffuser: Interpretable hierarchical planning via skill abstractions in diffusion-based task execution. arXiv preprint arXiv:2312./15553 2023.
[38] B. Liu, Y. Zhu, C. Gao, Y. Feng, Q Liu, Y. Zhu, and P. Stone. Libero: Benchmarking knowledge transfer for lifelong robot learning, Advances in Neural Information Processing Systems, 36, 2024.
[39] H. Liu, C. Li, Q. Wu, and Y. J. Lee. Visual instruction tuning. 2023.
[40] J. Luo, P. Dong, J. Wu. A. Kumar, X. Geng, and S. Levine. Action quantiee offline reinforcement learning for robotic skill learning. In Conference Robot Learning, pages PMLR, 2023.
12

---

[41] C. Lynch, M. Khansari, T. Xiao, V. Kumar. J. Tompson, S. Levine, and P. Sermanet. Learning latent plans from play. Conference on Robot Learning (CoRL), 2019. URL https:://aaaiiiii http://///aaiii..rr///////////// wtwg/a///////// 01973.
[42] Z. Mandi, F. Liu. K. Lee, and P. Abbeel. Towards more generalizable one-shot visual imitation learning. In 2022 International Conference on Robotics and Automation (ICRA), pages 2434 2444, IEEE, 2022.
[43] A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang, R. Kulkarni, L. Fei-Fei, S. Savarese. Y. Zhu, and
R. Martin-Martin What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108.0 03298, 2021.
[44] E. Mansimov and K. Cho. Simple nearest neighbor policy method for continuous control tasks, 2018. URL https://openrevierv net/forum? 1d-ByL48CC AW.
[45] F. Mentzer, D. Minnen, E. Agustsson, and M. Tschannen. Finite scalar quantization: Vq-vae made simple. arXiv preprint (xXi::3309.1500 2023.
[46] Octo Model Team, D. Ghosh, H. Walke, K. Pertsch, K. Black, O, Mees, S. Dasari, J. Hejna, C. Xu, J. Luo, T. Kreiman, Y. Tan, Sanketi, Q. Vuong, T. Xiao, D. Sadigh, C. Finn, and S. Levine. Octo: An open-source generalist robot policy. In Proceedings of Roboticx; Science and Systems, Delft, Netherlands, 2024.
[47] A. Padalkar, A. Pooley, A. Jain, Bewley, A. Herzog, Irpan, A. Khazatsky, A. Rai, A. Singh, A. Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:2310008864 2023.
[48] N. D. Palo and F. Johns. Dinobot: Robot manipulation via retrieval and alignment with vision foundation models. In IEEE International Conference on Robotics and Automation (ICRA), 2024.
[49] E. Perez, F. Strub, H. De Vries, V. Dumoulin, and A. Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[50] K. Pertsch, Y. Lee, and J. Lim. Accelerating reinforcement learning with learned skill priors. Conference on robot learning, pages 188-204. PMLR, 2021.
[51] D. A. Pomerleau, Alvinn: An autonomous land vehicle in a neural network. Advances in neural information processing systems, 1, 1988.
[ A. Radford, K. Narasimhan, Salimans, I. Sutskever, al. Improving language understanding by generative pre-training. 2018.
[53] A. Radford, Wu, R. Child, Luan, Amodei, and Sutskever. Language models are unsupervised multitask learners. 2019.
[54] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748 8763, PMLR, 2021.
[55] S. Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky,
J. Kay, J. T. Springenberg, T. Eccles, J. Bruce, A. Razavi, A. Edwards, N. Heess, Y. Chen, R. Hadsell, O. Vinyals, M. Bordbar, and N. de Freitas. A generalist agent, 2022.
[56] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[57] N. M. Shafiullah, Z. Cui, A. A. Altanzaya, and L. Pinto. Behavior transformers: Cloning & modes with one stone. Advances in neural information processing systems, 35:22955-22968, 2022.
[58] L. X. Shi, J. J. Lim, and Y. Lec. Skill-based model-based reinforcement learning. arXiv preprint arXiv:2207.07560, 2022.
[59] A. Singh, H. Liu, G. 7hou, A. Yu, N. Rhinchart, and S. Levine. Parrot: Data-driven hehavioral priors for reinforcement learning. In International Conference on Learning Representations, 2021. URL https:// wwtpereewweeeeeeeeeeeeeeeeeeeeee wet.//orum/kdd//uu--ooo
[60] K. Sohn, H. Lee, and X. Yan. Learning structured output representation using deep condi- tional generative models.
C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Asso- ciates, Inc., 2015. URL https::/procceeeeee wwwrroeeeiinngnnuuipp....p 2a55a44466ba5507777955000666 wttp:///2200000000PPrarpp..
13

---

[61] G. Swamy, S. Choudhury, J. A. Bagnell, and Z. S. Wu. Causal imitation learning under temporally correlated noise, 2022.
[62] G. Team. Gemini: A family of highly capable multimodal models, 2024.
[63] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava,
S. Bhosale. D. Bikel. L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu,
B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, Korenev,
Koura, M.-A. Lachaux, Lavril, Lee, Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, 1. Molybog. Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, Xu, Z. Yan. L. Zarov, Y. Zhang. A. Fan, M. Kambadur, S. Narang. A. Rodriguez, R. Stojnic,
Khabsa, I. Kloumann,
S. Edunov, and T. Scialom. Iaama 2: Open foundation and fine-tuned chat models, 2023.
[64] A. van den Oord, S. Dieleman, H. Zen, K. Simonyan, O. Vinyals, A. Graves, N. Kalchbrenner, A. Senior, and K. Kavukeuoglu. Wavenet: generative model for raw audio. Arxiv, 2016. URL https:
/ //arxiv... .ar.vvorrg org/ htrg///bbe6600.... 110...69.... 033999.
[65] A. Van Den Oord, O. Vinyals, al. Neural discrete representation learning. Advances neural information processing systems, 30, 2017.
[66] Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need, Advances in neural information processing systems, 30, 2017.
[67] W. Wan, Zhu, R. Shah, and Y. Zhu. Lotus: Continual imitation learning for robot manipulation through unsupervised skill discovery. arXiv preprint arXiv:231 1,00255, 2023.
[68] Z. Xian, N. Gkanatsios, T. Gervet, T.-W. Ke, K. Fragkiadaki. Chaineddiffuser: Unifying trajectory diffusion and keypose prediction for rohotic manipulation. In 7th Annual Conference on Robot Learning, 2023.
[69] J. Yu, X. J. Y. Koh, H. Zhang, R. Pang. J. Qin, A. Ku, Y. Xu, J. Baldridge, and Y. Wu. Voctor-quantized image modeling with improved VQGAN. In International Conference on Learning Representations, 2022. URL https: / //pprreveee/eeeeee///e/eee
[70] T. Yu, D. Quillen, Z. Hc, R. Julian, K. Hausman, C. Finn, and S. Levine. Mcta-world: A benchmark and evaluation for multi-task and meta reinforeement leamnnng. In Conference on robot learning, pages 1094-1100. PMLR, 2020.
[71] Y. Ze, G. Zhang, K. Zhang, C. Hu, M. Wang, and H. Xu. 3d diffusion policy. arXiv preprint axXiv:240..03554 2024.
[72] S. Zhang, L. Dong, X. Li. S. Zhang, X. Sun, S. Wang. J. Li, R. Hu, T. Zhang, F. Wu, al. Instruction tuning for large language models: A survey. arXiv preprint arXiv:2308. 10792, 2023.
[73] T. Z. Zhao, V. Kumar, S. Levine, and C. Finn. Learning fine-grained bimanual manipulation with low-cost hardware, 2023.
[74] R. Zheng, C.-A. Cheng, H. Daumé III, F. Huang, and A. Kolobov. Prise: Learning temporal action abstractions sequence compression problem. arXiv preprint arXiv:2402. 10450, 2024.
[75] R. Zheng, X. Wang, Y. Sun, S. Ma. J. Zhao, H. Xu, H. Daumé III, and F. Huang. Taco: Temporal latent cettion-driven contrastive loss for visual reinforeement learning. Advances Neural Information Processing Systems, 36, 2024.
[76] zhengyao jiang, Zhang, M. Janner, Y. Li, Rocktäschel, E. Grefenstette, and Y. Tian. Efficient planning in compact latent action space. In 3rd Offline RL Workshop: Offline RL. as a "Launchpad" 2022. URL https: //ppnreviee net/forum? met/forumiid--VVBS... -VBETTSTaa
[77] X. Zou, J. Yang, H. Zhang, F. Li, L. Li, J. Wang, L.. Wang, J. Gao, and J. Lee. Segment everything everywhere all at once. Advances in Neural Information Processing Systems, 36, 2024.
14

---

Appendix
A
Website
For further results and videos please website. git.ub...tt........ io
B Experiment Details
B.1
Hyperparameters:
We present hyperparameters in the following tables:
Table 3: Stage Parameters
Table 4: Stage Parameters
Parameter
Value
Parameter
Value
encoder dim
256
vocab size
1000
decoder dim
256
block size (n)
8
sequence length (T)
16/32
number of layers
6
encoder heads
4
number of heads
6
encoder layers
2
embedding dimension
384
decoder heads
4
attention dropout beam size temperature
0.1
decoder layers
4
5
attention dropout
0.1
1.0
fsq level conv layers downsampling factor
[8, 5, 5]
decoder loss scale
0/10
3
execution horizon (T,)
8
2/4
observation history
1
B.2 Architecture Implementation:
For vision encoder we used shallow Convolutional Neural Network (CNN), consisting of the first four layers of ResNet18 [27] followed by spatial softmax [36]. In encoder, we use causal convolution layers from [30]. For transformer blocks, we used the transformers library from hugging face https: hatees//hgggggaaac co/ docs/ tetecs//rasffoomrrsssse with appropriate masking for ensuring causality.
B.3 Baseline Implementation:
To ensure fair comparison of different model architectures, we use same input modalities and same observation & task encoders for all basclines, for both LIBERO and Metaworld experiments. VQ-BeT needs goal image, we instead give it task embedding as goal. Same as QueST, we concatenate observation embeddings for all modalities at any timestep and project them to respective model's hidden dimension.
Depending on the dataset, also tune some key hyperparameters for the baselines and present the results for best performing ones.
1. ResNet-T: Transformer trunk's 's hidden dimension and number of layers determines the model capacity. Original implementation [38] uses the hidden dimension of 64 with 4 layers. We observed improved performance for the hidden dimension of 256 with 6 layers and hence report all results for that. As per original implementation history of 10 timesteps.
observation
2. Diffusion Policy: The model capacity is determined by hidden dimension U-Net layers. Most widely used implementations use [256, 512, 1024], we ablate larger model with (256, 256, 512, 512, 1024] but did not observe any performance gains. We also ablate prediction (7) and execution horizon (T.) with 16. 32 and 8, 16 respectively and observed best performance for = 32, Ta 16 on LIBERO and 16, T,
MetaWorkd As per original paper ablations [15] an observation history of 1 was used.
15

---

3. VQ-BeT: Since LIBERO and MetaWorld are larger datasets as compared to the benchmarks in original VQ-BeT paper, we ablate some parameters to increase the model capacity. Specifically, the stage encoder by default is single MIP layer of dimension 128. We ablate this with 2, 4 layers and with 256, 512 dimensions but observed worse reconstruction loss with increase in capacity. We use residual-VO configuration of 32/2 & 1024 sized codebook which is close to the codebook size of 1000 for QueST. We use an observation window size of 10 and ablate the action window size (T) with 1, 1,5, 32. On LIBERO, the performance was lowest for T = 1, and highest for T = 5. VQ-BeT maps the whole input sequence to just one embedding leading to extreme compression for larger sequence length and thus performs worse with T 32.
B.4 Compute:
The models are implemented in PyTorch. For all our experiments use server consisting of 8 Nvidia RTX 1080Ti 10GB memory each. And all our models easily fit on one GPU for training.
B.5
Discussion on Baselines
Many recent works discrete latent variable models as mechanism to learn shared abstractions over continuous low-level skills. QueST, VQ-Be'T and PRISE all do this and perform two- staged learning. However, in this work we propose several important architecture choices leading 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-BeTT performs temporal abstraction by encoding a sequence of actions into one state-independent latent vector. It concatenates input actions and encodes the sequence to just one single latent encoding using MLP which does not explicitly model any temporal correlation. A sampled action sequence might contain multiple motion primitives of variable length and start point, and capturing them with single encoding is limiting it restricts abstraction at different levels of granularity. This is validated by the poorer (-14%) performance of VQ-BeT with a larger chunk size of 32. QueST flexibly captures this variability within encodings using its encoder, specifically designed to model temporal correlation within input actions. With an effective codebook size of C, QueST's effective latent space is C" while that of VQ-BeT is C. Its 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 and reliance 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 effectively encode new action sub-sequences that the model has not seen before (eg. stitching sequences in between two tasks in LIBERO-LONGE This is primarily why decoder finetuning is necessary in PRISE. On the contrary, QueST is more end-to-end as it lets the encoder handle temporal abstraction in the encoding phase itself and gracefully encodes new action-sequences combinations of its learned latent codes. Our few-shot results without decoder-finetuning finetuning shows that such a unified approach learns more generalized abstractions than the baselines and can more effectively represent new action sequences in unseen tasks.
suboptimal character-level tokenization, and it might struggle
Despite tuning the hyperparameters, the multitask performance of ACT very low (54%) on LIBERO. Moreover, as per BAKU [26], ACT's performance drastically reduces when number of demos are reduced to 35 task in MetaWorld. This suggests ACT is very data-hungry and hence struggles in low-data regime. Another major issue with ACT is tuning the action chunk size to which the method is very sensitive (-14% in metaworld fewshot with chunk size of 32). While report QueST results for chunk size 32, we did observe very little variation (<1% in LIBERO) with size 16,48 and 64, indicating robustness to this hyperparameter.
C Discussion on Ablations
For aiding this discussion we present the ablation results again in table 5 and table 6 below.
16

---

VQ
Obs. Cond.
Mirror Dec.
Ours
LIBERO-90
81.2 0.6
81.9 + 1.1
86.3 + 0.9
88.6 = 0.4
Few Shot
62.5 + 2.0
61.3 2.2
45.4 2.0
68.8 = 1.7
Table 5: Success rates after ablating design details of QueST.
Replacing FSQ with VQ still outperforms VQ-BeT few-shot setting suggesting that QueST's superior performance is not only due to better quantization scheme but also due to 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 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 closely related prior work SPiRL[58] that tried 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 s cross-attention architecture that allows for attending to all z-tokens and maintaining causality at the same time.
Non Causal do
Non Causal tho
Fully Non Causal
Ours
LIBERO-90 Few Shot
82.0 + 1.6
85.1 + 1.8
78.5 + 0.5
88.6 + 0.4
58.8 3.0
61.6 + 2.5
56.1 + 1.8
68.8 1.7
Table 6: Success rates after ablating the causality in QueST.
We observe that a fully-causal stage-1 most optimal and a non-causal decoder does not hurt much non-causal encoder does. This be explained with simplistic setting where the input to stage-1 are 2D trajectories of point agent. Consider an anti-clockwis circular trajectory and S-shaped one where the first half of the later overlaps with the first half (semi-circle) of the former. When both of these trajectory sequences are inputted the stage-1, a non-causa encoder will assign distinct sequences of z-tokens for both trajectories. But a causal encoder will assign same sequence of z-tokens for the first half of both trajectories and distinct later parts. This allows the model to re-use the z-tokens corresponding to semi-circle for creating other trajecttraeessriess that has semi-circle in them for example C-shaped or infinity-shaper trajectories.
Frozen the
Finetuned the
loss scale 10
loss scale 100
Few Shot 66.0 3.6 68.8
1.7
66.0 + 1.0
Table 7: Success rates for decoder finetuning settings in few-shot II.
Table 7 illustrates the impact of decoder finetuning in LIBERO-LONG fewshot IL setting. QueST outperforms all baselines even without finetuning the decoder. Finetuning decoder should not be necessary in this setting, as LIBERO-LONG tasks are combination of two tasks from LIBERO-90 (pretraining set). This highlights QueST's effectiveness in stitching trajectories using its learned skill-space. We report the finetuning results in the main paper, as they exhibit better performance.
D Additional Results
D.1 Fewshot IL
Fewshot Evaluation Protocol: In finetuning phase, we finetune ResNer-T, VQ-BeT & QueST for
100 epochs and ACT & Difffusion Policy for 200 epochs. For each task in MetaWorld, we evaluate
17

---

each method across 10 evenly spaced checkpoints for 5 seeds on 50 distinct initial states and report the results corresponding the best performing checkpoint. For Libero, found the final checkpoint to perform best for all methods and hence report results corresponding to it across 9 seeds.
Table 8: LIBERO 5-shot IL success rates across unseen 10 tasks, Results across random seeds.
Task ID
ResNet-T
ACT
Diffusion Policy
PRISE
VQ-ReT
QueST
1
53.8 = 116
20.0 6.0
32.0 = 5.5
26.7 = 6,4
16.5 + 11.9
56.4 + 5.5
2
65.7 16.9
33.3 13.1
37.8 5.9
48.3
3
70.2 1 7.8
67.7 1: 6.2
76.4 1 7.2
70.0 - 0.0
52.7 4 744
66.7 6.7
4
75.8 7.6
70.3 6.2
98.2 - 1.7
78.3 =
45.3 7.7
88.0
5
26.7 11.9
35.0
44.7 = 6.0
45.0 10.8
30.3 + 13.0
42.0 = 7.1
86.9
6.5
37.1 3.7
90.0
48.7 17.2
92.4 4.4
7
8.44
15.0
14.9 6.7
25.0 10 4.1
45.5
58.2 6.1
23.7 + 12.1
26.7 7.0
6.2
45.0
9.3
47.1 + 8.6
9
55.0 = 4.0
7.0
46.6 - 6.2
68.3 - 6.2
15.00:00
65.0 12.2
Table MetaWorld 5-shot success rates across 5 unseen tasks. Results across random seeds.
Task ID
ResNet-T
ACT
Diffusion Policy
PRISE
VQ-BeT
QueST
box-close-v2
63.2 - 5.2
67.2 + 6.2
68.0 = 1.6
60.8 + 6.6
76.3 9.6
84.0 = 7.3
disassemble-v2
68.8 2.0
83.2 3.2
81.3 3.8
74.1 7.3
92.7 1.9
76.4 26.0
hand-insert-v2
37.2 4.1
53.2 3.7 3.7
39.3 + 1.9
60.0 1 5.0
48.0 6.5
49.6 19.6.:00 6.4
pick-place-wall-v2
74.4 4.7
70.7 5.2
71.7 5.7
65.3 1.9
76.8 11.4
stick-pull-v2
8.8
7
71.31.9
67.5 5.6
62.0 = 11.4
72.8
D.2 Multitask II.
Table 10: LIBERO-90 multitask IL success rates across 90 tasks. Results across random seeds.
Task ID
ResNet-T
ACT
Diffusion Policy
PRISE
VQ-BeT
QueST
1.00
0.90
0.99
0.80
1.00
1.00
0.96
0.30
0.98
0.35
0.94
0.97
0.96
0.50
0.99
0.70
0.97
0.91
0.74
0.22
0.91
0.50
0.99
0.94
0.95
0.58
0.93
0.45
0.95
0.98
0.91
0.39
0.99
0.65
0.98
0.98
7
0.95
0.29
0.94
0.50
0.86
0.93
8
0.96
0.72
0.90
0.95
0.80
0.99
9
0.74
0.41
0.93
0.60
0.75
0.93
10
0.97
0.65
0.91
0.35
0.83
0.90
11
0.97
0.82
0.98
0.95
0.96
0.97
12
0.90
0.73
0.94
0.95
0.80
0.94
13
0.82
0.62
0.81
0.20
0.87
0.76
14
0.86
0.72
0.94
0.40
0.49
0.71
18

---

Task ID
ResNet-T
ACT
Diffusion Policy
PRISE
VQ-BeT
QueST
15
0.87
0.49
0.95
0.35
0.46
0.58
16
0.97
0.86
0.99
0.75
0.98
0.96
17
0.72
0.40
0.89
0.40
0.53
0.80
18
0.79
0.20
0.76
0.15
0.80
0.67
19
0.93
0.75
0.99
0.30
0.91
1.00
20
0.87
0.41
0.98
0.65
0.68
0.92
21
1.00
0.82
0.99
1.00
0.96
1.00
22
0.90
0.44
0.97
0.30
0.91
0.93
23
0.97
0.75
0.99
0.85
0.95
0.92
24
0.75
0.11
0.85
0.05
0.69
0.82
25
0.97
0.44
0.99
0.95
0.94
1.00
26
0.97
0.85
1.00
0.90
0.85
0.99
27
0.72
0.14
0.88
0.55
0.50
0.52
28
0.72
0.20
0.86
0.05
0.45
0.68
29
1.00
0.68
1.00
1.00
0.97
1.00
30
1.00
0.19
1.00
1.00
0.92
0.97
31
0.91
0.83
0.96
0.50
0.85
0.90
32
0.99
0.90
1.00
0.85
0.88
0.99
33
0.57
0.20
0.58
0.20
0.37
0.67
34
0.85
0.56
0.84
0.30
0.87
0.98
35
0.93
0.52
0.97
0.80
0.98
0.92
36
0.97
0.67
0.99
0.75
0.98
0.97
37
0.85
0.24
0.97
0.25
0.73
0.74
38
0.78
0.41
0.91
0.30
0.90
0.62
39
0.86
0.32
0.90
0.20
0.90
0.88
40
0.96
0.35
0.98
0.85
0.90
0.93
41
0.90
0.27
0.79
0.50
0.91
0.92
42
1.00
0.74
1.00
0.55
0.89
1.00
43
0.98
0.41
0.99
0.80
0.97
0.98
44
0.80
0.39
0.89
0.40
0.83
0.93
45
0.99
0.83
1.00
0.85
0.99
0.98
46
0.97
0.60
1.00
0.55
0.91
0.99
47
0.75
0.37
0.31
0.35
0.65
0.91
48
0.87
0.27
0.53
0.25
0.88
0.98
49
0.90
0.55
0.96
0.65
0.48
0.95
50
0.88
0.54
0.82
0.65
0.60
0.99
51
0.80
0.33
0.28
0.40
0.87
0.88
52
0.79
0.28
0.00
0.10
0.91
0.75
53
0.74
0.33
0.34
0.30
0.84
0.82
54
0.88
0.64
0.73
0.60
0.79
0.87
55
0.83
0.51
0.77
0.50
0.95
0.93
19

---

Task ID
ResNet-T
ACT
Diffusion Policy
PRISE
VQ-BeT
QueST
56
0.85
0.62
0.49
0.35
0.92
0.83
57
0.99
0.64
1.00
0.80
1.00
0.97
58
0.95
0.57
1.00
0.50
1.00
0.99
59
0.84
0.56
0.78
0.20
0.98
0.95
60
0.94
0.68
0.89
0.65
0.91
1.00
61
0.91
0.95
0.90
0.80
0.98
1.00
62
0.96
0.75
0.58
0.85
0.99
0.81
63
0.70
0.43
0.38
0.40
0.84
0.78
64
0.73
0.04
0.41
0.40
0.38
0.78
65
0.73
0.16
0.75
0.15
0.68
0.85
66
0.76
0.45
0.65
0.15
0.84
0.79
67
0.84
0.72
0.66
0.30
0.87
0.93
68
0.78
0.73
0.44
0.55
0.74
0.85
69
0.83
0.68
0.59
0.85
0.90
0.93
70
0.88
0.56
0.57
0.90
0.93
0.89
71
0.90
0.52
0.92
0.55
0.97
0.92
72
0.85
0.52
0.98
0.35
0.85
0.94
73
0.89
0.59
0.86
0.60
0.84
0.98
74
0.72
0.18
0.61
0.30
0.33
0.70
75
0.77
0.45
0.38
0.45
0.95
0.95
76
0.64
0.22
0.21
0.25
0.30
0.61
77
0.89
0.70
0.35
0.65
0.70
0.89
78
0.57
0.46
0.14
0.80
0.85
0.97
79
0.63
0.28
0.06
0.45
0.68
0.86
80
0.73
0.59
0.01
0.30
0.87
0.98
81
0.65
0.53
0.08
0.30
0.44
0.70
82
0.63
0.24
0.54
0.35
0.61
0.70
83
0.80
0.56
0.49
0.80
0.89
0.94
84
0.55
0.35
0.47
0.55
0.43
0.75
85
0.70
0.74
0.79
0.75
0.93
0.92
86
0.69
0.53
0.13
0.75
0.47
0.89
87
0.84
0.65
0.98
0.95
0.86
0.92
88
0.82
0.54
0.96
0.65
0.87
0.97
89
0.91
0.77
0.70
0.55
0.96
0.97
90
0.80
0.29
0.91
0.85
0.89
0.56
Table 11:
MetaWorka multitask II. success rates across 45 tasks. Results across 5 random seeds.
Task ID
ResNet-T
ACT
Diffusion Policy
VQBeT
QueST
assembly-v2 basketball-v2
0.73
0.97
0.88
0.82
1.00
0.76
0.80
0.78
0.82
0.68
20

---

Task ID
ResNet-T
ACT
Diffusion Policy
VQBeT
QueST
bin-picking-v2
0.89
1.00
0.96
0.20
0.94
button-prees-ttoddoww-22
1.00
1.00
1.00
1.00
1.00
button-pres-ttoodoon-wwlllll
1.00
1.00
1.00
1.00
1.00
button-press--2
1.00
1.00
1.00
1.00
1.00
button-preess-alll-y
1.00
1.00
0.98
0.98
0.98
coffee-button--2
1.00
1.00
1.00
1.00
1.00
coffee-pull-v?
0.90
0.92
0.96
0.82
0.98
coffee-push-v2
0.89
0.96
0.86
0.94
0.90
dial-turn-v2
0.98
0.99
1.00
1.00
1.00
door-close-v2
1.00
1.00
1.00
1.00
1.00
door-lock-v2
1.00
0.99
1.00
1.00
1.00
door-open-v2
0.96
0.95
0.96
0.94
0.94
door-unlock-v2
1.00
1.00
1.00
1.00
1.00
drawer-close-v2
1.00
1.00
1.00
1.00
1.00
drawer-open-v2
1.00
1.00
1.00
1.00
1.00
faucet-close-v2
1.00
1.00
1.00
1.00
1.00
faucet-open-v2
1.00
1.00
1.00
1.00
1.00
hammer-v2
0.95
1.00
0.98
1.00
0.94
handle-press-side-VV
1.00
1.00
1.00
1.00
1.00
handle-press-v2
1.00
1.00
1.00
1.00
1.00
handle-pull-side-22
0.69
0.94
0.78
0.74
0.98
handle-pull-v2
1.00
1.00
1.00
1.00
1.00
lever-pull-v2
0.94
0.93
0.84
0.80
0.92
peg-insert-side-t
0.81
0.94
0.90
0.76
0.86
peg-unplug-side-v2
0.88
0.91
0.88
0.92
0.90
pick-
0.62
0.89
0.74
0.34
0.76
pick-place-v2
0.67
0.71
0.76
0.74
0.78
plate-slide-back-side-2
1.00
1.00
1.00
1.00
1.00
plate-slide-back-22
1.00
1.00
1.00
1.00
1.00
plate-slide-side-22
0.98
1.00
0.98
0.98
1.00
slate-slide-22
1.00
1.00
1.00
1.00
1.00
push-back-v2
0.72
0.64
0.76
0.64
0.80
push-v2
0.84
0.90
0.84
0.76
0.92
push-wall-v2
0.92
0.98
0.94
0.94
1.00
reach-v2
0.39
0.37
0.32
0.28
0.36
reach-wall-v2
0.49
0.47
0.52
0.36
0.42
shelf-place-v2
0.65
0.85
0.66
0.76
0.88
soccer-v2
0.42
0.25
0.42
0.36
0.52
stick-push-v2
0.75
1.00
0.96
0.94
0.96
sweep-into-v2
0.90
0.92
0.88
0.90
0.84
sweep-v2
0.98
1.00
0.98
1.00
1.00
21

---

Task ID
ResNet-T
ACT
Diffusion Policy
VQBeT
QueST
window-close-v2
1.00
1.00
1.00
1.00
1.00
window-open-v2
1.00
1.00
1.00
1.00
1.00
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The main claims made in the abstract and introduction are that we achieve stronger performance than several state of the art baselines. We show several results justifying this claim in Figures 3 and 2.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We stated our limitation in the last section of the main paper.
Guidelines:
The answer NA means that the paper has 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 (c.g.. independence assumptions, noiseless settings, model sell-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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
23

---

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

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We include codes in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. wwtppppwpppppaecccccccccccccccee http//udddes/cocee
swimisssicicccy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification:
We have the training details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the averaged metrics different datasets, and the difference is significant.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
25

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
All the experiments can be run on one A6000 GPU.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that NeurIPS Code of Ethics https: / /ttp///eerrpp.eeeee//ee wccppublcc//eee hptpbblit/Ethiccccdddeleeees
didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification: This paper conforms with the ethics code.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration 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:
We discuss the broader impact in the main paper.
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.
26

---

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

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification: No new datasets/models proposed in this paper.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl 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.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer:
answerNA
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "MKzAUa41Bi", "reviewer": "Reviewer_cTdz", "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.", "questions": "- How many seeds were used in Figure 4 experiments? \n- 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?", "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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> 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.\n\nFor 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. \n\nFor 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. \n\nFor 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.\n\nIn 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.\n\n> The latency concerns could be relevant if their system is computationally slow.\n\nOur 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.\n\n> It would have been good to see results on a real robotic system instead of just in simulation.\n\nWe 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.\n\n> Regarding lack of error bars in Figure 4. How many seeds were used in Figure 4 experiments?\n\nUnfortunately 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.\n\n> 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?\n\nNot 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).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VoqeaaSxz0", "reviewer": "Reviewer_jrVG", "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.", "questions": "**Questions**\n- Equation 1, perhaps write out what round_std is.\n- 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?\n- 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?\n- Figure 2: Any intuition why few-shot is better than multi-task for QueST?\n\n**Possible typos**\n- Page 8, line 305: viz. means visualized?", "limitations": "- The current application is only on robotic manipulation but it will be interesting to extend this to other domains", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n- The paper is mostly well written", "weaknesses": "I am happy to increase my score if these points are addressed.\n\n**Comments**\n- 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.\n- 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).\n- It is not clear to me how this approach handles multimodality (in the sense of action distributions?)\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> 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? \n\nThanks 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.\n\n> 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\n\nWhile 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.\n\nRegarding 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.\n\n> It is not clear to me how this approach handles multimodality (in the sense of action distributions?)\n\nFirst, 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].\n\n> Ablation on \"causality\": I feel the word \"causality\" is misused... Is this not still modeling correlation only? How is this causal exactly?\n\nYou 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.\n\n> Equation 1, perhaps write out what round_std is.\n\nIt's the nearest integer rounding with straight-through gradients. We’ll update the manuscript.\n\n> 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?\n\nWe 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.\n\n> 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?\n\nYes, 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. \n\nThe 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.\n\n> Figure 2: Any intuition why few-shot is better than multi-task for QueST?\n\nThe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OTlOBJnL0Y", "reviewer": "Reviewer_VaLR", "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.\n- 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.\n- 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.", "questions": "- Are the updated results on Meta-World going to replace the current results in the main body of the paper?\n- Where can the “provided scripted policies” for collecting demonstrations on Meta-World as mentioned on the paper website be found?\n- 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.\n- How were the hyperparameters tuned for the proposed method, for the ablations, and for baselines?\n- How does the proposed method compare to the baselines in terms of latency?\n- Is the code going to be released publicly?\n\nUpdates:\n- Presentation: 2 -> 3\n- Score: 4 -> 6", "limitations": "- It would be nice to have a discussion of latency considerations.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Originality:\n- This work is a nice combination of existing methods (latent action abstraction, causal encoders/decoders, and Finite Scalar Quantisation).\n- The use of causal encoders/decoders for latent abstraction is an interesting idea.\n\nQuality:\n- LIBERO and Meta-World are reasonable benchmarks for validating the proposed method.\n- The method shows good performance when compared against several recent baselines: diffusion policy, ACT, VQ-BeT, and PRISE.\n- 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.\n- Results are reported for three random seeds for the experiments on LIBERO.\n- 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.\n\nClarity:\n- The paper is generally well-written.\n- The videos on the project website provide a nice indication of qualitative behaviour.\n\nSignificance:\n- 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:\n- 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.\n\nQuality:\n- There are no results for PRISE on Multitask LIBERO-LONG.\n- 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.\n- 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.\n- There is no analysis of inference latency of the proposed method and the baselines, which is important for practical applications.\n- 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.\n\nClarity:\n- 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.\n- 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.\n- It is not specified where the “provided scripted policies” for collecting demonstrations on Meta-World as mentioned on the paper website can be found.\n- Minor comment: Figure 3b is not described or mentioned in the text.\n\nSignificance:\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> 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\n\nMany 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. \n\nLike 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).\n\nPRISE 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.\n\n> There are no results for PRISE on Multitask LIBERO-LONG.\n\nWe were unable to reproduce the PRISE results when we ran their code so we were only able to include results presented in their paper.\n\n> The paper website states: “Depending on the dataset, we also tune some key hyperparameters for the baselines”. This is vague…\n\nBelow we summarize information about hyperparameter tuning we’ve added to the appendix\n- 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.\n- 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.\n- 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).\n- ACT: We tune model hidden dimension (256, 512); chunk size (16, 32); and kl weight (10,100).\n\n> 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.\n\nWe 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.\n\n> 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?\n\nOur 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.\n\n> It might be interesting to have an ablation experiment that uses Residual VQ as used by VQ-BeT instead of FSQ…\n\nWe 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.\n\n> 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.\n\nWe’ve moved this detail to the appendix and will add a line confirming that this is the case for all experiments.\n\n> It is not specified where the “provided scripted policies”…can be found.\n\nThey can be found in the official Metaworld codebase.\n\n> Figure 3b is not described or mentioned in the text.\n\nFixed\n\n> Are the updated results on Meta-World going to replace the current results in…the paper?\n\nYes\n\n> Is the code going to be released publicly?\n\nThis will happen alongside the camera ready version of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z5HCkxyoqQ", "reviewer": "Reviewer_utkR", "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.", "questions": "- Why the causality of the encoder transformer is so critical?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The overall paper is easy to read, with the description of the methodology being well-written.\n\n- 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. \n\n- A lot of ablation studies are conducted. However, the related analysis is not sufficient.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> 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.\n\nMany 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. \n\nLike 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).\n\nPRISE 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.\n\n> A lot of ablation studies are conducted. However, the related analysis is not sufficient.\n\nBelow we mention detailed analysis which has also been added to the paper.\n- 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. \n- 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. \n- 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.\n- We discuss the causality ablation in response to the next question.\n- 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.\n\n> Why is the causality of the encoder transformer so critical?\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We are grateful for the insightful feedback from all reviewers. The reviewers have recognized the novelty of our approach in modeling the inherent structure of manipulation action data through temporal correlation and causal-masking. A particularly noteworthy aspect of our research is the learning and decision-making framework based on state-independent abstractions, which has garnered significant interest. Our comprehensive evaluation, spanning 145 manipulation tasks across 3 diverse benchmarks and 3 distinct settings, against 5 established baselines demonstrates the robustness, versatility and superiority of our method. We are encouraged by the reviewers' assessment that QueST's exceptional performance and simplicity make it a promising candidate for wider adoption within the community.

Below we summarize some common concerns:
1. **Innovation/unique contribution**: QueST introduces a novel autoencoder architecture that flexibly captures diverse, variable-length motion primitives by representing them using a sequence of discrete codebook entries (skill tokens). It does so by modeling temporal correlation in action sequence data resulting in shareable and transferable abstractions, leading to superior multitask and fewshot performance.
1. **Latency**: 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.
2. **Latent space analysis**: On our project website, we provide t-SNE visualizations of the skill token embeddings. These visualizations demonstrate clear alignment of similar motion primitives (such as approaching, picking, and placing) across diverse tasks throughout the rollouts. Notably, this coherent skill-space emerges during stage-1 training, without any explicit skill or task labels, or semantic loss functions. The semantic consistency we observe is a direct result of our carefully designed architecture, which inherently promotes and facilitates the sharing of representations across tasks. This emergent structure in the latent space underscores the effectiveness of our approach in flexibley capturing generalized low-level skills without task-specific supervision.
3. **Limited evaluation/statistical significance**: We report LIBERO and MetaWorld multitask results on 4 and 5 seeds respectively. These results show very low variance (as per error bars) demonstrating strong statistical significance. We’ve expanded our fewshot evaluation to further support our findings: the few-shot LIBERO results now encompass 9 seeds, while few-shot MetaWorld utilize 5 seeds throughout. Moreover, we've performed t-tests across all our results to rigorously validate their statistical significance. Please check out the PDF attached with this response for detailed results.

Additionally we would like to point out that QueST performs almost similarly with and without decoder finetuning with both variants still outperforms all baselines in fewshot setting on both benchmarks. This suggests that QueST stage-1 can effectively extrapolate in learned skill space and combine the skill tokens to generalize to unseen action sequences. This highlights the potential for QueST’s encoder (decoder) as an universal tokenizer (detokenizer) for continuous actions once trained on a large enough dataset like OpenX Embodiment. This semantically-sound, task-agnostic, temporally abstracted tokenization can better facilitate the learning of large behavior models like RT2 and OpenVLA as compared to their currently used naive discretization schemes.

Meta-review: The authors have diligently addressed many of reviewers' concerns. Experimental details have been clarified, additional analysis provided, and questions regarding novelty and latency resolved. While the initial reviews showed some variance, nearly all reviewers ultimately recommended acceptance, with the authors' detailed rebuttal and supplementary materials playing a key role in raising scores. I entirely agree with this assessment and applaud the authors for their thorough responses and the helpful website (including the t-SNE visualization!). The core strengths of the paper are a clear contribution to this field of research.

Therefore, I recommend accepting this paper.
Decision: Accept (poster)