[2606.09758] Difference-Aware Retrieval Policies for Imitation Learning

Abstract: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 this https URL.

Published as a conference paper at ICLR 2026 D IFFERENCE -AWARE R ETRIEVAL P OLICIES FOR I MITATION L EARNING Quinn Pfeifer1 , Ethan Pronovost1 , Paarth Shah2 , Khimya Khetarpal3,4 , Siddhartha Srinivasa1 , Abhishek Gupta1,2 1 Paul G. Allen School of Computer Science & Engineering, University of Washington 2 Toyota Research Institute…

saved by

• David Yan

related reading

• Successor Features for Transfer in Reinforcement Learningproceedings.neurips.cc
• Ch. 21 - Imitation Learningunderactuated.mit.edu
• State of Robot Learning, December 2025vedder.io
• Reinforcement Learning via Implicit Imitation Guidancearxiv.org
• Should I Use Offline RL or Imitation Learning? – The Berkeley Artificial Intelligence Research Blogbair.berkeley.edu
• Learning to Imitate | SAIL Blogai.stanford.edu
• Just Ask for Generalization | Eric Jangevjang.com
• Behavioral cloning mysteryseohong.me
• 1011.0686arxiv.org
• Generative Adversarial Imitation Learningarxiv.org
• Generative Adversarial Imitation Learning.pdfcs.stanford.edu
• InfoGAIL: Interpretable Imitation Learning from Visual Demonstrationsarxiv.org