Actra: Optimized Transformer Architecture for Vision-Language-Action Models in Robot Learning

## Abstract

Abstract Vision-language-action models have gained significant attention for their ability to model trajectories in robot learning. However, most existing models rely on Transformer models with vanilla causal attention, which we find suboptimal for processing segmented multi-modal sequences. Additionally, the autoregressive generation approach falls short in generating multi-dimensional actions. In this paper, we introduce Actra, an optimized Transformer architecture featuring trajectory attention and learnable action queries, designed for effective encoding and decoding of segmented vision-language-action trajectories in robot imitation learning. Furthermore, we devise a multi-modal contrastive learning objective to explicitly align different modalities, complementing the primary behavior cloning objective. Through extensive experiments conducted across various environments, Actra exhibits substantial performance improvement when compared to state-of-the-art models in terms of generalizability, dexterity, and precision.

## 1 Introduction

Vision-language-action models (VLAs) have emerged as integral components of recent developments in robot learning. Previous multi-modality models, exemplified by vision-language models (VLMs), have demonstrated proficiency in handling both visual and textual inputs, successfully addressing a spectrum of tasks such as visual question answering, image captioning, and image retrieval. Distinctively, VLAs extend beyond the capabilities of VLMs by incorporating the ability to execute actions based on multi-modal inputs. This unique capability empowers VLAs to interpret language prompts, visually perceive their environment, and subsequently execute actions to fulfill the specified tasks. The potential applications of VLAs in robotics are not confined to controlled environments in traditional domains like manufacturing. They also prove their suitability for everyday tasks such as room cleaning and cooking , thanks to their dexterity and generalizability.

To accommodate multi-modal inputs, previous Transformer-based VLMs explored designing special self-attention schemes to better suit the unique properties of different modalities, such as UniLM , M6 , BLIP-2 and Octo . Consider the task of image captioning as an example, causal attention is not the best option to encode images because there is no clear causal relationship among the image patches. Thus, these VLMs allow bidirectional self-attention for the image tokens while maintaining causal attention for the text tokens.

Figure: (a) Causal Attention.
Refer to caption: x1.png

VLAs predominantly build upon the pioneering foundations laid by Decision Transformer and Trajectory Transformer . These two works frame reinforcement learning (RL) policies as sequence modeling problems, leveraging the expressive power of Transformer models. This paradigm has become a cornerstone across recent VLAs, but both models are Transformer decoders based on causal attention. Subsequent approaches, such as Gato and RT-1 , also adopt Transformer decoders as their network backbone, passing in different modalities as a single sequence. VIMA incorporates cross-attention mechanisms to condition the policy with multi-modal prompts, but the decoder stack still follows previous methods and uses causal attention.

On the contrary, we have identified that VLA trajectories in robotics exhibit unique properties better captured by a novel type of Transformer self-attention, as illustrated in Figure 1 & 2. Specifically, each language prompt, state, or action within a VLA trajectory can consist of multiple tokens, referred to as a “segment” in this paper. For instance, robot systems usually make use of several cameras, and as a result, a state is represented with a segment of tokens, each corresponding to a camera. Similar to state tokens, tokens for action dimensions also lack causal relationships with each other. Traditional causal attention hinders full information flow within a segment, as tokens are restricted from attending to the subsequent tokens. To overcome this limitation, we introduce trajectory attention, optimized for VLA trajectories. Trajectory attention possesses two key characteristics: inter-segment attention is causal, and intra-segment attention is bidirectional. Since a VLA model only needs to encode the language prompt and follow the corresponding instruction, causal attention is also unnecessary for the prompt segment. Consequently, we advocate for processing trajectories at the segment level, rather than merely at the token level.

To complement trajectory attention, we devise a segment-level decoding scheme that generates a segment as a whole. Drawing inspiration from DETR’s object query , we propose employing action queries to more effectively extract information for action generation. Concretely, we employ one learnable action query for each action dimension. Each action query aggregates the most relevant information in the trajectory for its corresponding action dimension and generates the most probable value for that dimension. Different action queries can execute this procedure in parallel, facilitating the simultaneous generation of all action dimensions. This represents a substantial acceleration in action generation speed compared to earlier approaches that generate one action dimension at a time, such as RT-2 . By combining trajectory attention and action queries, we introduce an optimized Transformer architecture for VLA trajectories, which we name Action-query-based Trajectory-attention Transformer, or Actra for short.

Training the policy network with behavior cloning is a common practice for VLAs. However, it lacks explicit alignment between different modalities. Contrastive learning, which has been widely adopted to enhance inter-modality interaction, gained prominence in VLMs through the success of CLIP . CLIP achieved state-of-the-art performance on image classification tasks by introducing a contrastive language-image pre-training approach on a large-scale dataset of image-text pairs. Despite its popularity in VLMs, contrastive learning remains relatively under-explored in VLAs. While R3M proposed a time contrastive learning objective, and VIP introduced value-implicit pretraining, both primarily focus on obtaining a robust visual encoder without emphasizing alignment among multiple modalities.

Therefore, we propose a VLA contrastive learning objective aimed at enhancing multi-modality encoding capability. Positive trajectories are obtained by introducing slight noise to the actions. The rationale is that a slightly noisy trajectory can still be valid, as a path in the close vicinity of an expert path can also lead to the target position. To construct negative trajectories for the three modalities of VLAs, we sample prompt, state, or action segments from other trajectories in the dataset and substitute the original segments with those mismatched ones. We then train Actra to encode positive and negative trajectories and contrast them with the widely adopted InfoNCE loss . This approach explicitly encourages multi-modality information exchange among the vision encoder, language encoder, and action decoder of Actra. With the enhanced multi-modality encoding capability, VLA’s performance on primary robotics tasks is further improved.

The three main contributions of this paper are:

- •
We introduce Actra, an optimized Transformer architecture featuring trajectory attention and action query, designed to model VLA trajectories on the segment level;
- •
We propose a VLA contrastive learning objective to explicitly enhance Actra’s multi-modality encoding capability, complementing the primary behavior cloning objective of robot imitation learning;
- •
Extensive experimental results across three environments demonstrate that Actra significantly outperforms state-of-the-art VLA models in terms of generalization, dexterity, and precision, showcasing the effectiveness of our method.

Figure: Figure 2: The architecture of Actra. Each square represents a token. A trajectory $\tau$ consists of a prompt segment $p_{1:4}$, state segments $s_{1:2,t}$, action segments $a_{1:3,t}$. For clarity, the trajectory is a simplified example and does not reflect the actual specifications of Actra. In the Transformer, vertical dashed lines divide the segments. Learnable action queries $q_{1:3,t}$ are inserted after each state segment to extract information for action generation. Each token embedding (orange dot) in the trajectory can attend to embeddings from all previous segments (horizontal arrows), as well as all the embeddings in its own segment (gray lines). Notably, action queries, which contain no trajectory information, are hidden from other tokens, but they can still attend to all preceding tokens. In addition to its primary function of decoding actions, Actra can also encode the entire trajectory by pooling the embeddings in the last segment (red box).
Refer to caption: x3.png

## 2 Related Work

#### Vision-language-action Model.

Vision-language-action models (VLAs) constitute a class of multi-modal models designed to generate actions based on specified language prompts and perceived environment. Coined by RT-2 , VLAs have garnered increasing attention due to their dexterity and generalizability in handling complex robotics tasks. Early attempts were based on existing vision-language models (VLMs), exemplified by CLIPort and BC-Z . Gato explored the use of a single Transformer model as the control policy for tasks spanning various domains, unifying multi-modal inputs into a single sequence. RT-1 stands as a dedicated robotics transformer for robotics tasks. Our model is also a VLA but with an optimized Transformer architecture.

#### Multi-modal Transformer.

Several VLMs, including UniLM , M6 , and Octo have endeavored to optimize Transformer’s self-attention for vision-language inputs. Despite these efforts, adapting self-attention to the multi-modal inputs of VLAs has been relatively unexplored. Gato and RT-1 maintain causal attention in Transformer decoders. VIMA proposes passing the prompt into the policy through cross-attention, but their Transformer decoder stack still employs causal attention. To the best of our knowledge, Actra is the first VLA designed to accommodate VLA trajectories with a unique self-attention mechanism.

First introduced in DETR , learnable object queries have shown promising results in extracting information for object detection. BLIP-2 used a similar strategy to extract visual embeddings for vision-language tasks. In our approach, we employ learnable action queries at the action-dimension level to extract information most relevant to individual action dimensions.

#### Multi-modal Contrastive Learning.

A series of VLMs, CLIP , ALIGN , Florence , FILIP , has demonstrated the significance of contrastive learning in enhancing multi-modal interaction. However, in VLAs like R3M and VIP , where contrastive learning has been adopted, the primary emphasis remains on improving visual representations. In contrast, our proposed VLA contrastive learning task explicitly compels the model to align the three modalities, allowing for more effective encoding of multi-modal inputs. More related work is included in Appendix A.8.

## 3 Our Method

### 3.1 Preliminaries

Markov Decision Process (MDP) comprise states ($s$) and actions ($a$) and it can be conditioned by a language prompt ($p$). In the context of imitation learning, a VLA trajectory within the language-conditioned MDP is denoted as $\tau=(p,s_{t=1},a_{t=1},\dots,s_{t=T},a_{t=T})$. Each element in the trajectory—$p$, $s_{t}$, or $a_{t}$—comprises a segment of tokens. For instance, a state $s_{t}$ corresponds to the segment $s_{1:M,t}=(s_{1,t},s_{2,t},\dots,s_{M,t})$, where each element is a token. Tokens in $p$ are standard NLP tokens. State tokens in $s_{t}$ correspond to images or object poses. Action tokens in $a_{t}$ contain 6D poses or 2D coordinates. Therefore, a trajectory at the token level is written as $\tau=(p_{1:L},s_{1:M,t=1},a_{1:N,t=1},\dots,s_{1:M,t=T},a_{1:N,t=T})$. The goal is to train a policy that can generate an optimal action based on the past trajectory $\pi_{\theta}(a_{t}|p,s_{\leq t},a_{<t})$.

### 3.2 Actra

In natural language generation (NLG), language models such as GPT employ Transformer decoders as the backbone. To prevent tokens from having visibility into subsequent tokens, a causal attention mask is applied in the Transformer decoder. Prior VLA models ; have followed this trend for action generation. While causal attention is well-suited for NLG, where language tokens are sequentially generated, it is not the optimal attention mechanism for modeling VLA trajectories.

#### Trajectory attention.

Images of the state $s_{t}$ from multiple cameras arrive simultaneously, lacking causality among themselves. They are determined solely by the preceding action $a_{t-1}$ and the environment. The same principle applies to actions: $a_{t}$ is only dependent on previous states and actions, and they are conditionally independent from each other. The action dimensions in an action do not exhibit a clear causal order. For instance, in a 3D coordinate, it is uncertain whether $a_{1,t}$ depends on $a_{2,t}$ or vice versa. Regarding the language prompt, as it is provided by the user, the model’s job is to encode and understand it rather than generate the prompt, akin to BERT . Vanilla causal attention might impede information flow within each segment of a VLA trajectory, prohibiting $s_{1,t}$ from attending to $s_{2:M,t}$, and $s_{2,t}$ from attending to $s_{3:M,t}$, and so forth. This similarly holds for prompts and actions.

To address the issue, we propose an optimized Transformer self-attention mechanism for language-conditioned VLA trajectories, termed trajectory attention. Trajectory attention exhibits two key properties: the inter-segment connections are causal, and the intra-segment connections are bidirectional. Its corresponding attention matrix is illustrated in Figure 5. Following the convention of the Transformer attention matrix, we designate the row index as the destination of self-attention and the column index as the source. Consequently, the causal attention matrix has all its lower triangle entries, $(i,j)$ for $i\geq j$, set to one, and the rest set to zero. Trajectory attention is achieved by unmasking the entries in the causal attention matrix corresponding to $(p_{i},p_{j})$, $(s_{i,t},s_{j,t})$ or $(a_{i,t},a_{j,t})$ for $i<j$. When compared with the original causal attention, there are $L(L-1)/2+T\big{(}M(M-1)/2+N(N-1)/2\big{)}$ additional entries joining the self-attention in every Transformer layer, which explains the effectiveness of trajectory attention. Consequently, Actra is designed to process VLA trajectories at the segment level, which aligns well with the MDP setting as it involves states and actions rather than individual tokens.

#### Action query.

Adapting to the segment-level trajectory attention mechanism, we introduce a segment-level decoding scheme based on learnable action queries. Most prior VLAs generate action dimensions autoregressively, where each action dimension depends on its preceding token embedding . However, this approach is suboptimal because the embedding of the preceding token is highly dependent on its input and may lack the most relevant information for the action dimension. For instance, when generating $a_{1,t}$, its preceding token is $s_{M,t}$. Although the embedding of $s_{M,t}$ can aggregate information from the past trajectory through self-attention, it is largely influenced by its corresponding input image and may not contain sufficient information about $a_{1,t}$. To overcome this limitation, we adopt learnable action queries $q_{1:N}$ for individual action dimensions $a_{1:N}$, inspired by DETR . Each action query $q_{i}$ is dedicated to one action dimension $a_{i}$ and is shared across all timesteps: $q_{i,t=1}=q_{i,t=2}=\dots=q_{i,t=T}$ for $i\in\{1\dots N\}$. We argue that this approach can find more relevant information for each action dimension because the action query $q_{i}$ can exclusively attend to information pertinent to $a_{i,t}$. Since action queries have no associated input token, their embeddings fully retain action dimension information. Moreover, distinct from autoregressive generation, action queries can extract information and generate all dimensions of an action segment in parallel. Consequently, the decoding procedure operates at the segment level. This significantly speeds up action generation. As the action queries are solely used for information extraction and do not hold any trajectory information, they are masked out from the attention matrix, ensuring that other tokens cannot see them through the self-attention mechanism.

Figure: Figure 3: VLA contrastive learning. Given the anchor in the first row, we construct a positive by adding slight noise to the actions. For each modality, we create negatives by injecting mismatched segments sampled from the dataset. Using Actra as a trajectory encoder, we contrast their trajectory embeddings as a complementary objective to robotic imitation learning.
Refer to caption: x4.png

#### Actra

Combining trajectory attention and action query, we introduce a novel Transformer variant named Actra. In Actra, all action tokens $a_{i,t}$ can fully attend to $(p,s_{t=1},a_{t=1},\dots,s_{t},a_{t})$, and all state tokens $s_{i,t}$ can fully attend to $(p,s_{t=1},a_{t=1},\dots,s_{t})$. Consequently, their embeddings are enhanced for VLA trajectories. Each action query $q_{i,t}$ aggregates this enriched information, collecting more pertinent information for its corresponding action dimension. This makes Actra a more suitable Transformer for action generation in VLA trajectories. The training process utilizes standard behavioral cloning in robotic imitation learning, optimizing the objective $\mathcal{L}_{\text{BC}}=\min_{\theta}\sum_{t=1}^{T}-\log\pi_{\theta}(a_{t}|p,s
_{\leq t},a_{<t})$ on offline expert trajectories.

### 3.3 VLA Contrastive Learning

The primary behavior cloning objective involves a generative decoding process, where Actra generates the next action based on the past trajectory. Although this implicitly requires the model to encode multi-modal inputs, we argue that the encoding capability can be further enhanced with an explicit objective. Several studies have demonstrated the efficacy of contrastive learning in aligning different modalities ; ; ; . Therefore, we propose a novel VLA contrastive learning objective to train Actra to encode and contrast different modalities, aiming to improve robotic imitation learning, as illustrated in Figure 3.

Concretely, we assume the anchor trajectory is $\tau=(p,s_{t=1},a_{t=1},\dots,s_{t=T},a_{t=T})$. To construct a positive, $\tau^{+}$, we introduce small Gaussian noise to the actions, following the intuition that a slightly perturbed path can still lead to the destination. Subsequently, we design three types of negatives, each corresponding to one modality of VLAs. For language prompts, we create a negative by selecting an incorrect prompt $p^{\prime}$ from the dataset and substituting it for the correct prompt. The resulting negative trajectory for the language modality is $\tau^{-}_{p}=(p^{\prime},s_{t=1},a_{t=1},\dots,s_{t=T},a_{t=T})$. Although prompts are not unique to trajectories, states and actions are unique, allowing us to find false state and action inputs within the batch. Given another trajectory $\tau^{\prime}=(p^{\prime},s^{\prime}_{t=1},a^{\prime}_{t=1},\dots,s^{\prime}_{
t=T},a^{\prime}_{t=T})$ in the same batch as $\tau$, we can mismatch the states and actions to assemble negatives for the vision and action modalities: $\tau^{-}_{s}=(p,s^{\prime}_{t=1},a_{t=1},\dots,s^{\prime}_{t=T},a_{t=T})$ and $\tau^{-}_{a}=(p,s_{t=1},a^{\prime}_{t=1},\dots,s_{t=T},a^{\prime}_{t=T})$. Finally, we contrast the positive $\tau^{+}$ against three types of negatives: $\tau^{-}=\{\tau^{-}_{p},\tau^{-}_{s},\tau^{-}_{a}\}$.

In VLA contrastive learning, our Transformer model encodes the entire trajectory into a sequence of embeddings. A max-pooling layer then aggregates them into a single embedding, as shown in Figure 2. Finally, we utilize the standard InfoNCE as the objective to train the model to distinguish between positive and negative embeddings:

$$ $\begin{split}&\mathcal{L}_{\text{VLA-CL}}(\tau,\tau^{+},\tau^{-})=-\log\mathbb {E}\Bigg{[}\frac{\exp(f(\tau)^{\intercal}f(\tau^{+}))}{\exp(f(\tau)^{\intercal }f(\tau^{+}))+\sum\limits_{i=1}^{B}\sum\limits_{j\in\{p,s,a\}}\exp(f(\tau)^{ \intercal}f(\tau^{-}_{i,j}))}\Bigg{]},\end{split}$ (1) $$

where $f(\cdot)$ is Actra coupled with the max-pooling layer; $B$ is the batch size. During training, we mix the VLA contrastive learning the primary behavior cloning objective to establish the overall training loss: $\mathcal{L}=\alpha\mathcal{L}_{\text{BC}}+\beta\mathcal{L}_{\text{VLA-CL}}$.

## 4 Experimental Results

**Table 1: Performance comparison of success rate (%) on different generalization levels. The numbers under the “Distractors” category specify the quantity of distractors present in the environment.**
|  | Unseen tasks | Distractors |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Params | Color | Size | Shape | Container | Both | Mean | 0 | 2-4 | 6-8 | Mean | Overall |
| RT-1 | 46M | 27.03 | 6.36 | 20.30 | 0.79 | 1.27 | 11.15 | 61.09 | 39.17 | 23.40 | 41.22 | 22.43 |
| VIMA-large | 695M | 26.00 | 26.00 | 17.20 | 30.75 | 19.33 | 23.86 | 47.93 | 41.47 | 36.33 | 41.91 | 30.63 |
| Gato-large | 309M | 46.00 | 74.00 | 42.00 | 44.40 | 40.00 | 49.28 | 76.40 | 73.33 | 62.67 | 70.80 | 57.35 |
| Actra-large | 198M | 72.00 | 91.00 | 52.40 | 63.43 | 70.67 | 69.90 | 90.93 | 90.53 | 79.07 | 86.84 | 76.25 |

### 4.1 Experimental Setup

We compare various VLA models across seven skills in three different environments. For the “pick and place” skill, we gather a large-scale dataset in the Maniskill2 environment to evaluate generalization capabilities. In the Franka Kitchen environment , we assess the dexterity of the models across five skills: turn knob, toggle switch, slide door open, swing door open, lift by handle. Precision of actions is benchmarked using the Push-T environment . We compare with the following state-of-the-art baselines in robot manipulation tasks: Gato : The prompt, state, and action tokens are concatenated into a single sequence and a Transformer decoder serves as the policy network. RT-1 : We upgrade the USE prompt encoder to a T5 encoder and modify the token learner to accommodate images from multiple cameras. VIMA : We adapt its original $\mathrm{SE}(2)$ poses to the action dimensions of the our environment. Since our prompt is textual, we also use the same T5 encoder. Diffusion: We also developed two types of diffusion-based VLA models for Franka kitchen and Push-T. One model is based on DDPM for continuous actions, and the other is based on D3PM for discrete actions.

### 4.2 Implementation Details

We define a family of Actra models and utilize the most suitable model configuration for different environments. Actra-small (19.4M) consists of 6 layers, 8 attention heads, and an embedding size of 512. Actra-base (43.3M) comprises 6 layers, 12 attention heads, and an embedding size of 768. Actra-large (198M) is composed of 10 layers, 20 attention heads, and an embedding size of 1280. The baseline models adopt the same configurations unless specified otherwise. To encode the prompt, we utilize T5 encoder , whose pretrained weights are loaded and frozen during training. For the vision encoder, we experimented with both ResNet and ViT , trained from scratch. However, as we did not observe a performance gain when using ViT, we opted for ResNet as the vision encoder to minimize memory footprint. We use the AdamW optimizer with learning rate $=1\times 10^{-4}$ and weight decay $1\times 10^{-4}$.

**Table 2: Performance comparison (%) in Franka Kitchen and Push-T.**
|  | Franka Kitchen | Continuous Action | Discrete Action | Push-T Discrete Action |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Config | Params | $p\text{1}$ | $p\text{2}$ | $p\text{3}$ | $p\text{4}$ | Mean | $p\text{1}$ | $p\text{2}$ | $p\text{3}$ | $p\text{4}$ | Mean | Config | Params | Score |
| Diffusion | base | 43M | 100 | 94 | 72 | 34 | 76 | 54 | 18 | 10 | 2 | 21 | base | 43M | 83.89 |
| VIMA | base | 113M | 92 | 90 | 80 | 66 | 82 | 90 | 72 | 48 | 16 | 57 | small | 26M | 91.09 |
| Gato | base | 43M | 100 | 98 | 84 | 62 | 86 | 100 | 88 | 68 | 38 | 74 | small | 19M | 90.43 |
| Actra | base | 43M | 100 | 100 | 94 | 70 | 91 | 100 | 94 | 68 | 52 | 79 | small | 19M | 94.11 |

### 4.3 Performance Comparison in Maniskill

In the Maniskill environment , we evaluate one of the most commonly utilized skills: “pick and place”. Its goal is to pick up a target object and place it into a container. A language prompt specifies which target object and container are intended, with one prompt corresponding to one task. To test generalizability, we limit the training set to 15 tasks and evaluate the models on 34 tasks. All items are randomly placed and therefore, by default, our evaluation includes placement generalization . We conduct 50 trials for each task to minimize variance, and each trial is limited to 100 timesteps before a timeout. The models use the large configuration due to the large scale of our dataset.

We compare the success rates on various generalization levels, involving unseen target objects, containers, and distractors, as detailed in Table 1 & Figure 4. A seen task means that its language prompt is in the training set. For unseen tasks, we establish five main generalization levels. The first three levels (“Color”, “Size”, “Shape”) involve target objects with unseen colors, sizes, and shapes. E.g., small marker vs large marker (“Size”), and apple vs bowl (“Shape”). The fourth level (“Container”) introduces unseen containers. The fifth level (“Both”) involves both unseen target objects and unseen containers. This is not the simple average of the first four levels, as they either have an unseen target object or an unseen container, not both. A distractor is an item that is neither the target object nor the container. They are randomly sampled from a pool of diverse items. The robot and environment are randomly initialized. For all five “unseen tasks” levels, we sample 2-4 distractors and place them randomly. In seen tasks, we explore whether the number of distractors can impact the models’ performance (“Distractors”). Detailed analysis can be found in Appendix A.4.

### 4.4 Performance Comparison in Franka Kitchen

Franka Kitchen includes five skills that span seven specific tasks within the scene. The “turn knob” skill involves turning the oven knob to activate either the top or bottom burner. “Toggle switch” involves turning on the light switch. “Slide door open” requires opening the slide cabinet, while “swing door open” involves opening either the left hinge cabinet or the microwave door by the door handle. The “lift by handle” skill entails moving the kettle by its handle. Performance is measured by the completion of multi-stage tasks, as summarized in Table 2. Results are averaged over 50 runs. In each run, the models are required to complete four random tasks within 280 steps. If $p\text{i}=1$, it means the model has completed $i$ tasks and thus reached the $i$-th stage; otherwise, $p\text{i}=0$. Since the scale of Franka Kitchen is relatively small, we compare the models in their base configuration. We also compare the performance of models using continuous and discrete actions in this environment. Some examples by Actra in Franka Kitchen are shown in Appendix A.7.

### 4.5 Performance Comparison in Push-T

In Push-T , the models need to push a T-shaped object until it aligns perfectly with the target position. This task requires precise control, as performance is measured by the overlapping area, with perfect alignment equating to a score of 1.0. The performance of the models is compared in Table 2. Results are averaged over 30 trials. The maximum number of steps the models can take in each trial is 200, so they need to push the object precisely while maintaining adequate speed. Because this is a 2D task, we found that using small models is sufficient, except for the diffusion-based model, which uses the base configuration. Since the cross-attention layers in VIMA make its default small model amount to 50.8M parameters, which is even larger than Actra-base, we use three Transformer blocks instead. We found that models fail to learn effective policies using continuous actions; therefore, we only report results of discrete actions. Several push-T examples by Actra are included in Appendix A.7.

### 4.6 Ablation Study

Figure: Figure 4: Comparison of success rate on different generalization levels. “Unseen” for unseen tasks; “Seen” for seen tasks. “Distractors” involve at least two distractors, and “Basic” involves no distractors.
Refer to caption: x5.png

In our ablation study, we assess the novel components of Actra, as detailed in Table 3. We categorize the target objects of seen tasks into three difficulty levels: easy, medium, and hard. The details regarding the design of the difficulty levels are provided in Appendix A.5. VLA contrastive learning proves effective in enhancing primary robot imitation learning. We found that using VLA contrastive learning as a “warmup” phase works best for Actra. That is, we set $\alpha=0$, $\beta=1$ for the first few epochs, and then $\alpha=1$, $\beta=0$ for the remainder of the training. This schedule allows the model to absorb necessary knowledge about multi-modality encoding initially without interfering with the primary objective later in the training. In our ablation study, we analyze each of the three negative types. We observe that the removal of action negatives has the most substantial impact on performance degradation. Notably, hard tasks exhibit a heightened dependence on VLA contrastive learning compared to easier tasks.

We further investigate the impact of ablating trajectory attention and action query in Actra. The removal of trajectory attention results in a noticeable decrease in success rates across all difficulty levels, underscoring its crucial role in processing segmented VLA trajectories. Similarly, the absence of action query leads to reduced success rates, highlighting its significance in enhancing information extraction for action generation. Interestingly, when trajectory attention and action query are used in combination, the performance benefit exceeds the sum of their individual contributions, indicating a synergistic effect that enhances Actra’s precision in action generation.

### 4.7 Qualitative Analysis

Figure: Figure 5: The attention matrix of trajectory attention (left, showing two timesteps) and the visualization of two trained matrices (middle and right, showing five timesteps). Brighter cells correspond to higher attention weights. The tick labels are shown for the last token of every segment.
Refer to caption: x6.png

We present a visualization of the attention matrices from the top layer of Actra, depicted in Figure 5. The explanation of the attention matrix can be found in the appendix. In the visualization, prompt tokens exhibit similar attention values. Notably, tokens from more recent timesteps receive a higher attention weight compared to those from earlier in the sequence. This aligns with our expectation that the latest timestep is the most informative one for generating the next action. Some of the attention weights above the main diagonal are strongly activated, indicating the additional attention connections facilitated by our trajectory attention are beneficial. In the right matrix, a clear distinction is observed between the attention weights produced by state tokens and query tokens. This distinction underscores that action queries extract information differently from state tokens, elucidating their role in improving action generation.

Upon further analysis of the evaluation trajectories, a crucial distinction between Actra and the baselines emerges: Actra masters “instantaneous regrasp”. In most cases, baseline models struggle to recognize failed grasps and initiate another attempt promptly. Even if they identify a failed grasp, the time taken to start a new attempt is usually prolonged. In contrast, our Actra model promptly detects a failed grasp and continuously attempts to grasp the object until successful. This ability significantly reduces the failure rate with each regrasp, contributing to the substantially higher success rate observed in Actra. Figure 6 presents keyframes from an instantaneous regrasp corresponding to the most challenging task of “pick blue tea box and place into clear box”.

Figure: Figure 6: An example of instantaneous regrasp. Four grasp attempts were completed within only 30 steps, a capability not observed in the baseline models. An elaborate description of the example is provided in Appendix A.6.
Refer to caption: x9.png

## 5 Conclusion

This paper introduces Actra, an optimized Transformer architecture designed for VLA trajectories in robotic tasks. Actra distinguishes itself from vanilla Transformer decoders through two key components: trajectory attention and action query. Trajectory attention harnesses the unique characteristics of VLA trajectories, facilitating enhanced information flow among tokens within each segment. This allows Actra to encode the sequence at the segment level, and we introduce action queries to enable a segment-level decoding procedure. We incorporate an additional VLA contrastive learning task to explicitly train the model to encode multi-modal information, improving alignment between different modalities. This further elevates Actra’s performance in robot imitation learning. Through comprehensive comparisons in various environments, our approach demonstrates significant performance gain over state-of-the-art models in terms of generalizability, dexterity, and precision. Detailed ablation study and qualitative analysis further validate the effectiveness of Actra.

## 6 Limitations

Due to the high computational cost of diffusion-based VLA models, we were unable to compare against them in our large-scale Maniskill environment. Although more powerful language and vision encoders could be explored for even better performance, we kept them identical across Actra and the baselines to ensure a fair comparison. Although we have evaluated Actra in various setups, our evaluations were primarily conducted using the Franka embodiment. Exploring multi-embodiment would be an interesting direction for future research.

## Appendix A Appendix

### A.1 Trajectory Attention Matrix

The trajectory attention is implemented with its corresponding attention mask, as illustrated in Figure 7. Output tokens on the left attend to input tokens at the top. For example, in the first row, the token $p_{1}$ can attend to tokens $(p_{1},p_{2},p_{3},p_{4})$ and no future tokens, hence other tokens starting from $s_{1}$ are masked out. For action generation, action queries can attend to all tokens up to the current state, while no other tokens can attend to action queries. Therefore, action query tokens are all masked from the input.

Although we can use the decoding attention matrix for both decoding and encoding, employing the right matrix can save compute for the action queries. Regardless of using the left or right matrix, the resulting embeddings are identical for the encoded trajectory in VLA contrastive learning, thanks to our modified positional embeddings.

Figure: (a) Trajectory attention matrix for decoding.
Refer to caption: x10.png

### A.2 Positional Embeddings

We introduce two additional learnable embeddings, replacing the original GPT positional embedding []. The first is segment embedding, aiding the Transformer in distinguishing between prompt, state, query, and action segments, and is shared across different timesteps. The second is timestep embedding, enabling the Transformer to identify the timestep of each token and is identical for all tokens within the same timestep. Combining these two embeddings allows Actra to identify individual segments rather than individual tokens, as the order of tokens within a segment is irrelevant. Finally, the input embedding is the summation of all embeddings, following regular Transformers.

### A.3 Training Objectives of Diffusion-based VLA Models

The training objectives in diffusion-based VLA models [, ] can be written as:

$$ $\begin{split}\mathcal{L}_{\text{DDPM}}&=\operatorname{MSE}\left(\varepsilon^{k },\varepsilon_{\theta}\left(\mathbf{x}^{0}+\varepsilon^{k},k\right)\right)\\ \mathcal{L}_{\text{D3PM}}&=\operatorname{CE}\left(\varepsilon^{k},\varepsilon_ {\theta}\left(\mathbf{x}^{0}+\varepsilon^{k},k\right)\right)\end{split}$ (2) $$

where $\mathbf{x}^{0}$ is the original action and $\varepsilon^{k}$ is the noise of the $k$-th iteration; $\varepsilon_{\theta}$ is the VLA model.

### A.4 Generalization Levels

For Actra, Gato, and VIMA, unseen shape proves to be the most challenging level, followed by unseen containers. VIMA also exhibits volatility when dealing with small objects from the “Size” level. RT-1, in particular, struggles with identifying the container when an unseen container is introduced. It’s important to note that not all seen target objects have a generalized version for every generalization level. For example, a strawberry is a seen target object, but there is no over-sized strawberry for the “Shape” level. Consequently, models may achieve a higher success rate on some generalization levels than on seen tasks. Furthermore, since unseen color and size are part of the mixture, the success rate in “Both” is not as low as in “Shape” and “Container”. The introduction of more distractors in the scene increases the likelihood of collisions and causes additional difficulty in grasping the objects. However, this negative effect is not severe enough to considerably degrade the performance.

### A.5 Difficulty Levels

In simple terms, easy tasks involve spherical, regular-sized target objects, such as a baseball. The medium difficulty level includes elongated or small target objects, such as a banana or strawberry. Hard tasks encompass oversized, non-spherical, or thin objects, such as a tea box or knife.

Easy tasks include spherical, regular-sized objects. The reason why round objects are easier to pick is that the robot arm can close the gripper in any direction. Size also has a big impact on the success rate because over-sized objects require more precise grasp poses. If a grasp is not precise, the two fingers of the gripper might have collision with the object and not be able to reach down on the object. Small objects can increase the difficulty because the gripper might miss them if the grasp is slightly off. An object like a remote controller or a banana should be picked up “across” the object, not “along” the object. Thus, we define the medium difficulty level as the objects that are too big, too small, or elongated. Hard tasks involve non-spherical, oversized, or thin objects, such as a tea box and a knife. A tea box is non-spherical and oversized and thus the robot arm can only grasp it precisely in parallel with the sides, not diagonally. A knife can be hard since it is very thin and close to the desk. The gripper might collide with the desk while grasping.

### A.6 Instantaneous Regrasp

More explanation for the example in Figure 6. The first grasp was unsuccessful as one finger of the gripper collided with the blue tea box and the grasp slipped. Subsequently, Actra swiftly initiated two additional grasps; however, the gripper closed too early, resulting in collisions with the tea box again. Shortly after the second and third failures, the gripper’s fingers successfully reached down to opposite sides of the blue tea box, completing the fourth regrasp. Remarkably, all four grasp attempts were executed within a mere 30 timesteps, a feat not observed in the baselines.

### A.7 Franka Kitchen & Push-T Examples

We provide some execution examples by Actra in Franka Kitchen (Figure 8) and Push-T (Figure 9).

Figure: Figure 8: Actra in Franka Kitchen. The model completes four random tasks. The top left image shows the initial state and the other three images are three different final states.
Refer to caption: x12.png

Figure: Figure 9: Actra in Push-T. The grey T is the object and the green T is its target position. The model controls the blue dot to push the T-shaped object towards the target position. The red cross is the cursor. The images on the left and right are two different initial states, and the middle image is the final state where the object perfectly overlaps with the target position.
Refer to caption: x13.png

### A.8 Additional Related Work

MOO [] introduced multi-modal prompt capability to RT-1, while Q-Transformer [] adapted RT-1 to the Q-learning setting. RoboFlamingo [] constructed a VLA based on the existing Flamingo VLM [, ]. ACT [] adopts the DETR framework for robotics tasks but utilizes fixed position embeddings at the timestep level. Another category of VLAs focuses on building high-level planners for long-horizon robotics tasks and abstracting away the low-level control policies, such as SayCan [], PaLM-E [], and ChatGPT for Robotics [].

In addition to the primary learning objective, auxiliary or pretraining objectives have proven useful in further enhancing model performance. The success of masked language modeling, as initially proposed in BERT [], has prompted the adoption of similar objectives in various domains. In computer vision models and VLMs, representative works like MAE [] and ViLBERT [] have employed comparable strategies. VLAs have also utilized masked modeling objectives for their vision encoders, such as MVP [], Voltron [], GR-1 []. While these approaches have proven beneficial for the vision encoder, they often overlook the crucial alignment between different modalities.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: - Could you elaborate on the low performance of Actra on CALVIN compared to recent sota methods? What factors do you think contributes to these differences and how would you bridge them in future work? 

- The training duration of 5-10 epochs shows promising initial results. Have you explored longer training runs to understand the full convergence behavior and potential performance ceiling of the approach? 

- How are the contributions of trajectory attention and learnable action queries different to prior work that employed the same already?
 
- Could you consider including additional scaling experiments to demonstrate how this approach performs across different model capacities and dataset sizes? This could help better understand its broader applicability as a VLM policy.
soundness: 2
strengths: - novel dynamics self-supervised loss to improve policy learning performance, which does not require pretraining and improves performance

- proposed architecture changes improve upon a wide range of transformer baselines in different benchmarks

- the paper is overall well written and easy to follow with several illustrative figures
summary: The paper proposes optimized Transformer Architectures for Policy Learning. The authors propose 3 contributions: a custom self-attention mask for policy learning, learnable action queries for each action dimension and an contrastive loss function to improve policy learning by having the model differentiate between plausible dynamics and unrealistic ones. The proposed architecture is tested on three common IL benchmarks including CALVIN and VIMA with several ablations.
weaknesses: - The terminology "VLA" is misleading since the model lacks vision-language pretraining. This is fundamentally a standard Transformer-based policy without language generation or image understanding capabilities.

- Performance issues:

    - CALVIN results are significantly below SOTA (1.x vs 3-4 average rollout length from methods like GR-1 or 3DDA), with existing methods like GR-1 already implementing similar trajectory attention mechanisms and learnable action queries
    - VIMA improvements are marginal, with lower performance compared to recent approaches like [4] that also utilize per-dimension action tokens
    - Low Diffusion baseline performance on Relay Kitchen and Push-T compared to originally reported metrics - what explains that discrepancy?
- Limited novelty in technical contributions:
    - Trajectory attention is an incremental modification seen in multiple prior works (Octo [1], GR-1 [2])
    - Learnable action queries are also incremental and were previously explored in GR-1, ACT [3] or [4] and similar research
    - No substantial architectural innovations except the self-supervised loss

- Methodological limitations:
    - Absence of scaling analysis across model sizes and environments
    - No large-scale experimental validation
    - Training duration of 5-10 epochs is insufficient to demonstrate convergence
    - Lacks comprehensive ablation studies on larger-scale settings
    - Design changes of generalist policies need evaluation on large-scale datasets to see if the proposed changes scale as intended


Sources:

[1]: Team, Octo Model, et al. "Octo: An open-source generalist robot policy." RSS 2024

[2]: Wu, Hongtao, et al. "Unleashing large-scale video generative pre-training for visual robot manipulation." ICLR 2024

[3]: Zhao, Tony Z., et al. "Learning fine-grained bimanual manipulation with low-cost hardware." RSS 2023

[4]: Li, Jiachen, et al. "Mastering robot manipulation with multimodal prompts through pretraining and multi-task fine-tuning." _arXiv preprint arXiv:2310.09676_ (2023).

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: How would the proposed attention mechanism affect the inference speed? 

What are the failure modes of the method? Providing additional analysis would help readers better understand the work.
soundness: 2
strengths: The method achieves good performance from the evaluation perspective; 

The paper is well structured and easy to follow.
summary: The paper studies the problem of learning VLA models for robot manipulation tasks. The proposed method optimizes the architecture of classical transformer models by introducing trajectory attention. The authors also propose to use learnable action queries to make the action decoding more focused on relevant information. The method is validated on several simulation benchmarks.
weaknesses: The design is not sufficiently evaluated through experiments.  

The comparison with baselines is insufficient, more prior works need to be included and compared, e.g., Behavior Transformers[1], CoTPC[2]. 

[1] Behavior Transformers: Cloning k modes with one stone, NeurIPS’22. 

[2] Chain-of-Thought Predictive Control, ICML’24.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Claim in abstract: “Additionally, the autoregressive generation approach falls short in generating multi-dimensional actions.” Can you elaborate why it AR has fall short in action generation? 
2. How many tokens are there to represent action at each timestep?
3. How does the sample complexity of the model compare to diffusion policy [1] (for a single task setup, such as that in table 5)? 
4. What’s the sequence length of the model? After adding CDL, how much slowdown is present for training? 

[1] Chi, Cheng, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. "Diffusion policy: Visuomotor policy learning via action diffusion." The International Journal of Robotics Research (2023): 02783649241273668.
soundness: 3
strengths: 1. The paper explores an interesting direction: whether the autoregressive mask pattern makes sense for VLAs. While it is less well supported by efficient inference methods (i.e., flash-attention, ring-attention, etc.), the reviewer believes that exploring such direction is beneficial for the community. 
2. Thorough ablation studies. It helps the reader understand how the change in attention pattern affects the model's performance.
summary: Vision language action models typically use causal attention, which may be suboptimal for processing multimodal sequences. The paper makes two contributions: 1) analysis of how the attention pattern of VLAs can be modified to improve performance, and 2) sequence-level contrastive learning helps. Simulation experiments on VIMABench, Maniskill2, and CALVIN show that the method outperforms AR models and shows the effectiveness of contrastive learning.
weaknesses: 1. Certain claims in the abstract or introduction may not be correct and need to be elaborated.

“Additionally, the autoregressive generation approach falls short in generating multi-dimensional actions” (Line 15-16)
Many existing VLAs (such as Octo [1] and OpenVLA [2]) are AR models. Many behavior cloning models [4] are also AR models. They all generate multi-dimensional actions and perform real robot tasks.  
 
“VLAs predominantly build upon the pioneering foundations laid by Decision Transformer and Trajectory Transformer … This paradigm has become a cornerstone across recent VLAs” (line 46-49). 
These two formulations are drastically different structures than RT-2, OpenVLA, etc, where there’s no interleaving of vision, action, and proprioception data and do not model long sequences. Instead, many recent approaches have gone towards VQA models such as LLaVA / BLIP / Transfusion. 

2. The paper can benefit from comparisons with the current state-of-the-art VLAs (Octo [1] and OpenVLA [2]). Additionally, many recent works such as Implicit behavior cloning [5] and diffusion policy [6] can work with <50 demonstrations. Multi-task diffusion-based VLA (Octo, scaling up and distilling down [7]) also shows promise in few-shot learning. It is not obvious that the proposed method is more sample-efficient then these methods.

3. Lack of ablations and experiments on a real robotics setup. 4. The paper does not report confidence intervals (standard deviation or standard error) on the results. It is hard to conclude from the reported results to claim that contrastive learning helps Actra. 
5. Scalability: many of the baselines (i.e. RT-1, Gato, VIMA) can benefit significantly from pre-training (in fact most of the evaluations for these models are done post large-scale pre-training). It is not obvious if this method is as scalable (both in terms of model sizes and dataset sizes) compared to the baselines. 
6. Past works have explored sequence-level contrastive learning (i.e. ConDT [3]). The paper can benefit from including more ablation studies over different types of sequence-level contrastive learning frameworks. 

[1] Team, Octo Model, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari et al. "Octo: An open-source generalist robot policy." arXiv preprint arXiv:2405.12213 (2024).

[2] Kim, Moo Jin, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov et al. "OpenVLA: An Open-Source Vision-Language-Action Model." arXiv preprint arXiv:2406.09246 (2024).

[3] Konan, Sachin G., Esmaeil Seraj, and Matthew Gombolay. "Contrastive decision transformers." In Conference on Robot Learning, pp. 2159-2169. PMLR, 2023.

[4] Radosavovic, Ilija, Bike Zhang, Baifeng Shi, Jathushan Rajasegaran, Sarthak Kamat, Trevor Darrell, Koushil Sreenath, and Jitendra Malik. "Humanoid locomotion as next token prediction." arXiv preprint arXiv:2402.19469 (2024).

[5] Florence, Pete, Corey Lynch, Andy Zeng, Oscar A. Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. "Implicit behavioral cloning." In Conference on Robot Learning, pp. 158-168. PMLR, 2022.

[6] Chi, Cheng, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. "Diffusion policy: Visuomotor policy learning via action diffusion." The International Journal of Robotics Research (2023): 02783649241273668.

[7] Ha, Huy, Pete Florence, and Shuran Song. "Scaling up and distilling down: Language-guided robot skill acquisition." In Conference on Robot Learning, pp. 3766-3777. PMLR, 2023.