Data Retrieval with Importance Weights for Few-Shot Imitation Learning

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. Retrieval-based imitation learning addresses this by extracting relevant samples from large, widely available prior datasets to augment a limited demonstration dataset. To determine the relevant retrieval-based approaches 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, it has two shortcomings. 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 overcomes the bias of previous selection rules, and by using reasonable modeling parameters, IWR effectively smooths estimates using all data points. Across both simulation and real-world evaluations on the Bridge dataset, IWR consistently improves performance of existing retrieval-based methods, despite only requiring minor modifications.