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

Yueen Ma , 1 Dafeng Chi , 2 Shiguang Wu , 2 Yuecheng Liu , 2 Yuzheng Zhuang , 2 Jianye Hao 2 , Irwin King 1 The Chinese University of Hong Kong 1 Huawei Noah's Ark Lab 2 {yema21, king}@cse.cuhk.edu.hk {chidafeng1, wushiguang, liuyuecheng1, zhuangyuzheng, haojianye}

@huawei.com

## 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 [8] 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 [5], thanks to their dexterity and generalizability.

To accommodate multi-modal inputs, previous Transformer-based VLMs [42] explored designing special self-attention schemes to better suit the unique properties of different modalities, such as UniLM [11], M6 [28], BLIP-2 [26] and Octo [33]. 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.

VLAs predominantly build upon the pioneering foundations laid by Decision Transformer [9] and Trajectory Transformer [22]. These two works frame reinforcement learning (RL) policies as sequence modeling problems, leveraging the expressive power of Transformer models. This paradigm has

Preprint. Under review.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-3ef897ec-7283-4360-8c59-b5d6d4f4eb02.jpg)

a₁

a₂

a₃

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-9112d033-1b48-47fe-a68e-a2ce9b695a4b.jpg)

(a) Causal Attention.

(b) Trajectory Attention.

Figure 1: A comparison between causal attention and trajectory attention is illustrated for an action, comprising three action dimensions, referred to as a 'segment'. Each orange dot represents a token embedding, and a line corresponds to a valid attention connection from the input embedding (bottom) to the output embedding (top). In trajectory attention, tokens can attend to not only the preceding tokens but also the subsequent tokens within the same segment, as indicated by the green lines.

a₁ a₂ a₃ become a cornerstone across recent VLAs, but both models are Transformer decoders based on causal attention. Subsequent approaches, such as Gato [38] and RT-1 [5], also adopt Transformer decoders as their network backbone, passing in different modalities as a single sequence. VIMA [24] 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 &amp; 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 [6], 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 [4]. By combining trajectory attention and action queries, we introduce an optimized Transformer architecture for VLA trajectories, which we name Act ion-query-based Tra jectory-attention Tra nsformer, 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 [34]. 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 [32] proposed a time contrastive learning objective, and VIP [31] 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

Figure 2: The architecture of Actra. Each square represents a token. A trajectory τ 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).

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-ec64c356-5713-408f-9763-aab37d8c2789.jpg)

and negative trajectories and contrast them with the widely adopted InfoNCE loss [41]. 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 multimodality 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.

## 2 Related Work

Vision-language-action Model. Vision-language-action models (VLAs) constitute a class of multimodal models designed to generate actions based on specified language prompts and perceived environment. Coined by RT-2 [4], VLAs have garnered increasing attention due to their dexterity and generalizability in handling complex robotics tasks. Early attempts were based on existing visionlanguage models (VLMs), exemplified by CLIPort [39] and BC-Z [21]. Gato [38] explored the use of a single Transformer model [42] as the control policy for tasks spanning various domains, unifying multi-modal inputs into a single sequence. RT-1 [5] 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 [11], M6 [28], and Octo [33] 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 [38] and RT-1 [5] maintain causal attention in Transformer decoders. VIMA [24] 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 [6], learnable object queries have shown promising results in extracting information for object detection. BLIP-2 [26] 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 [34], ALIGN [23], Florence [46], FILIP [45], has demonstrated the significance of contrastive learning in enhancing multi-modal interaction. However, in VLAs like R3M [32] and VIP [31], 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 τ = ( p, s t =1 , a t =1 , . . . , s t = T , a t = T ) . Each element in the trajectoryp , 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 , . . . , 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 τ = ( p 1: L , s 1: M,t =1 , a 1: N,t =1 , . . . , 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 π θ ( a t | p, s ≤ t , a &lt;t ) .

## 3.2 Actra

In natural language generation (NLG), language models such as GPT [35] 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 [9, 5] 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 [10]. 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 languageconditioned 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 ≥ 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 , p i j ) , ( s i,t , s j,t ) or ( a i,t , a j,t ) for i &lt; j . When compared with the original causal attention, there are L L ( -1) / 2+ T ( M M ( -1) / 2+ N N ( -1) / 2 ) 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

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.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-03958a67-2dd1-43da-9afe-830777974219.jpg)

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 [4]. 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 [6]. 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 = · · · = q i,t = T for i ∈ { 1 . . . 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.

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 , . . . , s t , a t ) , and all state tokens s i,t can fully attend to ( p, s t =1 , a t =1 , . . . , 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 L BC = min θ ∑ T t =1 -log π θ ( a t | p, s ≤ t , a &lt;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 [34, 23, 46, 45]. Therefore, we propose a novel VLA contrastive learning objective to

|              |        |       | Unseen tasks   | Unseen tasks   | Unseen tasks   | Unseen tasks   | Distractors   |   Distractors | Distractors   | Distractors   | Distractors   | 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.4  | 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         |

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.

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 τ = ( p, s t =1 , a t =1 , . . . , s t = T , a t = T ) . To construct a positive, τ + , 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 ′ from the dataset and substituting it for the correct prompt. The resulting negative trajectory for the language modality is τ -p = ( p , s ′ t =1 , a t =1 , . . . , 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 τ ′ = ( p , s ′ ′ t =1 , a ′ t =1 , . . . , s ′ t = T , a ′ t = T ) in the same batch as τ , we can mismatch the states and actions to assemble negatives for the vision and action modalities: τ -s = ( p, s ′ t =1 , a t =1 , . . . , s ′ t = T , a t = T ) and τ -a = ( p, s t =1 , a ′ t =1 , . . . , s t = T , a ′ t = T ) . Finally, we contrast the positive τ + against three types of negatives: τ -= { τ -p , τ -s , τ -a } .

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

$$\bar { \mathcal { L } } _ { \text{VLA-CL} } ( \tau, \tau ^ { + }, \tau ^ { - } ) = - \log \mathbb { E } \left [ \frac { \gamma } { \exp ( f ( \tau ) ^ { \intercal } f ( \tau ^ { + } ) ) } } { B } \sum _ { i = 1 } ^ { B } \sum _ { j \in \{ p, s, a \} } \exp ( f ( \tau ) ^ { \intercal } f ( \tau _ { i, j } ^ { - } ) ) } \right ], \quad ( 1 )$$

where f ( ) · 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: L = α L BC + β L VLA-CL.

## 4 Experimental Results

## 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 [15] to evaluate generalization capabilities. In the Franka Kitchen environment [16], 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 [14]. We compare with the following state-of-the-art baselines in robot manipulation tasks: Gato [38]: The prompt, state, and action tokens are concatenated into a single sequence and a Transformer decoder serves as the policy network. RT-1 [5]: We upgrade the USE prompt encoder to a T5 encoder [37] and modify the token learner to accommodate images from multiple cameras. VIMA [24]: We adapt its original 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 [19] for continuous actions, and the other is based on D3PM [2] for discrete actions.

## 4.2 Implementation Details

We define a family of Actra models and utilize the most suitable model configuration for different environments. Actrasmall (19.4M) consists of 6 layers, 8 attention heads, and an embedding size of 512. Actrabase (43.3M) comprises 6 layers, 12 attention heads, and an embedding size of 768. Actralarge (198M) is composed of 10 layers, 20 attention heads, and an embedding size of 1280.

Table 2: Performance comparison (%) in Franka Kitchen and Push-T.

|           | Franka Kitchen   | Franka Kitchen   | Continuous Action   | Continuous Action   | Continuous Action   | Continuous Action   | Continuous Action   | Discrete Action   | Discrete Action   | Discrete Action   | Discrete Action   | Discrete Action   | Push-T Discrete Action   | Push-T Discrete Action   | Push-T Discrete Action   |
|-----------|------------------|------------------|---------------------|---------------------|---------------------|---------------------|---------------------|-------------------|-------------------|-------------------|-------------------|-------------------|--------------------------|--------------------------|--------------------------|
| Model     | Config           | Params           | p 1                 | p 2                 | p 3                 | p 4                 | Mean                | p 1               | p 2               | p 3               | p 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                    |

The baseline models adopt the same configurations unless specified otherwise. To encode the prompt, we utilize T5 encoder [37], whose pretrained weights are loaded and frozen during training. For the vision encoder, we experimented with both ResNet [18] and ViT [12], 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 [29] with learning rate = 1 × 10 -4 and weight decay 1 × 10 -4 .

## 4.3 Performance Comparison in Maniskill

In the Maniskill environment [15], 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 [24]. 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 &amp; 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 [16] 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 i = 1 , it means the model has completed i tasks and thus reached the i -th stage; otherwise, p 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 [14], 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

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.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-3b2e66d8-b3a6-44fc-896d-c303180c483f.jpg)

| Model           | Overall      | Hard         | Medium       | Easy         |
|-----------------|--------------|--------------|--------------|--------------|
| Actra- large    | 90.5         | 83.5         | 89.7         | 93.8         |
| w/o V-negatives | 89.2 (-1.3)  | 79.0 (-4.5)  | 88.3 (-1.4)  | 93.7 (-0.1)  |
| w/o L-negatives | 89.5 (-1.0)  | 81.0 (-2.5)  | 88.6 (-1.1)  | 93.3 (-0.5)  |
| w/o A-negatives | 88.0 (-2.5)  | 78.0 (-5.5)  | 86.3 (-3.4)  | 93.3 (-0.5)  |
| w/o VLA CL      | 87.9 (-2.6)  | 75.0 (-8.5)  | 88.6 (-1.1)  | 91.3 (-2.5)  |
| w/o Traj Attn   | 83.7 (-6.8)  | 76.0 (-7.5)  | 83.7 (-6.0)  | 86.3 (-7.5)  |
| w/o Act Query   | 81.3 (-9.2)  | 69.0 (-14.5) | 80.9 (-8.8)  | 86.0 (-7.8)  |
| w/o Both        | 79.5 (-11.0) | 68.0 (-15.5) | 79.4 (-10.3) | 83.3 (-10.5) |

Table 3: Ablation study of the proposed components in Actra. We ablate each type of negative trajectory for VLA contrastive learning. Starting from the row 'w/o VLA CL', we remove VLA contrastive learning and only utilize the BC objective. We test them on seen tasks with 2-4 distractors. The values in the brackets represent the difference from the complete 'Actralarge '.

small model amount to 50.8M parameters, which is even larger than Actrabase , 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

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 α = 0 , β = 1 for the first few epochs, and then α = 1 , β = 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

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,

P₁ P₂ P₃ P₄ S₁ S₂ A₁ A₂ A₃ S₁ S₂ Q₁ Q₂ Q₃ A₁ A₂ A₃ 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.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-babaad5e-9b52-408a-8ff1-f58ad4226fd6.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-ec7c0e1a-0d49-4e62-81fb-0e5a72ef985d.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-b32df852-bff6-4860-a929-c3dfe28374b8.jpg)

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.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-b52082c2-fef3-42e4-a28a-d07996f0b040.jpg)

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'.

## 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.

## References

- [1] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, Roman Ring, Eliza Rutherford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooei, Marianne Monteiro, Jacob L. Menick, Sebastian Borgeaud, Andy Brock, Aida Nematzadeh, Sahand Sharifzadeh, Mikolaj Binkowski, Ricardo Barreira, Oriol Vinyals, Andrew Zisserman, and Karén Simonyan. Flamingo: a visual language model for few-shot learning. In NeurIPS, 2022.
- [2] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 17981-17993, 2021.
- [3] Anas Awadalla, Irena Gao, Josh Gardner, Jack Hessel, Yusuf Hanafy, Wanrong Zhu, Kalyani Marathe, Yonatan Bitton, Samir Yitzhak Gadre, Shiori Sagawa, Jenia Jitsev, Simon Kornblith, Pang Wei Koh, Gabriel Ilharco, Mitchell Wortsman, and Ludwig Schmidt. Openflamingo: An open-source framework for training large autoregressive vision-language models. CoRR, abs/2308.01390, 2023.
- [4] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choromanski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, Pete Florence, Chuyuan Fu, Montse Gonzalez Arenas, Keerthana Gopalakrishnan, Kehang Han, Karol Hausman, Alexander Herzog, Jasmine Hsu, Brian Ichter, Alex Irpan, Nikhil J. Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Lisa Lee, Tsang-Wei Edward Lee, Sergey Levine, Yao Lu, Henryk Michalewski, Igor Mordatch, Karl Pertsch, Kanishka Rao, Krista Reymann, Michael S. Ryoo, Grecia Salazar, Pannag Sanketi, Pierre Sermanet, Jaspiar Singh, Anikait Singh, Radu Soricut, Huong T. Tran, Vincent Vanhoucke, Quan Vuong, Ayzaan Wahid, Stefan Welker, Paul Wohlhart, Jialin Wu, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich. RT-2: vision-language-action models transfer web knowledge to robotic control. CoRR, abs/2307.15818, 2023.
- [5] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alexander Herzog, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Tomas Jackson, Sally Jesmonth, Nikhil J. Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Kuang-Huei Lee, Sergey Levine, Yao Lu, Utsav Malla, Deeksha Manjunath, Igor Mordatch, Ofir Nachum, Carolina Parada, Jodilyn Peralta, Emily Perez, Karl Pertsch, Jornell Quiambao, Kanishka Rao, Michael S. Ryoo, Grecia Salazar, Pannag R. Sanketi, Kevin Sayed, Jaspiar Singh, Sumedh Sontakke, Austin Stone, Clayton Tan, Huong T. Tran, Vincent Vanhoucke, Steve Vega, Quan Vuong, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich. RT-1: robotics transformer for real-world control at scale. In Robotics: Science and Systems, 2023.
- [6] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In ECCV (1), volume 12346 of Lecture Notes in Computer Science, pages 213-229. Springer, 2020.
- [7] Yevgen Chebotar, Quan Vuong, Alex Irpan, Karol Hausman, Fei Xia, Yao Lu, Aviral Kumar, Tianhe Yu, Alexander Herzog, Karl Pertsch, Keerthana Gopalakrishnan, Julian Ibarz, Ofir Nachum, Sumedh Sontakke, Grecia Salazar, Huong T. Tran, Jodilyn Peralta, Clayton Tan, Deeksha Manjunath, Jaspiar Singh, Brianna Zitkovich, Tomas Jackson, Kanishka Rao, Chelsea Finn, and Sergey Levine. Q-transformer: Scalable offline reinforcement learning via autoregressive q-functions. CoRR, abs/2309.10150, 2023.
- [8] Feilong Chen, Duzhen Zhang, Minglun Han, Xiuyi Chen, Jing Shi, Shuang Xu, and Bo Xu. VLP: A survey on vision-language pre-training. Int. J. Autom. Comput., 20(1):38-56, 2023.
- [9] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In NeurIPS, pages 15084-15097, 2021.
- [10] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT (1), pages 4171-4186. Association for Computational Linguistics, 2019.

| [11]   | Li Dong, Nan Yang, Wenhui Wang, Furu Wei, Xiaodong Liu, Yu Wang, Jianfeng Gao, Ming Zhou, and Hsiao-Wuen Hon. Unified language model pre-training for natural language under- standing and generation. In NeurIPS, pages 13042-13054, 2019.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 |
|--------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [12]   | Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR. OpenReview.net, 2021.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 |
| [13]   | Danny Driess, Fei Xia, Mehdi S. M. Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, Wenlong Huang, Yevgen Chebotar, Pierre Sermanet, Daniel Duckworth, Sergey Levine, Vincent Vanhoucke, Karol Hausman, Marc Toussaint, Klaus Greff, Andy Zeng, Igor Mordatch, and Pete Florence. Palm-e: An embodied multimodal language model. In ICML, volume 202 of Proceedings of Machine Learning Research, pages 8469-8488. PMLR, 2023.                                                                                                                                                                                                                                                                                                                                                               |
| [14]   | Pete Florence, Corey Lynch, Andy Zeng, Oscar A. Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Aleksandra Faust, David Hsu, and Gerhard Neumann, editors, Conference on Robot Learning, 8-11 November 2021, London, UK, volume 164 of Proceedings of Machine Learning Research, pages 158-168. PMLR, 2021.                                                                                                                                                                                                                                                                                                                                                                                                                                                               |
| [15]   | Jiayuan Gu, Fanbo Xiang, Xuanlin Li, Zhan Ling, Xiqiang Liu, Tongzhou Mu, Yihe Tang, Stone Tao, Xinyue Wei, Yunchao Yao, Xiaodi Yuan, Pengwei Xie, Zhiao Huang, Rui Chen, and Hao Su. Maniskill2: A unified benchmark for generalizable manipulation skills. In International Conference on Learning Representations, 2023.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 |
| [16]   | Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura, editors, 3rd Annual Conference on Robot Learning, CoRL 2019, Osaka, Japan, October 30 - November 1, 2019, Proceedings, volume 100 of Proceedings of Machine Learning Research, pages 1025-1037. PMLR, 2019.                                                                                                                                                                                                                                                                                                                                                                                                            |
| [17]   | Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In CVPR, pages 15979-15988. IEEE, 2022.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                             |
| [18]   | Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, pages 770-778. IEEE Computer Society, 2016.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
| [19]   | Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
| [20]   | Brian Ichter, Anthony Brohan, Yevgen Chebotar, Chelsea Finn, Karol Hausman, Alexander Herzog, Daniel Ho, Julian Ibarz, Alex Irpan, Eric Jang, Ryan Julian, Dmitry Kalashnikov, Sergey Levine, Yao Lu, Carolina Parada, Kanishka Rao, Pierre Sermanet, Alexander Toshev, Vincent Vanhoucke, Fei Xia, Ted Xiao, Peng Xu, Mengyuan Yan, Noah Brown, Michael Ahn, Omar Cortes, Nicolas Sievers, Clayton Tan, Sichun Xu, Diego Reyes, Jarek Rettinghouse, Jornell Quiambao, Peter Pastor, Linda Luu, Kuang-Huei Lee, Yuheng Kuang, Sally Jesmonth, Nikhil J. Joshi, Kyle Jeffrey, Rosario Jauregui Ruano, Jasmine Hsu, Keerthana Gopalakrishnan, Byron David, Andy Zeng, and Chuyuan Kelly Fu. Do as I can, not as I say: Grounding language in robotic affordances. In CoRL, volume 205 of Proceedings of Machine Learning Research, pages 287-318. PMLR, 2022. |
| [21]   | Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. BC-Z: zero-shot task generalization with robotic imitation learning. In CoRL, volume 164 of Proceedings of Machine Learning Research, pages 991-1002. PMLR, 2021.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
| [22]   | Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In NeurIPS, pages 1273-1286, 2021.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
| [23]   | Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc V. Le, Yun- Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. In ICML, volume 139 of Proceedings of Machine Learning Research, pages 4904-4916. PMLR, 2021.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |

| [24]   | Yunfan Jiang, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, and Linxi Fan. VIMA: general robot manipulation with multimodal prompts. CoRR, abs/2210.03094, 2022.                                                                                                                                                                                               |
|--------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [25]   | Siddharth Karamcheti, Suraj Nair, Annie S. Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang. Language-driven representation learning for robotics. In Robotics: Science and Systems, 2023.                                                                                                                                                                                                                    |
| [26]   | Junnan Li, Dongxu Li, Silvio Savarese, and Steven C. H. Hoi. BLIP-2: bootstrapping language- image pre-training with frozen image encoders and large language models. In ICML, volume 202 of Proceedings of Machine Learning Research, pages 19730-19742. PMLR, 2023.                                                                                                                                                         |
| [27]   | Xinghang Li, Minghuan Liu, Hanbo Zhang, Cunjun Yu, Jie Xu, Hongtao Wu, Chilam Cheang, Ya Jing, Weinan Zhang, Huaping Liu, Hang Li, and Tao Kong. Vision-language foundation models as effective robot imitators. CoRR, abs/2311.01378, 2023.                                                                                                                                                                                  |
| [28]   | Junyang Lin, Rui Men, An Yang, Chang Zhou, Ming Ding, Yichang Zhang, Peng Wang, Ang Wang, Le Jiang, Xianyan Jia, Jie Zhang, Jianwei Zhang, Xu Zou, Zhikang Li, Xiaodong Deng, Jie Liu, Jinbao Xue, Huiling Zhou, Jianxin Ma, Jin Yu, Yong Li, Wei Lin, Jingren Zhou, Jie Tang, and Hongxia Yang. M6: A chinese multimodal pretrainer. CoRR, abs/2103.00823, 2021.                                                             |
| [29]   | Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR (Poster). OpenReview.net, 2019.                                                                                                                                                                                                                                                                                                              |
| [30]   | Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task-agnostic visiolinguistic representations for vision-and-language tasks. In NeurIPS, pages 13-23, 2019.                                                                                                                                                                                                                                         |
| [31]   | Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. VIP: towards universal visual reward and representation via value-implicit pre- training. In ICLR. OpenReview.net, 2023.                                                                                                                                                                                                     |
| [32]   | Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: A universal visual representation for robot manipulation. In CoRL, volume 205 of Proceedings of Machine Learning Research, pages 892-909. PMLR, 2022.                                                                                                                                                                                     |
| [33]   | Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Charles Xu, Jianlan Luo, Tobias Kreiman, You Liang Tan, Dorsa Sadigh, Chelsea Finn, and Sergey Levine. Octo: An open-source generalist robot policy. https: //octo-models.github.io , 2023.                                                                                                                       |
| [34]   | Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In ICML, volume 139 of Proceedings of Machine Learning Research, pages 8748-8763. PMLR, 2021.                                                              |
| [35]   | Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.                                                                                                                                                                                                                                                       |
| [36]   | Ilija Radosavovic, Tete Xiao, Stephen James, Pieter Abbeel, Jitendra Malik, and Trevor Darrell. Real-world robot learning with masked visual pre-training. In CoRL, volume 205 of Proceedings of Machine Learning Research, pages 416-426. PMLR, 2022.                                                                                                                                                                        |
| [37]   | Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21:140:1-140:67, 2020.                                                                                                                                                                 |
| [38]   | Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-maron, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. Transactions on Machine Learning Research, 2022. Featured Certification. |
| [39]   | Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Cliport: What and where pathways for robotic manipulation. In CoRL, volume 164 of Proceedings of Machine Learning Research, pages 894-906. PMLR, 2021.                                                                                                                                                                                                                        |
| [40]   | Austin Stone, Ted Xiao, Yao Lu, Keerthana Gopalakrishnan, Kuang-Huei Lee, Quan Vuong, Paul Wohlhart, Brianna Zitkovich, Fei Xia, Chelsea Finn, and Karol Hausman. Open-world object manipulation using pre-trained vision-language models. CoRR, abs/2303.00905, 2023.                                                                                                                                                        |

| [41]   | Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, abs/1807.03748, 2018.                                                                                                                                                                                                                                              |
|--------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [42]   | Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, pages 5998-6008, 2017.                                                                                                                                                                                              |
| [43]   | Sai Vemprala, Rogerio Bonatti, Arthur Bucker, and Ashish Kapoor. Chatgpt for robotics: Design principles and model abilities. CoRR, abs/2306.17582, 2023.                                                                                                                                                                                                                             |
| [44]   | Hongtao Wu, Ya Jing, Chilam Cheang, Guangzeng Chen, Jiafeng Xu, Xinghang Li, Minghuan Liu, Hang Li, and Tao Kong. Unleashing large-scale video generative pre-training for visual robot manipulation, 2023.                                                                                                                                                                           |
| [45]   | Lewei Yao, Runhui Huang, Lu Hou, Guansong Lu, Minzhe Niu, Hang Xu, Xiaodan Liang, Zhenguo Li, Xin Jiang, and Chunjing Xu. FILIP: fine-grained interactive language-image pre-training. In ICLR. OpenReview.net, 2022.                                                                                                                                                                 |
| [46]   | Lu Yuan, Dongdong Chen, Yi-Ling Chen, Noel Codella, Xiyang Dai, Jianfeng Gao, Houdong Hu, Xuedong Huang, Boxin Li, Chunyuan Li, Ce Liu, Mengchen Liu, Zicheng Liu, Yumao Lu, Yu Shi, Lijuan Wang, Jianfeng Wang, Bin Xiao, Zhen Xiao, Jianwei Yang, Michael Zeng, Luowei Zhou, and Pengchuan Zhang. Florence: A new foundation model for computer vision. CoRR, abs/2111.11432, 2021. |
| [47]   | Tony Z. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. In Robotics: Science and Systems, 2023.                                                                                                                                                                                                              |

## 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 , p 1 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.

(a) Trajectory attention matrix for decoding.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-17627642-fb7c-43f4-abe2-8c5633e090d2.jpg)

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-183ebae0-771e-471b-88d4-b349f8216359.jpg)

P₁ P₂ P₃ P₄ S₁ S₂ Q₁ Q₂ Q₃ A₁ A₂ A₃ S₁ S₂ Q₁ Q₂ (b) Trajectory attention matrix for encoding.

P₁ P₂ P₃ P₄ S₁ S₂ A₁ A₂ A₃ S₁ S₂ Q₁ Q₂ Q₃ A₁ A₂ A₃ Figure 7: The attention matrix of Trajectory Attention. Dark boxes represent masked entries in the attention matrix. The left attention matrix is used for decoding during action generation while the right attention matrix is used for encoding the trajectory in VLA contrastive learning.

## A.2 Positional Embeddings

We introduce two additional learnable embeddings, replacing the original GPT positional embedding [35]. 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 [19, 2] can be written as:

$$\mathcal { L } _ { D D P M } & = \text{MSE} \left ( \varepsilon ^ { k }, \varepsilon _ { \theta } \left ( \mathbf x ^ { 0 } + \varepsilon ^ { k }, k \right ) \right ) \\ \mathcal { L } _ { D 3 P M } & = \text{CE} \left ( \varepsilon ^ { k }, \varepsilon _ { \theta } \left ( \mathbf x ^ { 0 } + \varepsilon ^ { k }, k \right ) \right )$$

Q₃

A₁

A₂

A₃

where x 0 is the original action and ε k is the noise of the k -th iteration; ε θ 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 &amp; Push-T Examples

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

## A.8 Additional Related Work

MOO[40] introduced multi-modal prompt capability to RT-1, while Q-Transformer [7] adapted RT-1 to the Q-learning setting. RoboFlamingo [27] constructed a VLA based on the existing Flamingo VLM [1, 3]. ACT [47] 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 [20], PaLM-E [13], and ChatGPT for Robotics [43].

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 [10], has prompted the adoption of similar objectives in various domains. In computer vision models and VLMs, representative works like MAE [17] and ViLBERT [30] have employed comparable strategies. VLAs have also utilized masked modeling objectives for their vision

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.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-d1f5fa68-44a7-4ee7-b917-72eec47470a1.jpg)

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.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.01147v1-figure-5442ac3d-2262-446b-ac19-77b5f4c3e63c.jpg)

encoders, such as MVP [36], Voltron [25], GR-1 [44]. While these approaches have proven beneficial for the vision encoder, they often overlook the crucial alignment between different modalities.
Yueen Ma Dafeng Chi Shiguang Wu Yuecheng Liu Yuzheng Zhuang Jianye Hao Irwin King
Actra: Optimized Transformer Architecture for Vision-Language-Action Models in Robot Learning
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.