# Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets

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

## Key facts

- Cited by 58 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 38 of its 48 references that are in the map.
- First posted on arXiv 18 Apr 2023.

## 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 38)

- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [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)
- [CLIPort: What and Where Pathways for Robotic Manipulation](https://mlsanity.com/paper/2109.12098/) (2021)
- [BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning](https://mlsanity.com/paper/2202.02005/) (2022)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training](https://mlsanity.com/paper/2210.00030/) (2022)
- [Transporter Networks: Rearranging the Visual World for Robotic Manipulation](https://mlsanity.com/paper/2010.14406/) (2020)
- [Learning Latent Plans from Play](https://mlsanity.com/paper/1903.01973/) (2019)
- [HG-DAgger: Interactive Imitation Learning with Human Experts](https://mlsanity.com/paper/1810.02890/) (2018)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [Bridge Data: Boosting Generalization of Robotic Skills with Cross-Domain Datasets](https://mlsanity.com/paper/2109.13396/) (2021)
- [Masked Visual Pre-training for Motor Control](https://mlsanity.com/paper/2203.06173/) (2022)
- [Real-World Robot Learning with Masked Visual Pre-training](https://mlsanity.com/paper/2210.03109/) (2022)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control](https://mlsanity.com/paper/1812.00568/) (2018)
- [One-Shot Visual Imitation Learning via Meta-Learning](https://mlsanity.com/paper/1709.04905/) (2017)
- [DART: Noise Injection for Robust Imitation Learning](https://mlsanity.com/paper/1703.09327/) (2017)
- [Goal-conditioned Imitation Learning](https://mlsanity.com/paper/1906.05838/) (2019)
- [The Unsurprising Effectiveness of Pre-Trained Vision Models for Control](https://mlsanity.com/paper/2203.03580/) (2022)
- [Parrot: Data-Driven Behavioral Priors for Reinforcement Learning](https://mlsanity.com/paper/2011.10024/) (2020)
- [Human-in-the-Loop Imitation Learning using Remote Teleoperation](https://mlsanity.com/paper/2012.06733/) (2020)
- [Coarse-to-Fine Imitation Learning: Robot Manipulation from a Single Demonstration](https://mlsanity.com/paper/2105.06411/) (2021)
- [Q-attention: Enabling Efficient Learning for Vision-based Robotic Manipulation](https://mlsanity.com/paper/2105.14829/) (2021)
- [Pre-Training for Robots: Offline RL Enables Learning New Tasks from a Handful of Trials](https://mlsanity.com/paper/2210.05178/) (2022)
- [Learning and Retrieval from Prior Data for Skill-based Imitation Learning](https://mlsanity.com/paper/2210.11435/) (2022)
- [Task-Embedded Control Networks for Few-Shot Imitation Learning](https://mlsanity.com/paper/1810.03237/) (2018)
- [Interactive Imitation Learning in Robotics: A Survey](https://mlsanity.com/paper/2211.00600/) (2022)
- [Retrieval-Augmented Reinforcement Learning](https://mlsanity.com/paper/2202.08417/) (2022)
- [Imitation Learning by Estimating Expertise of Demonstrators](https://mlsanity.com/paper/2202.01288/) (2022)

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

- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [Data Quality in Imitation Learning](https://mlsanity.com/paper/2306.02437/) (2023)
- [Robometer: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons](https://mlsanity.com/paper/2603.02115/) (2026)
- [RAM: Retrieval-Based Affordance Transfer for Generalizable Zero-Shot Robotic Manipulation](https://mlsanity.com/paper/2407.04689/) (2024)
- [Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning](https://mlsanity.com/paper/2408.14037/) (2024)
- [DINOBot: Robot Manipulation via Retrieval and Alignment with Vision Foundation Models](https://mlsanity.com/paper/2402.13181/) (2024)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)
- [Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks](https://mlsanity.com/paper/2407.12998/) (2024)
- [R+X: Retrieval and Execution from Everyday Human Videos](https://mlsanity.com/paper/2407.12957/) (2024)
- [FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning](https://mlsanity.com/paper/2408.16944/) (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)
- [What Matters in Learning from Large-Scale Datasets for Robot Manipulation](https://mlsanity.com/paper/2506.13536/) (2025)
- [Curating Demonstrations using Online Experience](https://mlsanity.com/paper/2503.03707/) (2025)
- [Learning to Discern: Imitating Heterogeneous Human Demonstrations with Preference and Representation Learning](https://mlsanity.com/paper/2310.14196/) (2023)
- [On the Effectiveness of Retrieval, Alignment, and Replay in Manipulation](https://mlsanity.com/paper/2312.12345/) (2023)
- [Mechanistic Finetuning of Vision-Language-Action Models via Few-Shot Demonstrations](https://mlsanity.com/paper/2511.22697/) (2025)
- [Learning from Imperfect Demonstrations with Self-Supervision for Robotic Manipulation](https://mlsanity.com/paper/2401.08957/) (2024)
- [Offline Imitation Learning Through Graph Search and Retrieval](https://mlsanity.com/paper/2407.15403/) (2024)
- [Selective Sampling and Imitation Learning via Online Regression](https://mlsanity.com/paper/2307.04998/) (2023)
- [Guiding Data Collection via Factored Scaling Curves](https://mlsanity.com/paper/2505.07728/) (2025)
- [Data Retrieval with Importance Weights for Few-Shot Imitation Learning](https://mlsanity.com/paper/2509.01657/) (2025)
- [Uni-Skill: Building Self-Evolving Skill Repository for Generalizable Robotic Manipulation](https://mlsanity.com/paper/2603.02623/) (2026)
- [NOIR: Neural Signal Operated Intelligent Robots for Everyday Activities](https://mlsanity.com/paper/2311.01454/) (2023)
- [cVLA: Towards Efficient Camera-Space VLAs](https://mlsanity.com/paper/2507.02190/) (2025)
- [RT-Cache: Training-Free Retrieval for Real-Time Manipulation](https://mlsanity.com/paper/2505.09040/) (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)
- [Factorizing Diffusion Policies for Observation Modality Prioritization](https://mlsanity.com/paper/2509.16830/) (2025)

## Abstract (arXiv)

Enabling robots to learn novel visuomotor skills in a data-efficient manner remains an unsolved problem with myriad challenges. A popular paradigm for tackling this problem is through leveraging large unlabeled datasets that have many behaviors in them and then adapting a policy to a specific task using a small amount of task-specific human supervision (i.e. interventions or demonstrations). However, how best to leverage the narrow task-specific supervision and balance it with offline data remains an open question. Our key insight in this work is that task-specific data not only provides new data for an agent to train on but can also inform the type of prior data the agent should use for learning. Concretely, we propose a simple approach that uses a small amount of downstream expert data to selectively query relevant behaviors from an offline, unlabeled dataset (including many sub-optimal behaviors). The agent is then jointly trained on the expert and queried data. We observe that our method learns to query only the relevant transitions to the task, filtering out sub-optimal or task-irrelevant data. By doing so, it is able to learn more effectively from the mix of task-specific and offline data compared to naively mixing the data or only using the task-specific data. Furthermore, we find that our simple querying approach outperforms more complex goal-conditioned methods by 20% across simulated and real robotic manipulation tasks from images. See https://sites.google.com/view/behaviorretrieval for videos and code.