RAPIDE : Tokenisation efficace des actions pour les modèles Vision-Language-Action

FAST: Efficient Action Tokenization for Vision-Language-Action Models

January 16, 2025

Auteurs: Karl Pertsch, Kyle Stachowicz, Brian Ichter, Danny Driess, Suraj Nair, Quan Vuong, Oier Mees, Chelsea Finn, Sergey Levine

cs.AI

Résumé

Les modèles de séquences autorégressifs, tels que les politiques vision-langage basées sur les Transformers, peuvent être extrêmement efficaces pour capturer des comportements robotiques complexes et généralisables. Cependant, de tels modèles nous obligent à choisir une tokenisation de nos signaux d'action continus, ce qui détermine comment les symboles discrets prédits par le modèle se mappent aux actions robotiques continues. Nous constatons que les approches actuelles de tokenisation des actions des robots, basées sur des schémas de regroupement simples par dimension et par pas de temps, ont généralement de mauvaises performances lors de l'apprentissage de compétences habiles à partir de données robotiques à haute fréquence. Pour relever ce défi, nous proposons un nouveau schéma de tokenisation des actions des robots basé sur la transformée en cosinus discrète. Notre approche de tokenisation, appelée Tokenisation des Séquences d'Actions dans l'Espace Fréquentiel (FAST), nous permet d'entraîner des politiques vision-langage autorégressives pour des tâches hautement habiles et à haute fréquence là où les méthodes de discrétisation standard échouent complètement. Basé sur FAST, nous lançons FAST+, un tokeniseur d'actions de robot universel, entraîné sur 1M de trajectoires d'actions de robot réelles. Il peut être utilisé comme un tokeniseur boîte noire pour un large éventail de séquences d'actions de robot, avec des espaces d'actions divers et des fréquences de contrôle variées. Enfin, nous montrons que, combinée avec le VLA pi0, notre méthode peut être mise à l'échelle pour s'entraîner sur 10 000 heures de données robotiques et égaler les performances des VLAs de diffusion, tout en réduisant le temps d'entraînement jusqu'à 5 fois.

English

Autoregressive sequence models, such as Transformer-based vision-language action (VLA) policies, can be tremendously effective for capturing complex and generalizable robotic behaviors. However, such models require us to choose a tokenization of our continuous action signals, which determines how the discrete symbols predicted by the model map to continuous robot actions. We find that current approaches for robot action tokenization, based on simple per-dimension, per-timestep binning schemes, typically perform poorly when learning dexterous skills from high-frequency robot data. To address this challenge, we propose a new compression-based tokenization scheme for robot actions, based on the discrete cosine transform. Our tokenization approach, Frequency-space Action Sequence Tokenization (FAST), enables us to train autoregressive VLAs for highly dexterous and high-frequency tasks where standard discretization methods fail completely. Based on FAST, we release FAST+, a universal robot action tokenizer, trained on 1M real robot action trajectories. It can be used as a black-box tokenizer for a wide range of robot action sequences, with diverse action spaces and control frequencies. Finally, we show that, when combined with the pi0 VLA, our method can scale to training on 10k hours of robot data and match the performance of diffusion VLAs, while reducing training time by up to 5x.