Rápido: Tokenización Eficiente de Acciones para Modelos de Visión-Lenguaje-Acción

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

January 16, 2025

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

cs.AI

Resumen

Los modelos de secuencia autoregresivos, como las políticas de visión-lenguaje basadas en Transformadores, pueden ser tremendamente efectivos para capturar comportamientos robóticos complejos y generalizables. Sin embargo, dichos modelos nos obligan a elegir una tokenización de nuestras señales de acción continuas, lo que determina cómo los símbolos discretos predichos por el modelo se asignan a acciones robóticas continuas. Encontramos que los enfoques actuales para la tokenización de acciones de robots, basados en esquemas de particionamiento simples por dimensión y por paso de tiempo, generalmente tienen un rendimiento deficiente al aprender habilidades diestras a partir de datos de robot de alta frecuencia. Para abordar este desafío, proponemos un nuevo esquema de tokenización de acciones de robots basado en la transformada discreta del coseno. Nuestro enfoque de tokenización, Tokenización de Secuencia de Acción en el Espacio de Frecuencia (FAST), nos permite entrenar VLAs autoregresivos para tareas altamente diestras y de alta frecuencia donde los métodos de discretización estándar fallan por completo. Basándonos en FAST, lanzamos FAST+, un tokenizador de acciones de robot universal, entrenado en 1M de trayectorias de acciones de robot reales. Puede ser utilizado como un tokenizador de caja negra para una amplia gama de secuencias de acciones de robot, con espacios de acción diversos y frecuencias de control. Finalmente, demostramos que, cuando se combina con el VLA pi0, nuestro método puede escalar para entrenar con 10k horas de datos de robot y igualar el rendimiento de los VLAs de difusión, al tiempo que reduce el tiempo de entrenamiento hasta en un 5x.

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.