FAST: Tokenização Eficiente de Ações para Modelos de Visão-Linguagem-Ação

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

Resumo

Modelos de sequência autoregressivos, como políticas de visão-linguagem baseadas em Transformadores, podem ser tremendamente eficazes para capturar comportamentos robóticos complexos e generalizáveis. No entanto, tais modelos exigem que escolhamos uma tokenização de nossos sinais de ação contínuos, o que determina como os símbolos discretos previstos pelo modelo se mapeiam para ações contínuas do robô. Observamos que abordagens atuais para tokenização de ações de robô, baseadas em esquemas simples de agrupamento por dimensão e por etapa de tempo, geralmente têm desempenho ruim ao aprender habilidades habilidosas a partir de dados de robô de alta frequência. Para enfrentar esse desafio, propomos um novo esquema de tokenização baseado em compressão para ações de robô, com base na transformada discreta do cosseno. Nossa abordagem de tokenização, Tokenização de Sequência de Ação no Espaço de Frequência (FAST), nos permite treinar VLAs autoregressivos para tarefas altamente habilidosas e de alta frequência, onde métodos de discretização padrão falham completamente. Com base no FAST, lançamos o FAST+, um tokenizador de ações de robô universal, treinado em 1M de trajetórias de ação de robô reais. Pode ser usado como um tokenizador de caixa-preta para uma ampla gama de sequências de ação de robô, com espaços de ação diversos e frequências de controle. Por fim, demonstramos que, quando combinado com o VLA pi0, nosso método pode escalar para treinamento em 10 mil horas de dados de robô e igualar o desempenho dos VLAs de difusão, reduzindo o tempo de treinamento em até 5 vezes.

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.