Apprentissage par imitation en contexte via la prédiction du jeton suivant

In-Context Imitation Learning via Next-Token Prediction

August 28, 2024

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

cs.AI

Résumé

Nous explorons comment améliorer les modèles de prédiction du prochain jeton pour réaliser un apprentissage par imitation en contexte sur un vrai robot, où le robot exécute de nouvelles tâches en interprétant les informations contextuelles fournies lors de la phase d'entrée, sans mettre à jour ses paramètres de politique sous-jacents. Nous proposons l'In-Context Robot Transformer (ICRT), un transformateur causal qui effectue des prédictions autorégressives sur les trajectoires sensorimotrices sans se fier à des données linguistiques ou à une fonction de récompense. Cette formulation permet une exécution flexible et sans entraînement de nouvelles tâches au moment du test, réalisée en incitant le modèle avec des trajectoires sensorimotrices de la nouvelle tâche composées d'observations d'images, d'actions et de tuples d'états, collectées via une téléopération humaine. Des expériences avec un robot Franka Emika démontrent que l'ICRT peut s'adapter à de nouvelles tâches spécifiées par des incitations, même dans des configurations d'environnement différentes à la fois de l'incitation et des données d'entraînement. Dans un environnement multitâche, l'ICRT surpasse significativement les modèles actuels de prédiction du prochain jeton de pointe en robotique en généralisant à des tâches non vues. Le code, les points de contrôle et les données sont disponibles sur 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/