Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

VAT: Vision Action Transformer by Unlocking Full Representation of ViT

Wenhao Li, Chengwei Ma, Weixin Mao

TL;DR

This work introduces Vision Action Transformer (VAT), a ViT-augmented policy architecture that leverages the full Hierarchy of ViT representations by injecting and updating action tokens at every transformer layer. VAT comprises parallel Vision and Action Modules with cross-attention and FiLM-based task conditioning, enabling deep fusion of perception and control. On LIBERO benchmarks, VAT achieves state-of-the-art 98.15% average success, and ablations show the full representation trajectory is crucial for performance, especially on long-horizon tasks. The approach demonstrates that preserving and integrating intermediate ViT representations can substantially improve robotic imitation learning and suggests a broader principle for embodied AI: exploit the complete feature hierarchy of vision models.

Abstract

In robot learning, Vision Transformers (ViTs) are standard for visual perception, yet most methods discard valuable information by using only the final layer's features. We argue this provides an insufficient representation and propose the Vision Action Transformer (VAT), a novel architecture that is extended from ViT and unlocks the full feature hierarchy of ViT. VAT processes specialized action tokens with visual features across all transformer layers, enabling a deep and progressive fusion of perception and action generation. On a suite of simulated manipulation tasks, VAT achieves a 98.15\% average success rate across four LIBERO benchmarks, establishing a new state-of-the-art by outperforming prior methods like OpenVLA-OFT. Our work presents not only a powerful model for imitation learning but also demonstrates the critical importance of leveraging the complete ''representation trajectory'' of vision models to advance robotic policy. The GitHub URL for the project code is https://github.com/sellerbubble/VAT.

VAT: Vision Action Transformer by Unlocking Full Representation of ViT

TL;DR

This work introduces Vision Action Transformer (VAT), a ViT-augmented policy architecture that leverages the full Hierarchy of ViT representations by injecting and updating action tokens at every transformer layer. VAT comprises parallel Vision and Action Modules with cross-attention and FiLM-based task conditioning, enabling deep fusion of perception and control. On LIBERO benchmarks, VAT achieves state-of-the-art 98.15% average success, and ablations show the full representation trajectory is crucial for performance, especially on long-horizon tasks. The approach demonstrates that preserving and integrating intermediate ViT representations can substantially improve robotic imitation learning and suggests a broader principle for embodied AI: exploit the complete feature hierarchy of vision models.

Abstract

In robot learning, Vision Transformers (ViTs) are standard for visual perception, yet most methods discard valuable information by using only the final layer's features. We argue this provides an insufficient representation and propose the Vision Action Transformer (VAT), a novel architecture that is extended from ViT and unlocks the full feature hierarchy of ViT. VAT processes specialized action tokens with visual features across all transformer layers, enabling a deep and progressive fusion of perception and action generation. On a suite of simulated manipulation tasks, VAT achieves a 98.15\% average success rate across four LIBERO benchmarks, establishing a new state-of-the-art by outperforming prior methods like OpenVLA-OFT. Our work presents not only a powerful model for imitation learning but also demonstrates the critical importance of leveraging the complete ''representation trajectory'' of vision models to advance robotic policy. The GitHub URL for the project code is https://github.com/sellerbubble/VAT.

Paper Structure

This paper contains 19 sections, 8 equations, 4 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Method
  4. Experiments
  5. Experimental Configuration of VAT
  6. Model Layer Skipping
  7. Attention Heatmap Visualization
  8. Comprehensive Ablation And Generalization Analysis
  9. Conclusion
  10. Ethics Statement
  11. Reproducibility Statement
  12. Experiments on Training Configurations
  13. Experiments on Camera Views
  14. Experiments on Loss Function
  15. Experiments on Trainable Parameters
  16. ...and 4 more sections

Figures (4)

  • Figure 1: VAT architecture within a single layer. A standard ViT block (Vision Module, left) processes vision tokens. In parallel, a new Action Module (right) updates action tokens by cross-attention to vision tokens. Task-specific information is injected via a FiLM layer, and the action module mirrors the vision module's structure but with its own trainable parameters.
  • Figure 2: Results of VAT Layer Skipping Experiments
  • Figure 3: Attention heatmap of VAT on LIBERO-Spatial
  • Figure 4: Attention heatmap of VAT on LIBERO-Object, Goal and 10