In-Context Imitation Learning via Next-Token Prediction (2408.15980v2)

Abstract: 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/

• The paper introduces the ICRT model that uses next-token prediction on sensorimotor trajectories to perform in-context imitation learning without parameter updates.
• It demonstrates substantial improvements in real-world tasks, achieving 76.7% success in pick-and-place and 93.3% in poking compared to baselines.
• The framework enables flexible, multi-task robotic control, reducing task-specific training overhead and promoting rapid adaptation to novel environments.

In-Context Imitation Learning via Next-Token Prediction

The paper "In-Context Imitation Learning via Next-Token Prediction" introduces the In-Context Robot Transformer (ICRT), a transformative approach in the field of robotic control. This approach aims to extend the capabilities of next-token prediction models, typically successful in language and vision domains, into practical robot learning applications. ICRT is designed to leverage in-context learning, adapting to new tasks through prompt-based demonstrations without further parameter updates.

Framework and Methodology

ICRT is constructed upon a causal transformer architecture, catering to autoregressive prediction on sensorimotor trajectories. By training on these trajectories—comprising image observations, actions, and proprioceptive states—the model circumvents the need for linguistic data or reward functions. The proposed framework allows for the robot to be prompted with teleoperated human demonstrations and consequently execute tasks in previously unseen environments. The versatility and robustness of ICRT are highlighted through its ability to manage multi-task environments, achieving a significant edge over current state-of-the-art models.

Contributions

The paper delineates the contributions in the following aspects:

1. ICRT Model Introduction: ICRT operates as a next-token prediction model, utilizing robot sensorimotor data as prompts to achieve in-context imitation learning in varying configurations.
2. Multi-Task Robot Dataset and Training Paradigm: The authors provide a specialized dataset named ICRT-Multi-Task (ICRT-MT), fostering the model's multi-task and in-context learning capabilities.
3. Empirical Validation: Physical experiments using a Franka Emika robot varied across different levels of task complexity to evaluate ICRT's effectiveness. The results manifest ICRT's favorable performance on unseen task generalization.

Experimental Setup and Results

The experimental setup involves real-world robotic tasks emphasizing two action primitives: pick-and-place and poking. Each task contains five levels of complexity, measuring how effectively ICRT can discern task-specific actions amidst distractors.

Key Results:

• Pick-and-Place: ICRT demonstrated a substantial improvement over baseline models. For instance, ICRT achieved a success rate of 76.7% on the task of picking up and placing objects, whereas the goal-conditioned policy managed only 33.3%, and Octo struggled at 5%.
• Poking: ICRT excelled particularly well in poking tasks, averting distractors and identifying correct objects with a success rate of 93.3%, compared to the goal-conditioned policy at 6.7%.

Implications and Future Directions

Practical: The introduction of ICRT facilitates an intuitive and efficient pathway for robotic policy learning, significantly reducing the overhead of task-specific training. This framework could be pivotal in advancing real-world robotic applications where flexibility and rapid adaptation are crucial.

Theoretical: The concept of in-context learning in robotics opens new research avenues. It underscores the potential of transformer-based architectures in learning task-agnostic representations, heralding a paradigm shift from traditional fine-tuning methods to more versatile, context-sensitive models.

Future Developments:

1. Enhanced Generalization: While the model shows promising results, extending its capabilities to entirely new task primitives remains an ongoing challenge. Future research could explore increasing model capacity and dataset diversity to achieve broader generalization.
2. Cross-morphology Transfer: Investigating how the model can be adapted to varying robot morphologies without re-training could significantly boost its applicability.
3. Inference Optimization: A noted bottleneck is the low inference frequency of the ICRT-Llama2 variant. Improving this aspect will be crucial for operational efficiency in real-time applications.

Conclusion

The research presented in "In-Context Imitation Learning via Next-Token Prediction" offers a novel and effective approach for real-world robot learning. By leveraging in-context learning through sensorimotor trajectories, ICRT achieves robustness and flexibility in task execution, outperforming existing next-token prediction models. This foundational work paves the way for future enhancements in robotic generalization and adaptation, underscoring the potential for broader applications in dynamic and complex environments.