[Paper Review] Keypoint Action Tokens Enable In-Context Imitation Learning in Robotics

Norman Di Palo, Edward Johns|arXiv (Cornell University)|Mar 28, 2024

Robot Manipulation and Learning4 citations

TL;DR

This paper introduces Keypoint Action Tokens (KAT), a method that enables off-the-shelf text-pretrained Transformers like GPT-4 Turbo to perform few-shot in-context imitation learning in robotics by converting visual observations and action trajectories into sequences of keypoint-based text tokens. Despite no fine-tuning on robotics data, KAT achieves state-of-the-art performance on real-world tasks with as few as 10 demonstrations, demonstrating that large language models can generalize to vision-and-action sequence-to-sequence tasks through pattern learning alone.

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 language models for embodied tasks. Videos are available at https://www.robot-learning.uk/keypoint-action-tokens.

Motivation & Objective

• To enable few-shot imitation learning in robotics using pre-trained language models without any fine-tuning on robotics data.
• To address the challenge of data scarcity in robot learning by leveraging the emergent few-shot pattern recognition capabilities of large text-based Transformers.
• To develop a general-purpose framework that maps visual observations and action sequences into a text-compatible token format for sequence-to-sequence learning.
• To evaluate whether off-the-shelf LLMs can perform as well as or better than specialized imitation learning models (e.g., diffusion policies) in low-data regimes.
• To investigate the scalability and performance of KAT across diverse real-world tasks, including spatial generalization, novel object generalization, and 6-DoF trajectory control.

Proposed method

• Keypoint Action Tokens (KAT) represent visual observations and action trajectories as sequences of keypoint coordinates, which are then converted into natural language tokens (e.g., 'left hand at (x,y,z)') for input to a text-based Transformer.
• The method uses a pre-trained Vision Transformer to extract a fixed number of keypoints (K) from each observation frame, encoding spatial and temporal structure into a sequence of descriptive tokens.
• Demonstrations and the current observation are concatenated into a prompt, where the input sequence includes both keypoint tokens and action tokens from expert demonstrations.
• The text-pretrained Transformer (e.g., GPT-4 Turbo) performs in-context learning by generating the next sequence of action tokens based on the prompt, which are then decoded into robot control commands.
• The framework operates entirely in inference mode, requiring no fine-tuning or additional training on robotics data.
• The method supports long-context inputs, enabling the processing of multiple demonstrations and complex trajectories within a single prompt.

Figure 1: An illustration of our pipeline. KAT transforms a visual observation into a sequence of keypoint tokens . From it, the text-pretrained Transformer (LLM) predicts action tokens , that are then executed as a trajectory of poses.

Experimental results

Research questions

• RQ1 Can off-the-shelf text-pretrained Transformers perform few-shot in-context imitation learning in robotics without any fine-tuning?
• RQ2 How well does KAT generalize to unseen objects, spatial configurations, and multi-modal behaviors with only 10 demonstrations?
• RQ3 Does the performance of KAT improve with newer generations of LLMs, even without any robotics-specific training?
• RQ4 Can KAT achieve performance comparable to or better than state-of-the-art diffusion policies in low-data regimes?
• RQ5 What are the limitations of KAT when scaling to longer demonstration sequences or real-time online control?

Key findings

• KAT achieves state-of-the-art performance on real-world robotic tasks such as Align T, Wipe Plate, Sweep, and Bottle Upright using only 10 demonstrations and no fine-tuning.
• GPT-4 Turbo outperforms both GPT-3 Turbo and GPT-2 in few-shot imitation learning, demonstrating that newer LLMs generalize better to vision-and-action sequences.
• The method performs on par or better than diffusion policies in low-data regimes, indicating that in-context learning with LLMs can be highly efficient for imitation learning.
• Performance improves with larger context windows in LLMs, suggesting that longer input sequences enhance pattern recognition and trajectory generation.
• The framework enables immediate deployment of learned skills after demonstration, as no model retraining is required.
• KAT shows strong generalization to novel objects and spatial configurations, indicating robustness to distributional shift in real-world settings.

Figure 2: Illustration of the extraction pipeline of keypoint tokens .