# FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning

arXiv: https://arxiv.org/abs/2408.16944 · HTML: https://mlsanity.com/paper/2408.16944/

## Key facts

- Cited by 33 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 36 of its 39 references that are in the map.
- First posted on arXiv 29 Aug 2024.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 36)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Voyager: An Open-Ended Embodied Agent with Large Language Models](https://mlsanity.com/paper/2305.16291/) (2023)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning](https://mlsanity.com/paper/2306.03310/) (2023)
- [DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset](https://mlsanity.com/paper/2403.12945/) (2024)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [BridgeData V2: A Dataset for Robot Learning at Scale](https://mlsanity.com/paper/2308.12952/) (2023)
- [Mobile ALOHA: Learning Bimanual Mobile Manipulation with Low-Cost Whole-Body Teleoperation](https://mlsanity.com/paper/2401.02117/) (2024)
- [GMFlow: Learning Optical Flow via Global Matching](https://mlsanity.com/paper/2111.13680/) (2021)
- [Real-World Robot Learning with Masked Visual Pre-training](https://mlsanity.com/paper/2210.03109/) (2022)
- [Goal-Conditioned Imitation Learning using Score-based Diffusion Policies](https://mlsanity.com/paper/2304.02532/) (2023)
- [Interactive Language: Talking to Robots in Real Time](https://mlsanity.com/paper/2210.06407/) (2022)
- [Any-point Trajectory Modeling for Policy Learning](https://mlsanity.com/paper/2401.00025/) (2023)
- [RT-H: Action Hierarchies Using Language](https://mlsanity.com/paper/2403.01823/) (2024)
- [Language Conditioned Imitation Learning over Unstructured Data](https://mlsanity.com/paper/2005.07648/) (2020)
- [LIV: Language-Image Representations and Rewards for Robotic Control](https://mlsanity.com/paper/2306.00958/) (2023)
- [Language-Driven Representation Learning for Robotics](https://mlsanity.com/paper/2302.12766/) (2023)
- [Learning to Act from Actionless Videos through Dense Correspondences](https://mlsanity.com/paper/2310.08576/) (2023)
- [RoboCLIP: One Demonstration is Enough to Learn Robot Policies](https://mlsanity.com/paper/2310.07899/) (2023)
- [Language Embedded Radiance Fields for Zero-Shot Task-Oriented Grasping](https://mlsanity.com/paper/2309.07970/) (2023)
- [RoboTAP: Tracking Arbitrary Points for Few-Shot Visual Imitation](https://mlsanity.com/paper/2308.15975/) (2023)
- [Learning and Retrieval from Prior Data for Skill-based Imitation Learning](https://mlsanity.com/paper/2210.11435/) (2022)
- [Data Quality in Imitation Learning](https://mlsanity.com/paper/2306.02437/) (2023)
- [Can Foundation Models Perform Zero-Shot Task Specification For Robot Manipulation?](https://mlsanity.com/paper/2204.11134/) (2022)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [Vision-based Manipulation from Single Human Video with Open-World Object Graphs](https://mlsanity.com/paper/2405.20321/) (2024)

## Built on by (most cited first, 30 of 33)

- [Humanoid Policy ~ Human Policy](https://mlsanity.com/paper/2503.13441/) (2025)
- [IGOR: Image-GOal Representations are the Atomic Control Units for Foundation Models in Embodied AI](https://mlsanity.com/paper/2411.00785/) (2024)
- [Robometer: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons](https://mlsanity.com/paper/2603.02115/) (2026)
- [Motion Tracks: A Unified Representation for Human-Robot Transfer in Few-Shot Imitation Learning](https://mlsanity.com/paper/2501.06994/) (2025)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)
- [DataMIL: Selecting Data for Robot Imitation Learning with Datamodels](https://mlsanity.com/paper/2505.09603/) (2025)
- [RICL: Adding In-Context Adaptability to Pre-Trained Vision-Language-Action Models](https://mlsanity.com/paper/2508.02062/) (2025)
- [AMPLIFY: Actionless Motion Priors for Robot Learning from Videos](https://mlsanity.com/paper/2506.14198/) (2025)
- [Mechanistic Finetuning of Vision-Language-Action Models via Few-Shot Demonstrations](https://mlsanity.com/paper/2511.22697/) (2025)
- [Data Retrieval with Importance Weights for Few-Shot Imitation Learning](https://mlsanity.com/paper/2509.01657/) (2025)
- [GenFlowRL: Shaping Rewards with Generative Object-Centric Flow in Visual Reinforcement Learning](https://mlsanity.com/paper/2508.11049/) (2025)
- [cVLA: Towards Efficient Camera-Space VLAs](https://mlsanity.com/paper/2507.02190/) (2025)
- [Latent Policy Steering with Embodiment-Agnostic Pretrained World Models](https://mlsanity.com/paper/2507.13340/) (2025)
- [Robotic VLA Benefits from Joint Learning with Motion Image Diffusion](https://mlsanity.com/paper/2512.18007/) (2025)
- [COLLAGE: Adaptive Fusion-based Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2508.01131/) (2025)
- [Retrieve, Don't Retrain: Extending Vision Language Action Models to New Tasks at Test Time](https://mlsanity.com/paper/2606.15631/) (2026)
- [Skills Made to Order: Efficient Acquisition of Robot Cooking Skills Guided by Multiple Forms of Internet Data](https://mlsanity.com/paper/2409.15172/) (2024)
- [Factorizing Diffusion Policies for Observation Modality Prioritization](https://mlsanity.com/paper/2509.16830/) (2025)
- [Policy Compatible Skill Incremental Learning via Lazy Learning Interface](https://mlsanity.com/paper/2509.20612/) (2025)
- [Grounding Robot Generalization in Training Data via Retrieval-Augmented VLMs](https://mlsanity.com/paper/2603.11426/) (2026)
- [Task-agnostic Lifelong Robot Learning with Retrieval-based Weighted Local Adaptation](https://mlsanity.com/paper/2410.02995/) (2024)
- [Quality over Quantity: Demonstration Curation via Influence Functions for Data-Centric Robot Learning](https://mlsanity.com/paper/2603.09056/) (2026)
- [Difference-Aware Retrieval Policies for Imitation Learning](https://mlsanity.com/paper/2606.09758/) (2026)
- [ATHENA: Accelerated Multi-Task Heterogeneous Influence Functions for Robot Data Curation](https://mlsanity.com/paper/2606.16208/) (2026)
- [Flow as Flow: Modeling Robot Velocity Fields as Probability Velocity Fields for Flow-Based Object Manipulation](https://mlsanity.com/paper/2606.23090/) (2026)
- [RoboTok: A Scalable Data Engine for Internet Demonstration Video Retrieval and Dexterous Manipulation Learning](https://mlsanity.com/paper/2609.03199/) (2026)
- [Using Temperature Sampling to Effectively Train Robot Learning Policies on Imbalanced Datasets](https://mlsanity.com/paper/2510.19373/) (2025)
- [TSD: A Physics-Inspired Trajectory Saliency Detector for Efficient Imitation Learning](https://mlsanity.com/paper/2606.23371/) (2026)
- [Reuse Before You Retrieve: Diagnosing Headroom and Complementarity for Test-Time Augmentation of Embodied Multimodal Policies](https://mlsanity.com/paper/2608.17484/) (2026)
- [Hierarchical Skill Retrieval for Data-Efficient Adaptation of Vision-Language-Action Models](https://mlsanity.com/paper/2608.24042/) (2026)

## Abstract (arXiv)

Few-shot imitation learning relies on only a small amount of task-specific demonstrations to efficiently adapt a policy for a given downstream tasks. Retrieval-based methods come with a promise of retrieving relevant past experiences to augment this target data when learning policies. However, existing data retrieval methods fall under two extremes: they either rely on the existence of exact behaviors with visually similar scenes in the prior data, which is impractical to assume; or they retrieve based on semantic similarity of high-level language descriptions of the task, which might not be that informative about the shared low-level behaviors or motions across tasks that is often a more important factor for retrieving relevant data for policy learning. In this work, we investigate how we can leverage motion similarity in the vast amount of cross-task data to improve few-shot imitation learning of the target task. Our key insight is that motion-similar data carries rich information about the effects of actions and object interactions that can be leveraged during few-shot adaptation. We propose FlowRetrieval, an approach that leverages optical flow representations for both extracting similar motions to target tasks from prior data, and for guiding learning of a policy that can maximally benefit from such data. Our results show FlowRetrieval significantly outperforms prior methods across simulated and real-world domains, achieving on average 27% higher success rate than the best retrieval-based prior method. In the Pen-in-Cup task with a real Franka Emika robot, FlowRetrieval achieves 3.7x the performance of the baseline imitation learning technique that learns from all prior and target data. Website: https://flow-retrieval.github.io