QueST: Self-Supervised Skill Abstractions for Learning Continuous Control

Authors: atharva  mete
haotian  xue
albert  wilcox
yongxin  chen
animesh  garg
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 02:50:32

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.

Main paper: 
 introduction: 
 One of the grand goals of robotic learning is a general-purpose model that can learn from complex multitask demonstration data and generalize to new tasks in a zero-shot or few-shot manner. While such general-purpose models have become 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, we 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 a surge of work towards the goal of learning generalist policies from large, diverse datasets. Several papers have used techniques such as 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 a carefully discretized action space combined with a GPT-style transformer leads to impressive capabilities modeling multimodal behavior distributions [57, 16, 35] . In another vein, several works have attempted to 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 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 a 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'.A 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 phonetics 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 a unique encoder-decoder architecture primarily designed to impart causal inductive bias in action sequence data into the latent space. Such formulation enables us 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 a GPT-like transformer outperforms state-of-the-art baselines on 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 a detailed ablation and sensitivity study to validate our key architectural design decisions. 
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 a transformer: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 a 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 energy-based models [23, 29] or diffusion models [15, 71, 14, 68, 28] . The Behavior Transformer (BeT) line of work [57, 16, 35] shows that transformer-based categorical policies in carefully discretized action spaces do a good job handling multimodal demonstrator distributions and QueST builds upon this by contributing a 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 a transformer trained as a CVAE [60] to output chunks of actions performs well for a 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 a long line of work which handles sequential correlations through temporally-extended action abstractions [41, 50, 59, 35, 74, 76, 30] .In the past, robot learning researchers have approached multi-task decision making settings using a 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 a large focus on multi-task languageconditioned 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 on 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 a low-data regime. On the other hand LVMs like QueST are designed to learn sharable representations that can be applied to new tasks.LVMs, modeled by a paired encoder-decoder, have found extensive applications in computer vision [65, 32, 8] and generative models [56, 21, 13] . Recent studies also demonstrate the utility of latent space representations in robot decision-making, spanning offline RL [58, 50, 40] , 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 a 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 a self-supervised auto-encoder to learn skill codes, but their functionality relies on Model Predictive Control (MPC) using state prediction with ground-truth state-conditioned objective functions, which make it difficult to apply to real-world manipulation tasks. VQ-BeT [35] bears the strongest resemblance to QueST, also using a pre-trained discrete latent skill space to discretize the action space for a transformer-based 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.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 no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
problem setting: 
 We consider a datasetD = {{(O 1 , a 1 ), . . . , (O Ti , a Ti )} N k i=0 , L k } M k=0consisting of M robot interaction trajectories where a t is a continuous-valued action and O t is a tuple consisting of a highdimensional sensory observation. The data is collected via either human teleoperation or scripted policies for M different task each with a label L k . In our setting, O t consist of RGB image observations from the front camera and gripper camera (if available) along with proprioceptive state of the agent. L k is a natural language description of the k th task but can also be a one-hot encoding.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 a simple scalar quantization scheme. Here the input representation e is projected to very few dimensions (typically 3 to 5) which are then bounded and rounded, creating an implicit codebook. z = round_ste(f (e)), where f is the bounding functionGiven a feature vector e ∈ R d to quantize, instead of learning a parameterized codebook [65] and quantizing e by matching the nearest neighbor in the codebook, FSQ quantizes e by first bounding Proprioception Proprioception ProprioceptionTask CLIP/one-hotz 1 z 2 z 3 z nCausal Conv-1D it into certain range with f (e.g. f = ⌊α/2⌋ ⊙ tanh(e)) and then rounding each dimension into integer numbers directly with straight-through gradients (round_ste). α ∈ Z d defines the width of the codebook for each dimension (e.g. α = [8, 5, 5, 5], d = 4). Finally, it is easy to see that the size of the quantization space is Π d i=1 α i . An MLP can be used to transform z into continuous space of required dimension further.A common problem with vector-quantized codebooks (VQ) [65] is the under-utilization of the codebook. Recent works have attempted to 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 in our main experiments, but since many prior works in this space use VQ we also perform an ablation with it (see section 5.6).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We include codes in the supplementary materials.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
method: 
 In this section, we describe the key ideas behind Quantized Skill Transformer. In Section 4.1 we present our encoder-decoder architecture, which is designed to provide the flexibility to learn a wide range of skills with inductive biases to ensure that the learned skills are useful. In Section 4.2, we detail our skill prior, which we train to autoregressively predict codebook skills. Our full pipeline is shown in Figure 1 .As a motivating example, consider the task of lifting a pot and placing it on a stove beside. This consist of 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 are semantically meaningful so that they can be reused for new tasks, for example reusing the reaching skill for an object lifting task. In order to address these desiderata, we introduce the novel autoencoder architecture shown in Figure 1 consisting of an encoder ϕ θ and decoder ψ θ .The input to the encoder ϕ θ is an action sequence a t:t+T -1 sampled from the dataset, which we pass through several 1D 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-attention layers for sequence modeling. With a downsampling factor of F , the encoder outputs in total n = T /F embeddings. The embeddings are then quantized using FSQ as per the equation 1 into n discrete latent codes {z i } termed as skill tokens:(z 1 , . . . , z n ) = FSQ(ϕ θ (a t , . . . , a t+T -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 t cannot depend on input from any future timesteps. We found this inductive bias to 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 to maintain causality while attending to all codes, the decoder ψ θ 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 Z of skill codes, ψ θ reconstructs the original actionEQUATIONAs in [35] , the autoencoder is trained by minimizing the ℓ 1 reconstruction loss:EQUATIONUnlike prior work which often conditions on the state as well as 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 we show that our model learns generalizable abstractions that are shared and can be transferred across tasks.After training the encoder ϕ θ and decoder ψ θ , we train a skill prior π φ (Z|e, o) to predict skills Z = z 1:n corresponding to the demonstrator action distribution conditioned on a task embedding e and a length h sequence of image observations and proprioception inputs, o = (i t-h , p t-h ), . . . , (i t , p t ).We encode image observations with a 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 o t for a timestep t 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 nonlinear 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 π φ (Z|T o t-h:t , e) autoregressively as:EQUATIONwhere 
 is a 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 a sequence of demonstrator actions a t:t+T -1 and use the trained encoder ϕ θ to extract a latent skill vector Z t = z 1:n according to Equation 2. Then, we optimize π φ using the following negative log-likelihood loss:EQUATIONThe full skill prior pipeline is shown in Figure 1 .Few-Shot Finetuning: For few-shot finetuning on new tasks, we use a model pre-trained on large set of tasks and finetune it 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 a boost in the performance. Specifically, we finetune the decoder using following decoder loss:EQUATION)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.At inference time, QueST uses the skill prior π φ alongside the decoder ψ θ to sample actions. Conditioned on the encoded observation history T o t-h:t and task embedding e, we use top-k sampling with a temperature of τ to autoregressively sample a skill vector Ẑ ∼ π φ (•|T o t-h:t , e) from the skill prior. In practice, we find k = 5 and τ = 1 to work well across all environments. Then, we use the decoder to map the skill vector back to the action space, producing a sequence of predicted actions ât:t+T -1 = ψ θ ( Ẑ). In a receding horizon fashion, we execute the first T a ≤ T actions before replanning.We ablate the use of causal layers in various parts of our network in Table 2 . We see that removing causality from any part of our architecture leads to worse performance, suggesting that causal masking imparts the model with a helpful inductive bias for modeling robot action data.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 a modest downsampling factor of F = 4, and in Figure 4b we see that QueST does well with a 1024 codebook vectors. For more discussion on ablations please refer Appendix C. 
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 IL. Lastly, we perform some ablations to empirically justify our model design choices.We use the following benchmark suites to evaluate in the settings discussed above: LIBERO [38] is a lifelong learning benchmark featuring several task suites consisting of a variety of language-labeled rigid-and articulated-body manipulation tasks. Specifically, we evaluate on the LIBERO-90 suite, which consists of 90 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 in more detail below, we use the LIBERO benchmark to study the multitask IL, few-shot transfer and long-horizon IL 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.MetaWorld [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 5 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, we use learned task embeddings e for task conditioning.Baselines: We compare to the following baselines, which include similar discrete LVM pipelines as well as state of the art imitation learning algorithms: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.2. The UNet-based diffusion policy from [15] , which uses a 1D convolutional UNet to map samples from a Gaussian prior to action samples from the demonstrator distribution according to a learned denoising process. [73] , which trains a transformer as a CVAE [60] to predict action chunks. [35] , which learns a discrete latent space using a VQ-VAE [65] and uses a transformer to predict discrete latent codes. [74] , which first quantizes observation-action pairs and performs temporal abstraction using byte pair encoding (BPE) to learn a skill token vocabulary, which it uses as an action space for few-shot learning. For a detailed discussion between QueST and the baseline methods, please refer to Appendix B.5.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 MetaWorld).In Figure 2a , we present the average success rate across 90 tasks in LIBERO-90 against the aforementioned baselines. Quantized Skill Transformer achieves state-of-the-art results on LIBERO-90 benchmark, outperforming the baseline VQ-BeT and Diffusion Policy by a margin of 8 and 13% respectively. We attribute its performance to 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 to 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 a simpler benchmark, all methods perform almost similar in Multitask-IL setting which is 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 a continuous latent variable with the same dimensionality as the actions. However, the consistent outperformance of QueST over the diffusion policy provides compelling evidence for the benefits of using a bottlenecked latent variable. This bottleneck encourages the model to learn shared representations, resulting in enhanced performance. While both VQ-BeT and PRISE employ a latent bottleneck, VQ-BeT's architecture neglects the inherent inductive biases in the action data, which we believe results in a less well-structured latent space. PRISE incorporates this using a latent forward transition model, but is bottlenecked by the use of BPE which we posit is not suitable for such a dynamic latent space.In this setting we take the pretrained model from section 5.2 and test its 5-shot performance on unseen tasks from LIBERO-LONG and held-out set in ML45. We sample only five demonstrations for each task, generate the skill tokens using pretrained encoder and use them to finetune the skill prior and the decoder as described in Section 4.2. We also present the results without finetuning the decoder (frozen ψ θ in figure 2b & 3b) to validate its generalization to unseen skill tokens sequences.Figure 2b shows the average success rate for 5-shot IL across 8 unseen tasks in LIBERO-LONG. QueST achieves SOTA performance, surpassing all other baselines by an absolute margin of 14%. Table 1 : Success rates after ablating design details of QueST. We present the mean across four random seeds and error tolerances show the standard error.Though we see a 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 a fair comparison, we also tried fine-tuning the decoder of VQ-BeT but did not observe any gains from it. VQ-BeT struggles in this setting as it heavily relies on offset head to output continuous action corrections which requires more data-samples for sufficient coverage in the continuous action space.Figure 3b shows the average success rate for 5-shot IL across 5 unseen tasks in MetaWorld. 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 5 demonstration examples. For few-shot evaluation protocol please refer Appendix D.1.In this setting, we aim to purely study and compare the performance of our model on long-horizon tasks. We train the model (both stages) solely on LIBERO-LONG complete dataset (50 demonstrations 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 a good latent space in LVMs. QueST still outperforms all LVM baselines by a large margin demonstrating its long-horizon modeling capabilities.Overall, we see that our model outperforms baselines like VQ-BeT in multitask settings, showing stronger modelling capacity. At the same time, it has the correct latent structure to outperform baselines like diffusion in few shot settings, especially even with frozen decoder, indicating strong generalization capabilities of learned skill-space.Our pipeline runs at 33Hz, which is more than suffcient for vision-based real-robot control where most camera systems run at 30 fps. For comparison, our implementation of Resnet- We validate the proposed architecture by ablating some of its key design decisions. All the ablations are 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.2. Observation Conditioned Decoder: 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 1 . We see that conditioning on observations leads to a deterioration in performance.3. Mirrored Decoder: Following a typical autoencoder design, we use a 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 1 . We see that this method performs worse, suggesting attending to all quantized codes in z, as our decoder does, is important for faithfully predicting actions.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 n z-embeddings at a 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-90. For better comprehension, we encourage readers to 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 a discrete latent space. The key idea behind QueST is to represent action sequences as a series of codebook vectors, and we 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-ofthe-art baselines in multitask and few-shot learning settings. Our results highlight the usefulness of QueST's encoder (decoder) as semantically-sound, task-agnostic tokenizer (detokenizer) for continuous actions, and its potential to leverage Large Multi-modal Language Models in stage-2.Limitations: While the benchmarks we consider encompass a 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 a 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-Embodiment [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, to enhance abstraction learning.A statement on societal impact: This paper works towards the broader goal of automating a wide range of 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 a handful of companies. It is important that we in the machine learning community advocate for equitable use of the technology we develop.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to 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 of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
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 on 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, confidence intervals, or statistical significance tests, at 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 run with given experimental conditions).• 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). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to 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 is 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: 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)?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 a 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: For further results and videos please see our website. https://quest-model.github.ioWe present hyperparameters in the following tables: For vision encoder we used a shallow Convolutional Neural Network (CNN), consisting of the first four layers of ResNet18 [27] followed by a spatial softmax [36] . In encoder, we use causal convolution layers from [30] . For transformer blocks, we used the transformers library from hugging face https://huggingface.co/docs/transformers/ with appropriate masking for ensuring causality.To ensure fair comparison of different model architectures, we use same input modalities and same observation & task encoders for all baselines, for both LIBERO and Metaworld experiments. VQ-BeT needs a 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, we also tune some key hyperparameters for the baselines and present the results for best performing ones.1. ResNet-T: Transformer trunk'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 we use an observation history of 10 timesteps. As per original paper ablations [15] an observation history of 1 was used.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 1 encoder by default is a single MLP 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-VQ 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, 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.The models are implemented in PyTorch. For all our experiments we use a server consisting of 8 Nvidia RTX 1080Ti 10GB memory each. And all our models easily fit on one GPU for training.Many recent works use discrete latent variable models as a mechanism to learn shared abstractions over continuous low-level skills. QueST, VQ-BeT and PRISE all do this and perform two-staged learning. However, in this work we propose several important architecture choices leading to QueST's strong performance. Specifically, QueST encodes actions to a sequence of n encodings (skill tokens) using a novel autoencoder that captures temporal correlation within an action sequence with causal convolution and masked self-attention layers.Like QueST, VQ-BeT performs temporal abstraction by encoding a sequence of actions into one state-independent latent vector. It concatenates input actions and encodes the sequence to just one single latent encoding using an 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 a single encoding is limiting as 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 n 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 n 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 on continuous offset prediction (whereas QueST can achieve high success rate using only the output from its action decoder).PRISE performs temporal abstraction by learning discrete codes for state-action pairs and using BPE to group common token sequences into higher-level skills. However, BPE is known to suffer with evolving language leading to a suboptimal character-level tokenization, and it might struggle to effectively encode new action sub-sequences that the model has not seen before (eg. stitching sequences in between two tasks in LIBERO-LONG). 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 as combinations of its learned latent codes. Our few-shot results without decoder-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.Despite tuning the hyperparameters, the multitask performance of ACT is very low (54%) on LIBERO. Moreover, as per BAKU [26] , ACT's performance drastically reduces when number of demos are reduced to 35-per-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 we 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.For aiding this discussion we present the ablation results again in • Replacing FSQ with VQ still outperforms VQ-BeT in few-shot setting suggesting that QueST's superior performance is not only due to a better quantization scheme but also due to it architecture that flexibly maps an input sequence to multiple embeddings and allows for efficient transfer.• It's tempting to ground the mapping between z-tokens and actions with observation tokens with an intuition that z-tokens will define a coarse set of actions and observation tokens will aid finer action decoding. But we observe worse performance with this. We hypothesize that the reconstruction objective forces encoder and decoder for most optimal quantization at the bottleneck layer but with extra observation information the decoder might focus more on observation tokens in turn hurting the quantization. This observation goes hand-in-hand with a closely related prior work SPiRL [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 cross-attention architecture that allows for attending to all z-tokens and maintaining causality at the same time. 6 : Success rates after ablating the causality in QueST.• We 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 later 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 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 has semi-circle in them for example C-shaped or infinity-shaped trajectories.Frozen ψ θ Finetuned ψ θ 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 IL.• 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.Fewshot Evaluation Protocol: In finetuning phase, we finetune ResNet-T, VQ-BeT & QueST for 100 epochs and ACT & Diffusion Policy for 200 epochs. For each task in MetaWorld, we evaluate each method across 10 evenly spaced checkpoints for 5 seeds on 50 distinct initial states and report the results corresponding to the best performing checkpoint. For Libero, we found the final checkpoint to perform best for all methods and hence report results corresponding to it across 9 seeds. Justification: 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 .Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Correspondence to: amete7@gatech.edu 38th Conference on Neural Information Processing Systems (NeurIPS2024). 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: The only major weakness of the author’s work is the limited evaluations. If the authors can justify using just three seeds across experiments, that would build more confidence. Otherwise, as this system targets robotic learning, it would have been good to see results on a real robotic system instead of just in simulation. The authors use large transformer models, so latency concerns could be relevant if their system is computationally slow. We are also concerned with the lack of error bars in Figure 4, which appears only to use a single seed for the ablation experiments, casting doubts on any conclusion the authors make from these results. The authors should clarify these details and run additional experiments to show the robustness of their results if they have not.
questions: questions: - How many seeds were used in Figure 4 experiments? 
- If the latent variables are used as key and query values, is it fair to say the generated skills can be imagined as some interpolation between a subspace of the fixed-sized vectors constructed by the positional encodings?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: I am happy to increase my score if these points are addressed.

**Comments**
- Should probably include analysis on what the latent $z$'s ended up learning. Do they actually have some temporal abstraction going on? Can we associate them? It appears that this is in the website but I believe this should be in the main paper as one premise is that the latent representation is leveraging the structure of the actions.
- This work appears to rely on the fact that the demonstrations have some form of structures---probably a limitation. Perhaps in this experimentation setting it is a fair assumption, but what if we have partially observed setting (e.g. items with dynamic inertial properties).
- It is not clear to me how this approach handles multimodality (in the sense of action distributions?)
- Ablation on "causality": I feel the word "causality" is misused because this is simply a mask enforcing whether or not to look at the future data. Is this not still modelling correlation only? How is this causal exactly?
questions: questions: **Questions**
- Equation 1, perhaps write out what round_std is.
- Page 5, lines 179-180: What is the intuition of using cross attention between positional embedding and skill tokens but not self-attention with positional encoding?
- Page 5, lines 186-187: Do the authors think that this is because the actions can be invariant (e.g. in velocity control the sequence of action is likely similar if we are reaching downwards to the items)? How would stochasticity of the dynamics play into learning this?
- Figure 2: Any intuition why few-shot is better than multi-task for QueST?

**Possible typos**
- Page 8, line 305: viz. means visualized?
limitations: limitations: - The current application is only on robotic manipulation but it will be interesting to extend this to other domains
paper_summary: paper_summary: - This work introduces a new method called QueST that uses a latent variable model to learn a set of discrete temporal action abstractions / motion primitives / skills for imitation learning. This is done by training an auto-encoder to encode a sequence of actions using a causality preserving encoder, discretising the embedding using Finite Scalar Quantization (FSQ) to obtain a sequence of discrete latents, and then decoding the latents using a causality preserving decoder to recover the sequence of actions.
- A policy is learned by predicting the latent actions from a history of observations and a language instruction while keeping the latent decoder frozen. An exception to the latter is when fine-tuning a policy in a few-shot setting. At test time, top-k sampling for sampling an optimal sequence of latents.
- The method is validated on the LIBERO and Meta-World environments and compared to several recent approaches: ResNet-T, diffusion policy, ACT, VQ-Bet, and PRISE.
strength_weakness: strength_weakness: Originality:
- Two of the baselines, VQ-BeT and PRISE, also learn temporal action abstractions. It would be nice to have a more detailed and clearer discussion of the differences between these three methods than is currently provided, e.g., differences and similarities could be summarised in a table.

Quality:
- There are no results for PRISE on Multitask LIBERO-LONG.
- The paper website states: “Depending on the dataset, we also tune some key hyperparameters for the baselines”. This is vague and it is unclear why this was not done for every dataset.
- It is not clear whether all the baselines are based on new experiments or whether some of the results are shown as reported in previous works.
- There is no analysis of inference latency of the proposed method and the baselines, which is important for practical applications.
- It might be interesting to have an ablation experiment that uses Residual VQ as used by VQ-BeT instead of FSQ to assess if this choice plays a significant role for the difference in performance.

Clarity:
- Some information about how the hyperparameters of the baselines were tuned are on the paper website, but important information like this should be in the main body of the paper or the appendix.
- The paper website states that “To ensure fair comparison of different model architectures, we use same input modalities and same observation & task encoders for all baselines…”. This is important information that should be in the main body or the appendix of the paper. It is also unclear whether this only applies to the Meta-World experiments or also the LIBERO experiments.
- It is not specified where the “provided scripted policies” for collecting demonstrations on Meta-World as mentioned on the paper website can be found.
- Minor comment: Figure 3b is not described or mentioned in the text.

Significance:
- The proposed method feels a bit like an incremental evolution of VQ-BeT and PRISE. Even though these have only come out very recently, it would still be good to have a clearer positioning with respect to these two works as already mentioned above.
questions: questions: - Are the updated results on Meta-World going to replace the current results in the main body of the paper?
- Where can the “provided scripted policies” for collecting demonstrations on Meta-World as mentioned on the paper website be found?
- Are all results for the baselines based on new experiments or are any of the results shown as reported in previous works? In the latter case, it would be useful for this to be indicated in the result tables.
- How were the hyperparameters tuned for the proposed method, for the ablations, and for baselines?
- How does the proposed method compare to the baselines in terms of latency?
- Is the code going to be released publicly?

Updates:
- Presentation: 2 -> 3
- Score: 4 -> 6
limitations: limitations: - It would be nice to have a discussion of latency considerations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - The overall pipeline that first learns the abstract representations of behaviors and then models the behavior with the autoregressive model is common. I believe the main difference of this paper contains two points: (1) using FSQ to learn discrete latent space; (2) implement the encoder with a causal transformer instead of a non-causal transformer. So the innovations seem not enough. 

- A lot of ablation studies are conducted. However, the related analysis is not sufficient.
questions: questions: - Why the causality of the encoder transformer is so critical?
limitations: limitations: Yes

Comments: title: Gentle reminder!
comment: Thank you for your suggestion and further questions.

> Have you considered running an experiment to examine how the performance of VQ-BeT changes when increasing the size of the latent space?

There is a slight misunderstanding here. If by increasing latent space you mean trying a larger codebook size then yes, we do try an effective codebook size of 4096 for VQ-BeT but did not observe any performance gain (<1% variation in LIBERO-90). This follows what VQ-BeT authors report for their codebook size ablations (Table 12 in VQ-BeT paper). What we meant by a "smaller latent space" is the fact that VQ-BeT concatenates input actions and encodes the sequence to just one single latent encoding using an MLP. A sampled action sequence might contain multiple motion primitives of variable length and start point, and capturing them with a single encoding is limiting as 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 'n' encodings using its encoder, specifically designed to model temporal correlation within input actions. With a codebook size of C, QueST’s effective latent space is C^n while that of VQ-BeT is C, hence a smaller latent space. While we report QueST results for chunk size 32, we did observe very little variation (<1% in LIBERO-90) with size 16,48 and 64, indicating robustness to this hyperparameter (a major issue in tuning chunking based methods like ACT, MT-ACT).

> Could you elaborate on to what extent this is an issue? Temporal action abstraction should allow the reuse of existing abstractions for new tasks, i.e., it might not be necessary to encode “new actions”.

The reuse of existing abstractions indeed happens, hence PRISE reports a non-zero success rate in the few-shot setting. However, new tasks will definitely contain new action sub-sequences that the model has not seen before (eg. stitching sequences in between two tasks in LIBERO-LONG). BPE relies on frequency statistics from the pretraining data and hence might lead to inefficient tokenization of such new sub-sequences. 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. Our few-shot results without decoder-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.

> Minor comment: I would suggest adding a footnote or similar to indicate that the results shown for PRSE are taken from the original PRISE paper.

We'll update this in the main paper.

We hope we have been able to address your questions in the above clarifications. Kindly let us know if you have additional questions or concerns that stand between us and a higher score.
Thanks for the feedback and for raising your score. Your feedback has been a great help to improve the presentation and clarity of the paper.

> Regarding to the analysis…

This is a good point. We’ll move a subset of these plots to the main body of the paper as well as some analysis describing how they demonstrate QueST’s capabilities to learn shared skills across several tasks.

> Regarding structure…

Thank you for raising this concern. An example of a dataset without temporal/trajectory data is the ARNOLD[1] benchmark that has current observation as state (s) and next gripper keypoint as action (a) data. While our proposed method does require temporal data, we believe that the same architecture could work for (s,a) pairs. In fact, we initially did test the architecture on ARNOLD, where the input to the encoder was (s,a) pairs and the decoder output was actions (a), and were able to achieve fairly low reconstruction loss. However, we currently leave this application of the architecture to future work. Since we haven’t evaluated QueST rigorously in such a setting, we’ve added a sentence to the Problem Setting section (3.1) emphasizing this assumption. 

Additionally, we’d like to point out that this assumption is extremely common in recent behavior cloning literature. All of our SOTA baselines (VQ-BeT, PRISE, ACT, Diffusion Policy) make a similar assumption, and several recent large scale robotics datasets (Open-X Embodiment, DROID, BridgeData, etc) and popular behavior cloning benchmarks (Robomimic, Mimicgen, LIBERO, Metaworld, CALVIN, RLBench, Franka Kitchen, D4RL, etc.) are compatible with this assumption. Thus, while this assumption may limit the method’s applicability to some settings, it enables scalability and still fits in well with the field as a whole.

Thanks again for your valuable feedback. We would appreciate hearing about any further limitations we can address in order to further increase the score.

[1]: Gong, Ran, et al. "ARNOLD: A benchmark for language-grounded task learning with continuous states in realistic 3D scenes." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.
Thank you for the rebuttal.

> “Unlike QueST, VQ-BeT’s relatively small latent space limits its expressive capacity”

Have you considered running an experiment to examine how the performance of VQ-BeT changes when increasing the size of the latent space?

> “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”

Could you elaborate on to what extent this is an issue? Temporal action abstraction should allow the reuse of existing abstractions for new tasks, i.e., it might not be necessary to encode “new actions”.

> “We were unable to successfully recreate PRISE results so we report the results from their paper”

Minor comment: I would suggest adding a footnote or similar to indicate that the results shown for PRSE are taken from the original PRISE paper.
Thank you for the response. In short I have raised my score.

Regarding to the analysis: I feel the main paper should include this because that is verifies the claim that the model is able to learn skill abstractions. Even if the "performance" is not as well I believe in that case it would still deliver the point.

Regarding structure: Yes, I meant tasks like manipulation/physical tasks have inherit structure, and the method itself is able to leverage this through the data. Would this method still work if the dataset is not stored as trajectories? If I understand correctly this method requires temporal data.
We kindly request all reviewers to go through the rebuttals and let us know if their concerns have been addressed. Please let us know if there are any more questions/clarifications. Thank you.