# Difference-Aware Retrieval Policies for Imitation Learning

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

## Key facts

- Cited by 1 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 27 of its 41 references that are in the map.
- First posted on arXiv 8 Jun 2026.

## 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, 27 of 27)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://mlsanity.com/paper/2406.02523/) (2024)
- [MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations](https://mlsanity.com/paper/2310.17596/) (2023)
- [The Surprising Effectiveness of Representation Learning for Visual Imitation](https://mlsanity.com/paper/2112.01511/) (2021)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Lipschitz Continuity in Model-based Reinforcement Learning](https://mlsanity.com/paper/1804.07193/) (2018)
- [Keypoint Action Tokens Enable In-Context Imitation Learning in Robotics](https://mlsanity.com/paper/2403.19578/) (2024)
- [In-Context Imitation Learning via Next-Token Prediction](https://mlsanity.com/paper/2408.15980/) (2024)
- [Feedback in Imitation Learning: The Three Regimes of Covariate Shift](https://mlsanity.com/paper/2102.02872/) (2021)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)
- [Instant Policy: In-Context Imitation Learning via Graph Diffusion](https://mlsanity.com/paper/2411.12633/) (2024)
- [FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning](https://mlsanity.com/paper/2408.16944/) (2024)
- [Learning Smooth Humanoid Locomotion through Lipschitz-Constrained Policies](https://mlsanity.com/paper/2410.11825/) (2024)
- [CCIL: Continuity-based Data Augmentation for Corrective Imitation Learning](https://mlsanity.com/paper/2310.12972/) (2023)
- [L2C2: Locally Lipschitz Continuous Constraint towards Stable and Smooth Reinforcement Learning](https://mlsanity.com/paper/2202.07152/) (2022)
- [REGENT: A Retrieval-Augmented Generalist Agent That Can Act In-Context in New Environments](https://mlsanity.com/paper/2412.04759/) (2024)
- [SEABO: A Simple Search-Based Method for Offline Imitation Learning](https://mlsanity.com/paper/2402.03807/) (2024)
- [Bayesian Gaussian mixture model for robotic policy imitation](https://mlsanity.com/paper/1904.10716/) (2019)
- [The Unreasonable Effectiveness of Discrete-Time Gaussian Process Mixtures for Robot Policy Learning](https://mlsanity.com/paper/2505.03296/) (2025)

## Built on by (most cited first, 1 of 1)

- [Value-Preserving Action Quantization](https://mlsanity.com/paper/2610.17252/) (2026)

## Abstract (arXiv)

Parametric imitation learning via behavior cloning can suffer from poor generalization to out-of-distribution states due to compounding errors during deployment. We show that reusing the training data during inference via a semi-parametric retrieval-based imitation learning approach can alleviate this challenge. We present Difference-Aware Retrieval Policies for Imitation Learning (DARP), a semi-parametric retrieval-based imitation learning approach that addresses this limitation by reparameterizing the imitation learning problem in terms of local neighborhood structure rather than direct state-to-action mappings. Instead of learning a global policy, DARP trains a model to predict actions based on k-nearest neighbors from expert demonstrations, their corresponding actions, and the relative distance vectors between neighbor states and query states. DARP requires no additional assumptions beyond those made for standard behavior cloning -- it does not require additional data collection, online expert feedback, or task-specific knowledge. We demonstrate consistent performance improvements of 15-46% over standard behavior cloning across diverse domains, including continuous control and robotic manipulation, and across different representations, including high-dimensional visual features. Code and demos are available at https://weirdlabuw.github.io/darp-site/.