On the Generalization Capabilities, Design Choices and Limitations of Keypoint Imitation Learning

Technical Summary: On the Generalization Capabilities, Design Choices and Limitations of Keypoint Imitation Learning

Problem Statement
Robotic manipulation in unstructured environments requires policies that can generalize to unseen objects, poses, and scene variations. While imitation learning is a promising approach, standard RGB-based policies typically demand large datasets to handle such variations effectively. This motivates the exploration of intermediate representations, such as keypoints, which offer compact, fine-grained semantic encoding. Although visual foundation models now enable one-shot keypoint extraction without task-specific training, it remains unclear how to optimally integrate these representations into imitation learning pipelines, which extraction strategies are most effective, and whether they genuinely outperform alternative object-centric representations like depth maps or segmentation masks.

Methodology
The authors propose a three-stage Keypoint Imitation Learning (KIL) pipeline:

1. Keypoint Specification: A human operator manually annotates 3–6 semantic keypoints on a reference image for each object category.
2. Representation Extraction: Using visual foundation models, the system extracts corresponding keypoints from every observation frame during training and inference. The 2D positions are lifted to 3D using calibrated RGB-D camera intrinsics and extrinsics. The paper investigates three extraction strategies:
   • Image Matching (IM): Computing descriptors for reference keypoints and finding best matches via cosine similarity in observation frames.
   • Instance Matching (IN): Extending IM by first generating instance masks via an open-vocabulary detector, then restricting the descriptor search to specific object masks.
   • Tracking (T): Using a dedicated tracking model initialized by instance matching to follow keypoints over time.
3. Policy Learning: The resulting sequence of 3D keypoints is encoded by a transformer encoder and fed into a diffusion policy (action head) to predict robot actions. The policy takes a history of observations and proprioceptive states as input.

Key Contributions

• Design Guidelines: The paper consolidates previous KIL methods and provides practical guidelines by systematically comparing extraction methods (IM, IN, T), foundation models (RADIO, DIFT, DINO), and policy architectures (transformer encoders vs. direct concatenation).
• Large-Scale Evaluation: The authors conduct over 2,000 real-world rollouts across two robot platforms and five distinct manipulation tasks (e.g., placing shoes/mugs, pouring bottles, picking pens).
• Multi-Instance Extension: The work extends KIL to tasks involving multiple instances of the same object category, addressing a limitation in previous single-object studies.
• Limitation Analysis: The study explicitly identifies and quantifies the limitations of current foundation models, particularly regarding robustness to large object orientation changes and occlusion.

Results

• Generalization Performance: KIL achieved an overall success rate of 75% across five tasks, significantly outperforming a standard RGB-based diffusion baseline (47%). The performance gap widened dramatically under scene variations (KIL: 70% vs. RGB: 10%).
• Comparison with Object-Centric Baselines: KIL performed on par with S2-diffusion (73%), an object-centric baseline utilizing depth maps and segmentation masks, suggesting that keypoints do not inherently offer superior generalization over other intermediate representations in all contexts.
• Extraction Methods: Image Matching, Instance Matching, and Tracking yielded similar performance (70–75% average success). Image Matching was selected as the default due to its simplicity and lower computational cost.
• Foundation Models: For image matching, RADIOv2.5-B (85% on specific tasks) outperformed DINOv3-B and DIFT. However, DIFT showed higher robustness when combined with instance matching or tracking.
• Encoder Design: A dedicated transformer encoder performed on par with a simple concatenation approach (85% success) but significantly outperformed a variant that appended cosine similarity scores to 3D positions (which degraded to 7% due to overfitting on similarity scores).
• Multi-Object Tasks: In tasks requiring the manipulation of two objects, KIL achieved a task completion rate of 61–63%, outperforming the RGB baseline (48%) and matching S2-diffusion (58%).
• Orientation Robustness: A critical finding was the sensitivity of KIL to object orientation. When object poses spanned a full ±180° range, KIL performance dropped by 55% (Image Matching) and 28% (Tracking). In contrast, RGB and S2 baselines dropped by only 10% and 9% respectively, indicating that current foundation models lack robustness to large orientation changes.

Significance and Claims
The paper concludes that KIL is a data-efficient approach that offers significant generalization benefits over raw RGB policies, particularly in scenes with distractors or background variations. However, the authors modestly claim that KIL does not universally outperform alternative object-centric representations like S2-diffusion. The performance of KIL is heavily dependent on the capabilities of the underlying foundation models used for extraction; specifically, KIL inherits the limitations of these models regarding occlusion handling and orientation invariance. The work clarifies that while keypoints provide a compact representation, their utility is bounded by the robustness of the extraction pipeline, suggesting that future improvements may require end-to-end fine-tuning of extraction models or hybrid representations.