Data Retrieval with Importance Weights for Few-Shot Imitation Learning

2025-09-01

AI Summary

Retrieval-based imitation learning suffers from high estimation variance and noise sensitivity in few-shot settings, while neglecting the underlying prior data distribution. To address these issues, we propose an importance-weighted retrieval mechanism that reweights retrieved samples by the ratio of the target distribution to the prior distribution—enabling low-bias, robust data selection. Theoretically, we show that conventional nearest-neighbor retrieval corresponds to the limiting case of Gaussian kernel density estimation (KDE), elucidating its inherent high variance and distributional mismatch. Our work is the first to incorporate probabilistic ratio modeling into retrieval and efficiently compute weights via KDE. We integrate this mechanism into a retrieval-augmented end-to-end imitation learning framework. Experiments across multiple simulated environments and the real-world Bridge robot dataset demonstrate significant performance gains. The method requires only lightweight modifications to existing pipelines, delivering consistent improvements—validating both its effectiveness and broad applicability.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Intelligent Robots: State Estimation
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Querying, indexing, and retrieval in Web-related graphs

Abstract

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.

Problem

• Addressing high-variance nearest neighbor noise in retrieval-based imitation learning
• Mitigating bias by accounting for prior data distribution during retrieval
• Improving few-shot imitation learning through importance-weighted data selection

Innovation

• Importance Weighted Retrieval using Gaussian KDEs
• Estimates target-prior data distribution probability ratio
• Smooths retrieval estimates with all data points