次のトークン予測を通じたインコンテキスト模倣学習

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

要旨

次のトークン予測モデルを強化し、リアルロボット上でコンテキスト内での模倣学習を実行する方法を探求します。ロボットは、入力フェーズ中に提供された文脈情報を解釈して新しいタスクを実行し、基本的なポリシーパラメータを更新せずに行います。私たちはIn-Context Robot Transformer（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/