Keypoint Action Tokens Enable In-Context Imitation Learning in Robotics (2403.19578v3)

Abstract: We show that off-the-shelf text-based Transformers, with no additional training, can perform few-shot in-context visual imitation learning, mapping visual observations to action sequences that emulate the demonstrator's behaviour. We achieve this by transforming visual observations (inputs) and trajectories of actions (outputs) into sequences of tokens that a text-pretrained Transformer (GPT-4 Turbo) can ingest and generate, via a framework we call Keypoint Action Tokens (KAT). Despite being trained only on language, we show that these Transformers excel at translating tokenised visual keypoint observations into action trajectories, performing on par or better than state-of-the-art imitation learning (diffusion policies) in the low-data regime on a suite of real-world, everyday tasks. Rather than operating in the language domain as is typical, KAT leverages text-based Transformers to operate in the vision and action domains to learn general patterns in demonstration data for highly efficient imitation learning, indicating promising new avenues for repurposing natural LLMs for embodied tasks. Videos are available at https://www.robot-learning.uk/keypoint-action-tokens.

• The paper introduces a novel imitation learning framework that tokenizes visual observations and robot actions to repurpose large pre-trained Transformers.
• It demonstrates efficient few-shot learning with competitive performance on real-world manipulation tasks, handling spatial generalization and multi-modal behaviors.
• The research paves the way for cross-domain transfer of Transformer capabilities from NLP to robotics, addressing data scarcity and accelerating robot learning.

Keypoint Action Tokens for Efficient Imitation Learning in Robotics

Introduction

The paper introduces a novel approach to few-shot imitation learning by utilizing off-the-shelf, large text-pretrained Transformers without additional training. This method, dubbed Keypoint Action Tokens (KAT), converts visual observations and trajectories of actions into sequences of tokens, which are subsequently ingested by a Transformer. A key achievement of KAT is demonstrating the capacity of these models, initially developed for natural language processing tasks, to efficiently transfer their sequence-to-sequence prediction capabilities to the domain of visual imitation learning in robotics. This approach allows for rapid learning from a minimal set of demonstrations and instant deployment of learned skills, presenting a significant step forward in robot learning efficiency by leveraging large models trained in data-rich domains.

Methodology

The methodology centers around the formulation of imitation learning as a sequence-to-sequence prediction problem, with visual observations and actions represented as sequences of tokens. The paper introduces two critical components:

• Visual Observation Tokenization: Utilizes Vision Transformers (ViTs) to transform visual observations into sequences of 3D keypoint tokens. These tokens encapsulate crucial geometric and semantic information from the observation images, effectively reducing the dimensionality of the input data.
• Action Sequence Tokenization: Transforms end-effector trajectories into sequences of action tokens. Actions are represented as triplets of 3D points in the robot's end-effector frame, then tokenized into character sequences. This conversion ensures that both observations and actions exist within the same representational space, aiding the Transformer's pattern recognition tasks.

Leveraging off-the-shelf, large, text-pretrained Transformers, this framework repurposes these models as efficient imitation learning machines. Through tokenizing observation and action sequences, the paper showcases that these Transformers can learn and generalize complex behaviors from remarkably few demonstrations.

Experiments and Results

Experimental validations were conducted across various real-world manipulation tasks. These experiments compare KAT's performance against state-of-the-art imitation learning methods and introduce novel baselines designed to isolate the contributions of keypoint and action token representations. The results underscore KAT's efficacy in few-shot learning challenges, achieving competitive or superior success rates in tasks requiring spatial generalization, novel object interaction, multi-modal behavior learning, and execution of 6-DoF action spaces.

Key findings include:

• KAT performs on par or better than existing state-of-the-art methods in low-data regimes.
• The approach provides promising avenues for the use of text-based Transformers across different data modalities, particularly in robotics, where data scarcity is a prominent challenge.
• The methodology suggests that the continued advancement of large Transformers can substantially benefit robot learning indirectly, by providing a means to efficiently learn and generalize from limited demonstrations.

Discussion and Future Work

The paper positions KAT not only as a step forward in leveraging LLMs for robotics but also highlights several potential areas for future exploration. The implications of this research extend to the broader utilization of pre-trained models across various domains, advocating for a cross-disciplinary approach to tackling data scarcity and learning efficiency issues. Moreover, it raises pertinent questions regarding the scalability of in-context learning as the number of demonstrations increases and suggests potential pathways for integrating dynamic keypoint extraction methods and exploring model fine-tuning as mechanisms to further enhance performance.

The work done on Keypoint Action Tokens fundamentally challenges and expands the scope of imitation learning in robotics. By harnessing the power of pre-trained Transformers, originally developed for text, this method offers a novel perspective on accelerating robot learning capabilities, serving as a base for future endeavors in the field.