Обучение с имитацией в контексте через предсказание следующего токена

In-Context Imitation Learning via Next-Token Prediction

August 28, 2024

Авторы: Letian Fu, Huang Huang, Gaurav Datta, Lawrence Yunliang Chen, William Chung-Ho Panitch, Fangchen Liu, Hui Li, Ken Goldberg

cs.AI

Аннотация

Мы исследуем способы улучшения моделей предсказания следующего токена для выполнения обучения с имитацией в контексте на реальном роботе, где робот выполняет новые задачи, интерпретируя контекстную информацию, предоставленную во время фазы ввода, без обновления своих основных параметров политики. Мы предлагаем Ин-Контекстный Робот Трансформер (ICRT), причинный трансформер, выполняющий авторегрессионное предсказание на траекториях сенсорно-моторных действий без использования лингвистических данных или функции вознаграждения. Эта формулировка обеспечивает гибкое и обучение-независимое выполнение новых задач на этапе тестирования, достигаемое путем подачи модели сенсорно-моторных траекторий новой задачи, состоящих из наблюдений изображений, действий и кортежей состояний, собранных с помощью человеческой телеоперации. Эксперименты с роботом Franka Emika показывают, что ICRT способен адаптироваться к новым задачам, указанным с помощью подсказок, даже в конфигурациях среды, отличных как от подсказки, так и от данных обучения. В среде с множественными задачами ICRT значительно превосходит текущие передовые модели предсказания следующего токена в робототехнике по обобщению на невиденные задачи. Код, контрольные точки и данные доступны на 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/