Aprendizado por Imitação em Contexto via Previsão do Próximo Token

In-Context Imitation Learning via Next-Token Prediction

August 28, 2024

Autores: Letian Fu, Huang Huang, Gaurav Datta, Lawrence Yunliang Chen, William Chung-Ho Panitch, Fangchen Liu, Hui Li, Ken Goldberg

cs.AI

Resumo

Exploramos como melhorar os modelos de previsão da próxima ficha para realizar aprendizado por imitação em contexto em um robô real, onde o robô executa novas tarefas interpretando informações contextuais fornecidas durante a fase de entrada, sem atualizar seus parâmetros de política subjacentes. Propomos o Transformer de Robô em Contexto (ICRT), um transformer causal que realiza previsão autoregressiva em trajetórias sensoriomotoras sem depender de dados linguísticos ou função de recompensa. Essa formulação permite a execução flexível e sem treinamento de novas tarefas no momento do teste, alcançada solicitando ao modelo trajetórias sensoriomotoras da nova tarefa compostas por observações de imagem, ações e tuplas de estados, coletadas por teleoperação humana. Experimentos com um robô Franka Emika demonstram que o ICRT pode se adaptar a novas tarefas especificadas por prompts, mesmo em configurações de ambiente que diferem tanto do prompt quanto dos dados de treinamento. Em um ambiente de configuração multitarefa, o ICRT supera significativamente os modelos de previsão da próxima ficha mais avançados no estado atual da arte em robótica na generalização para tarefas não vistas. O código, checkpoints e dados estão disponíveis em https://icrt.dev/

English

We explore how to enhance next-token prediction models to perform in-context imitation learning on a real robot, where the robot executes new tasks by interpreting contextual information provided during the input phase, without updating its underlying policy parameters. We propose In-Context Robot Transformer (ICRT), a causal transformer that performs autoregressive prediction on sensorimotor trajectories without relying on any linguistic data or reward function. This formulation enables flexible and training-free execution of new tasks at test time, achieved by prompting the model with sensorimotor trajectories of the new task composing of image observations, actions and states tuples, collected through human teleoperation. Experiments with a Franka Emika robot demonstrate that the ICRT can adapt to new tasks specified by prompts, even in environment configurations that differ from both the prompt and the training data. In a multitask environment setup, ICRT significantly outperforms current state-of-the-art next-token prediction models in robotics on generalizing to unseen tasks. Code, checkpoints and data are available on https://icrt.dev/