# Data Retrieval with Importance Weights for Few-Shot Imitation Learning

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

## Key facts

- Cited by 10 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 28 of its 32 references that are in the map.
- First posted on arXiv 1 Sep 2025.

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [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)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [BridgeData V2: A Dataset for Robot Learning at Scale](https://mlsanity.com/paper/2308.12952/) (2023)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations](https://mlsanity.com/paper/2310.17596/) (2023)
- [Scaling Up and Distilling Down: Language-Guided Robot Skill Acquisition](https://mlsanity.com/paper/2307.14535/) (2023)
- [The Surprising Effectiveness of Representation Learning for Visual Imitation](https://mlsanity.com/paper/2112.01511/) (2021)
- [ALOHA Unleashed: A Simple Recipe for Robot Dexterity](https://mlsanity.com/paper/2410.13126/) (2024)
- [Telescoping Density-Ratio Estimation](https://mlsanity.com/paper/2006.12204/) (2020)
- [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)
- [Density Ratio Estimation via Infinitesimal Classification](https://mlsanity.com/paper/2111.11010/) (2021)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [Imitation Learning by Estimating Expertise of Demonstrators](https://mlsanity.com/paper/2202.01288/) (2022)
- [Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning](https://mlsanity.com/paper/2408.14037/) (2024)
- [Sim-and-Real Co-Training: A Simple Recipe for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/2503.24361/) (2025)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)
- [Experience Replay with Likelihood-free Importance Weights](https://mlsanity.com/paper/2006.13169/) (2020)
- [Robot Data Curation with Mutual Information Estimators](https://mlsanity.com/paper/2502.08623/) (2025)
- [Importance Sampling Policy Evaluation with an Estimated Behavior Policy](https://mlsanity.com/paper/1806.01347/) (2018)
- [FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning](https://mlsanity.com/paper/2408.16944/) (2024)
- [Learning to Discern: Imitating Heterogeneous Human Demonstrations with Preference and Representation Learning](https://mlsanity.com/paper/2310.14196/) (2023)
- [How to Leverage Diverse Demonstrations in Offline Imitation Learning](https://mlsanity.com/paper/2405.17476/) (2024)

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

- [Robometer: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons](https://mlsanity.com/paper/2603.02115/) (2026)
- [Mechanistic Finetuning of Vision-Language-Action Models via Few-Shot Demonstrations](https://mlsanity.com/paper/2511.22697/) (2025)
- [Tether: Autonomous Functional Play with Correspondence-Driven Trajectory Warping](https://mlsanity.com/paper/2603.03278/) (2026)
- [Grounding Robot Generalization in Training Data via Retrieval-Augmented VLMs](https://mlsanity.com/paper/2603.11426/) (2026)
- [Closed Loop Dynamic Driving Data Mixture for Real-Synthetic Co-Training](https://mlsanity.com/paper/2605.21372/) (2026)
- [TSD: A Physics-Inspired Trajectory Saliency Detector for Efficient Imitation Learning](https://mlsanity.com/paper/2606.23371/) (2026)
- [Domain Arithmetic: One-Shot VLA Adaptation under Environmental Shifts](https://mlsanity.com/paper/2607.00666/) (2026)
- [Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF](https://mlsanity.com/paper/2607.07693/) (2026)
- [Hierarchical Skill Retrieval for Data-Efficient Adaptation of Vision-Language-Action Models](https://mlsanity.com/paper/2608.24042/) (2026)
- [ReWeight: Leveraging Human Data for VLA Post-Training via Demonstration Retrieval and Sample Weighting](https://mlsanity.com/paper/2609.13851/) (2026)

## Abstract (arXiv)

While large-scale robot datasets have propelled recent progress in imitation learning, learning from smaller task specific datasets remains critical for deployment in new environments and unseen tasks. One such approach to few-shot imitation learning is retrieval-based imitation learning, which extracts relevant samples from large, widely available prior datasets to augment a limited demonstration dataset. To determine the relevant data from prior datasets, retrieval-based approaches most commonly calculate a prior data point's minimum distance to a point in the target dataset in latent space. While retrieval-based methods have shown success using this metric for data selection, we demonstrate its equivalence to the limit of a Gaussian kernel density (KDE) estimate of the target data distribution. This reveals two shortcomings of the retrieval rule used in prior work. First, it relies on high-variance nearest neighbor estimates that are susceptible to noise. Second, it does not account for the distribution of prior data when retrieving data. To address these issues, we introduce Importance Weighted Retrieval (IWR), which estimates importance weights, or the ratio between the target and prior data distributions for retrieval, using Gaussian KDEs. By considering the probability ratio, IWR seeks to mitigate the bias of previous selection rules, and by using reasonable modeling parameters, IWR effectively smooths estimates using all data points. Across both simulation environments and real-world evaluations on the Bridge dataset we find that our method, IWR, consistently improves performance of existing retrieval-based methods, despite only requiring minor modifications.