PaperlayerView exact revision on arXiv

Source-linked AI summary

In-Context Imitation Learning via Next-Token Prediction

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

arXiv:2408.15980v2cs.ROcs.AI

TL;DR

Pretrained robot policies still struggle to perform unseen tasks in new environments without additional training. ICRT uses sensorimotor trajectories as prompts for a causal transformer, and physical experiments show generalization to unseen objects and novel environment configurations, particularly when multiple tasks are possible.

  • Problem

    Pretrained robot models can show some generalization yet still struggle with unseen tasks in different environments without new demonstrations and fine-tuning.

    Source: §I · p. 1

  • Method

    ICRT autoregressively predicts from robot sensorimotor trajectories, using teleoperated demonstrations as prompts without updating model parameters.

    Source: §I · p. 1, §III · p. 2

  • Results

    ICRT generalizes learned motion primitives to unseen objects and novel environment configurations, particularly when multiple tasks are present.

    Source: §VII · p. 6

  • Takeaways & Limitations

    Physical experiments on a Franka Emika robot show that prompts can specify unseen tasks at varying levels of generalization.

    Source: §I · p. 2

    Generalization to completely unseen action primitives remains unclear, and ICRT assumes a fixed robot morphology and impedance controller.

    Source: §VII · p. 6

Abstract

from arXiv · show

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/

I. INTRODUCTION

Large pretrained robot policies still struggle with unseen tasks in new environments without further data collection or fine-tuning. ICRT applies in-context learning to robot trajectories, while its multi-task data design encourages the model to use prompts to identify tasks and objects.

I. INTRODUCTION

  • Existing pretrained robot models often need new teleoperated demonstrations or data from hand-crafted primitives, followed by fine-tuning, to complete unseen tasks.This adds workflow complexity in real-world environments.Source: §I · p. 1
  • ICRT uses trajectories of image observations, proprioceptive states, and actions as prompts to perform unseen tasks in different environment configurations without additional training.The trajectories implicitly encode the task primitives and objects to interact with.Source: §I · p. 1
  • ICRT operates directly on raw robot trajectories, avoids key-point or key-frame identification, and supports one or more trajectories in a long context window.The authors contrast this with few-shot imitation methods that use more complicated losses or prior knowledge.Source: §I · p. 1
  • Multi-task environments where the same initial observation permits different tasks encourage the model to use prompts to identify the task and the object.In other multi-task data, unique objects in each environment can let the model infer the task from the observation alone.Source: §I · pp. 1–2
  • The paper introduces ICRT, a multi-task dataset and training paradigm, and physical experiments on a Franka Emika robot to evaluate prompt-specified unseen tasks.The introduction frames these as contributions toward real-robot in-context learning.Source: §I · p. 2

A. Multi-Task Imitation Learning for Robotics

  • Behavior cloning, energy-based models, and diffusion models are established imitation-learning approaches, but typically require a separate model per task.Multi-task policies can also be distilled from task-specific models after training.Source: A. Multi-Task Imitation Learning for… · p. 2
  • Recent robot next-token-prediction policies use transformers trained on large multi-task datasets; Octo and OpenVLA are identified as state-of-the-art examples.Octo uses goal images and language instructions, whereas OpenVLA uses language instructions and visual observations.Source: A. Multi-Task Imitation Learning for… · p. 2

B. In-Context Learning

Prior in-context imitation methods use task matching, object segmentation, keypoint representations, or auxiliary losses, while ICRT prompts a transformer directly with robot trajectories. The problem setting evaluates whether a model can use a few new teleoperated demonstrations to handle unseen objects and novel environment configurations without further training.

B. In-Context Learning

  • Prior in-context imitation methods include contrastive task encoders, visual servoing with object segmentation, keypoint-action tokens, and video prompting with auxiliary losses; ICRT uses raw images without extra perception modules.The cited visual-servoing approach requires an additional object segmentation model, while Vid2Robot requires many auxiliary losses.Source: B. In-Context Learning · p. 2
  • ICRT directly prompts a transformer with sensorimotor trajectories, unlike simulated approaches that rely on full states or known reward functions.ICRT does not model rewards, uses a longer context window, and demonstrates in-context learning in physical experiments using image observations.Source: B. In-Context Learning · p. 2

III. PROBLEM STATEMENT

  • ICRT's test setting uses a few new human-teleoperated demonstrations as prompts to handle unseen tasks in novel environment configurations without additional training.Environment configuration refers to the objects present and their spatial arrangement.Source: §III · p. 2
  • The evaluation varies test-time environment configuration to assess selection of the motion primitive and interaction object; new tasks use unseen objects but training-set primitives.A task is characterized by a motion primitive and the objects used with it.Source: §III · p. 3
  • The experiments assume diverse training demonstrations from fixed cameras and a wrist camera, and that each test task is within the robot's reachable workspace.Each trajectory includes RGB observations, proprioception, and action.Source: §III · p. 3

IV. APPROACH

ICRT is trained on grouped visuomotor trajectories, with same-task trajectories concatenated so that prompts provide context for action prediction. The data design emphasizes multi-task scenes where the current observation alone does not determine the task.

A. Data Formulation

  • Each visuomotor trajectory consists of camera images, proprioceptive robot states, and actions; the state is absolute end-effector pose and actions are pose deltas plus continuous gripper control.Trajectories are grouped into disjoint task sets, using semantic labels where available.Source: A. Data Formulation · p. 3
  • DROID contributes roughly 2,000 pretraining trajectories after sampling 10,000 demonstrations and grouping language-instruction labels by CLIP text-embedding similarity.The grouping uses a 0.9 cosine-similarity threshold and requires at least four trajectories per task group.Source: A. Data Formulation · p. 3
  • ICRT-MT contains 1,098 trajectories across 29 tasks and six motion primitives, with multiple possible tasks from each environment observation to encourage prompt use.The listed primitives are picking, pick-and-place, stacking, pushing, poking, and opening and closing drawers.Source: A. Data Formulation · p. 3
  • Single-task scenes can let a model infer the task from the current observation rather than the prompt, motivating ICRT-MT's multiple-task-per-scene design.ICRT-MT was collected using the DROID setup.Source: A. Data Formulation · p. 3
  • For each epoch, trajectories from each task are shuffled and concatenated; each batch samples a 512-step subsequence, where up to five same-task trajectories usually appear, and the first k are prompts including one complete trajectory.Source: A. Data Formulation · p. 3

B. Model Architecture

ICRT projects image, proprioception, and action inputs into a shared sequence for a causal transformer, then decodes predicted states and actions. At inference, demonstrations and current observations condition autoregressive action prediction, with KV caching used for faster execution.

B. Model Architecture

  • A pretrained vision encoder, modality-specific projectors, and a causal transformer form ICRT's three main model components.The method overview depicts visual and proprioceptive features combined into a state representation, with actions encoded separately.Source: B. Model Architecture · p. 4
  • The vision encoder processes multiple camera views, while attention pooling combines visual tokens with proprioception into a state token and an MLP embeds actions.The resulting state and action tokens are passed to the transformer.Source: B. Model Architecture · p. 4
  • MLP decoder heads produce state and action outputs from the transformer, which is trained to predict shifted proprioceptive states and actions.The model predicts the next h actions at each time step and uses temporal ensembling to execute the final action.Source: B. Model Architecture · p. 4
  • The authors also test a pretrained Llama2-7B initialized transformer, fine-tuned with rank-32 LoRA because of the domain gap and compute limits.The Llama2-Base model is randomly initialized with 12 layers and a latent dimension of 768; full fine-tuning was not feasible.Source: B. Model Architecture · p. 4
  • Training applies action-prediction L1 loss only after prompt trajectories, rather than on the prompt itself.This loss placement is intended to provide supervision for responding to trajectory prompts at test time.Source: B. Model Architecture · p. 4
  • At inference, teleoperated demonstrations and current observations condition the model's next-action prediction, which is repeated as updated observations arrive.The example pipeline shows a prompt demonstration used to perform the task in a different environment.Source: B. Model Architecture · p. 4, Figure 4 · p. 5
  • KV caching reduces sequence-processing complexity from quadratic to linear in sequence length; ICRT is reported to infer at 39.6 Hz for real-time closed-loop control.The cache stores previous outputs so the model computes outputs only for new tokens.Source: B. Model Architecture · p. 5

V. EXPERIMENTS

Experiments test whether prompt trajectories support novel tasks across pick-and-place and poking, using rollouts that differ from the prompts. ICRT outperforms goal- or language-conditioned baselines, though object confusions and grasp or placement errors remain.

Experiment Design

  • The evaluation uses six unseen tasks for each of two primitives: pick-and-place and poking.Each task has five difficulty tiers, and each tier is attempted once; pick-and-place tiers add distractors or placement choices, while poking tiers add distractor objects.Source: §V · p. 5

Results

  • ICRT outperforms two state-of-the-art goal- or language-conditioned robot foundation models in both pick-and-place and poking tasks.Table I reports mean success rates for each task and overall, with standard errors; evaluation covers six unseen tasks per primitive and 60 trials overall.Source: Table I · p. 5

Prompt Generation

  • Rollouts use a different environment setup from the prompt demonstrations, testing task-relevant understanding rather than action copying.Three demonstrations are collected per task, and a random demonstration is used as the test prompt.Source: §V · p. 5

Results

  • Goal-conditioned policies identify pick targets without distractors but degrade as distractors increase and may fail when a goal image does not specify how to act in the current environment.For poking, where the goal position can resemble the start configuration, ICRT uses the prompt trajectory to identify the task and reach the target while ignoring distractors.Source: §V · pp. 5–6
  • OpenVLA may need more than 50 demonstrations per task, while ICRT still sometimes misses grasps, selects the wrong object, or places objects incorrectly.ICRT struggles when a distractor has the same color but a different shape.Source: §V · p. 6

VI. ABLATIONS

This section points to additional ablations of core design choices, with results in Table II and further studies on the project website.

VI. ABLATIONS

  • Additional experiments ablate core design choices, with results presented in Table II.Further ablation studies are available on the official website.Source: §VI · p. 6

A. Model Initialization

The model-initialization ablation finds that the language-pretrained variant has lower training loss but worse performance than smaller counterparts. Dataset ablations indicate that ICRT-MT's shared-start task structure is important for in-context learning.

A. Model Initialization

  • ICRT-Llama2-7B has lower training loss but performs worse than smaller counterparts, possibly because its inference frequency is lower.The reported frequencies are 10.7 Hz for ICRT-Llama2 and 39.6 Hz for the smaller counterpart.Source: A. Model Initialization · p. 6

B. Training Dataset

  • Training on the DROID subset is insufficient to complete any test task, while the ICRT-MT structure is particularly beneficial for in-context learning.ICRT-MT lets multiple tasks be performed from the same initial observation, requiring the model to rely on context to identify the task and object.Source: B. Training Dataset · p. 6
  • ICRT (MT) performs similarly to DROID-pretrained ICRT on pick-up and place, surpasses it on putting a radish in a grey bowl, but is weaker on poking.The authors suggest pretraining on a large, diverse dataset may improve visual-feature and control alignment.Source: B. Training Dataset · p. 6

C. No Prompt Loss

ICRT performs better when action loss is applied only after the prompt trajectories. The conclusion also reports generalization to unseen objects and configurations, while identifying limits in primitive and robot-morphology transfer.

C. No Prompt Loss

  • Applying action-prediction loss only after the prompt trajectories significantly improves ICRT’s performance.The authors hypothesize that prompt loss can encourage unconditional generation from observations, reducing attention to the prompt when multiple tasks are possible.Source: C. No Prompt Loss · p. 6

VII. LIMITATIONS AND CONCLUSION

  • Generalization to completely unseen action primitives remains unclear; the reported results cover only some primitives resembling those in training.The authors identify primitive-level generalization as a direction for future work.Source: §VII · p. 6
  • ICRT assumes a fixed robot morphology and impedance controller, leaving transfer across robot morphologies for future investigation.Source: §VII · p. 6
  • ICRT-Llama2’s low inference frequency may contribute to its low performance, and the authors aim to speed up inference.Source: §VII · p. 6
  • ICRT generalizes learned motion primitives to unseen objects and novel configurations, particularly when multiple tasks are present.The conclusion describes this generalization as using robot sensorimotor trajectories as context.Source: §VII · p. 6

A. Scene Illustrations

The appendix describes test scenes distinct from the prompts and reports repeatability and prompt-type experiments for the black-dog placement task. It also outlines additional ablations and an evaluation of two unseen primitives.

A. Scene Illustrations

B. Ablation Studies

1) Repeatability Experiments:

2) Prompt Trajectories:

  • Prompt types and prompt counts yield similar success rates in the tested task, suggesting performance is not sensitive to these prompt variations.The tested single prompts vary in distractors or placement; experiments also use two or three prompts.Source: Table IV · p. 9, 2) Prompt Trajectories · p. 10
  • Table V covers three tasks using two unseen primitives, with success rates calculated over five trials per task.Source: Table V · p. 10

3) Unseen Primitives:

The appendix reports results on primitives resembling training primitives, an ablation of joint DROID and ICRT-MT training, and implementation details including action parameterization and hyperparameter tables.

3) Unseen Primitives:

  • ICRT achieves decent success rates on three tasks using two unseen primitives that resemble training primitives.The tasks include grasping and dropping a toy tiger, grasping and dropping an unseen blue sponge, and placing the sponge to the right of the tiger.Source: 3) Unseen Primitives · p. 10

4) Co-training:

  • Single-stage co-training with DROID and ICRT-MT converges more slowly and performs poorly on responding to prompts and completing tasks.The authors hypothesize that DROID’s 200 tasks versus ICRT-MT’s 29 contribute to a dataset imbalance that makes task learning harder.Source: 4) Co-training · p. 10, Table VI · p. 10

C. Detailed results

D. Hyperparameters

E. Parameterization

  • The action representation combines relative end-effector translation, 6DoF rotation, and absolute gripper position into a 10-dimensional vector.Rollouts also use temporal ensembling and receding-horizon control with an action horizon of 10.Source: E. Parameterization · p. 11

F. System Information

The system-information section reports training durations and inference frequency for ICRT variants on specified hardware.

F. System Information

  • Using the proposed formulation, which can leverage the KV cache, ICRT-Llama2 runs naively at 10Hz; its reported inference speed is averaged over 100 steps.Tests were performed on a workstation with an NVIDIA RTX 3090Ti, Intel i5-12400F and 64GB memory.Source: F. System Information · p. 12
  • ICRT pre-training on DROID takes 56 minutes, and fine-tuning on ICRT-MT takes 18 hours.Source: F. System Information · p. 12
  • ICRT-Llama7B takes roughly 28 hours to fine-tune.Source: F. System Information · p. 12
Loading 2408.15980v2…