Keypoint-Aware RAG for Robotic Manipulation: In-Context Constraint Learning via Large-Scale Retrieval

Abstract

Recent advances in robotic manipulation leverage foundation models pre-trained on internet-scale data, where keypoint-based representations have shown promising results in spatial reasoning. However, existing approaches primarily focus on zero-shot generalization or human-collected demonstrations, with limited exploration of large-scale robotic datasets. In this work, we propose Keypoint-Aware Retrieval Augmented Generation (KARAG), a simple yet novel framework that synergistically integrates visual-language models (VLMs) with robotic datasets through retrieval-augmented generation (RAG). Our framework bridges the retrieval and generation phases via in-context learning with keypoint-aware constraints, enabling simultaneous utilization of internet-scale knowledge and structured robotic datasets. Extensive experiments in both simulated and real-world environments demonstrate that KARAG significantly enhances the stability and accuracy of VLM-generated outputs without requiring human demonstrations or additional training, achieving 10%–20% success rate improvements in real-world scenarios and 12%–46% improvements in simulation over the baseline. Furthermore, we present an algorithm for converting large robotic datasets into Keyframe-Keypoint-Trajectory representations to facilitate retrieval. Our dataset and implementation are publicly available at https://github.com/RobertAckleyLin/KARAG/.