TraceVLA: Visual Trace Prompting Melhora a Consciência Espaço-Temporal para Políticas Robóticas Generalistas

TraceVLA: Visual Trace Prompting Enhances Spatial-Temporal Awareness for Generalist Robotic Policies

December 13, 2024

Autores: Ruijie Zheng, Yongyuan Liang, Shuaiyi Huang, Jianfeng Gao, Hal Daumé III, Andrey Kolobov, Furong Huang, Jianwei Yang

cs.AI

Resumo

Embora os grandes modelos de visão-linguagem-ação (VLA) pré-treinados em extensos conjuntos de dados de robôs ofereçam políticas generalistas promissoras para a aprendizagem robótica, eles ainda enfrentam dificuldades com a dinâmica espaço-temporal na robótica interativa, tornando-os menos eficazes no manuseio de tarefas complexas, como manipulação. Neste trabalho, introduzimos a indução visual de rastreamento, uma abordagem simples, porém eficaz, para facilitar a consciência espaço-temporal dos modelos VLA para previsão de ações, codificando visualmente trajetórias estado-ação. Desenvolvemos um novo modelo TraceVLA, por meio do ajuste fino do OpenVLA em nosso próprio conjunto de dados coletado de 150 mil trajetórias de manipulação de robôs usando a indução visual de rastreamento. Avaliações do TraceVLA em 137 configurações no SimplerEnv e 4 tarefas em um robô físico WidowX demonstram desempenho de ponta, superando o OpenVLA em 10% no SimplerEnv e 3,5 vezes em tarefas de robô real, exibindo generalização robusta em diferentes encarnações e cenários. Para validar ainda mais a eficácia e generalidade de nosso método, apresentamos um modelo VLA compacto baseado em 4B Phi-3-Vision, pré-treinado no Open-X-Embodiment e ajustado fino em nosso conjunto de dados, que rivaliza com a linha de base 7B OpenVLA, ao mesmo tempo em que melhora significativamente a eficiência de inferência.

English

Although large vision-language-action (VLA) models pretrained on extensive robot datasets offer promising generalist policies for robotic learning, they still struggle with spatial-temporal dynamics in interactive robotics, making them less effective in handling complex tasks, such as manipulation. In this work, we introduce visual trace prompting, a simple yet effective approach to facilitate VLA models' spatial-temporal awareness for action prediction by encoding state-action trajectories visually. We develop a new TraceVLA model by finetuning OpenVLA on our own collected dataset of 150K robot manipulation trajectories using visual trace prompting. Evaluations of TraceVLA across 137 configurations in SimplerEnv and 4 tasks on a physical WidowX robot demonstrate state-of-the-art performance, outperforming OpenVLA by 10% on SimplerEnv and 3.5x on real-robot tasks and exhibiting robust generalization across diverse embodiments and scenarios. To further validate the effectiveness and generality of our method, we present a compact VLA model based on 4B Phi-3-Vision, pretrained on the Open-X-Embodiment and finetuned on our dataset, rivals the 7B OpenVLA baseline while significantly improving inference efficiency.