Data Retrieval with Importance Weights for Few-Shot Imitation Learning

📄 arXiv: 2509.01657v1 📥 PDF

作者: Amber Xie, Rahul Chand, Dorsa Sadigh, Joey Hejna

分类: cs.RO, cs.AI

发布日期: 2025-09-01

备注: Conference on Robot Learning 2025

💡 一句话要点

提出重要性加权检索(IWR)方法，提升少样本模仿学习中数据检索的性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 少样本学习 模仿学习 数据检索 重要性加权 高斯核密度估计

📋 核心要点

  1. 现有基于检索的模仿学习方法依赖于最近邻估计，易受噪声影响，且忽略了先验数据分布。
  2. 论文提出重要性加权检索(IWR)方法，通过高斯核密度估计目标数据和先验数据分布的比率，作为重要性权重。
  3. 实验表明，IWR在模拟环境和真实机器人数据集上，均能有效提升现有检索方法的性能。

📝 摘要（中文）

大规模机器人数据集推动了模仿学习的最新进展，但从较小的特定任务数据集学习对于在新环境和未见任务中的部署仍然至关重要。一种少样本模仿学习方法是基于检索的模仿学习，它从大型、广泛可用的先验数据集中提取相关样本，以扩充有限的演示数据集。为了确定来自先验数据集的相关数据，基于检索的方法最常计算先验数据点到目标数据集中点的潜在空间中的最小距离。虽然基于检索的方法已经证明了使用这种度量进行数据选择的成功，但我们证明了它等价于目标数据分布的高斯核密度估计(KDE)的极限。这揭示了先前工作中使用的检索规则的两个缺点。首先，它依赖于易受噪声影响的高方差最近邻估计。其次，它在检索数据时没有考虑先验数据的分布。为了解决这些问题，我们引入了重要性加权检索(IWR)，它使用高斯KDE估计重要性权重，即目标数据和先验数据分布之间的比率，用于检索。通过考虑概率比率，IWR试图减轻先前选择规则的偏差，并且通过使用合理的建模参数，IWR有效地使用所有数据点来平滑估计。在模拟环境和Bridge数据集的真实评估中，我们发现我们的方法IWR始终提高现有基于检索的方法的性能，尽管只需要进行少量修改。

🔬 方法详解

问题定义：论文旨在解决少样本模仿学习中，如何从大量先验数据集中检索到与目标任务最相关的数据，以提升模仿学习的性能。现有方法主要基于最近邻搜索，容易受到噪声数据的影响，并且没有充分利用先验数据的分布信息。

核心思路：论文的核心思路是利用重要性加权的思想，通过估计目标数据分布和先验数据分布的比值，作为检索数据时的权重。这样可以减轻由于先验数据分布不均匀带来的偏差，并利用所有数据点进行平滑估计，降低噪声的影响。

技术框架：IWR方法的整体框架如下：1) 使用高斯核密度估计(KDE)分别估计目标数据集和先验数据集的概率密度函数。2) 计算每个先验数据点的重要性权重，即目标数据概率密度与先验数据概率密度的比值。3) 在检索数据时，使用计算得到的重要性权重对先验数据进行加权，选择权重较高的样本。

关键创新：论文的关键创新在于引入了重要性加权机制，将数据检索问题转化为概率密度比值估计问题。通过考虑目标数据和先验数据的分布差异，能够更准确地选择与目标任务相关的数据，从而提升模仿学习的性能。

关键设计：IWR方法中，高斯核密度估计的带宽参数是关键的设计。论文中提到使用合理的建模参数，可以有效地平滑估计，降低噪声的影响。此外，重要性权重的计算方式也至关重要，需要保证数值稳定性，避免出现除零错误。

📊 实验亮点

实验结果表明，IWR方法在模拟环境和真实机器人数据集（Bridge dataset）上均取得了显著的性能提升。与现有的基于检索的方法相比，IWR能够更准确地选择与目标任务相关的数据，从而提升模仿学习的性能。具体提升幅度在不同任务和数据集上有所不同，但总体上均优于基线方法。

🎯 应用场景

该研究成果可应用于机器人模仿学习、自动驾驶、游戏AI等领域。在这些领域中，往往难以获取大量的目标任务数据，而存在大量的先验数据。通过IWR方法，可以有效地利用这些先验数据，提升模型在目标任务上的泛化能力和学习效率。未来，该方法可以进一步扩展到其他机器学习任务中，例如迁移学习、领域自适应等。

📄 摘要（原文）

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.