Data Retrieval with Importance Weights for Few-Shot Imitation Learning
作者: Amber Xie, Rahul Chand, Dorsa Sadigh, Joey Hejna
分类: cs.RO, cs.AI
发布日期: 2025-09-01
备注: Conference on Robot Learning 2025
💡 一句话要点
提出重要性加权检索(IWR)方法，提升小样本模仿学习的数据检索质量。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 模仿学习 小样本学习 数据检索 重要性加权 高斯核密度估计
📋 核心要点
- 现有基于检索的模仿学习方法依赖于最近邻估计，易受噪声影响，且忽略了先验数据分布。
- 论文提出重要性加权检索(IWR)，通过高斯核密度估计目标数据和先验数据分布的比率，作为重要性权重。
- 实验表明，IWR在模拟和真实机器人数据集上，均能有效提升现有检索方法的性能。
📝 摘要（中文）
大规模机器人数据集推动了模仿学习的最新进展，但从较小的特定任务数据集学习对于在新环境和未见任务中的部署仍然至关重要。一种小样本模仿学习方法是基于检索的模仿学习，它从大型、广泛可用的先验数据集中提取相关样本，以扩充有限的演示数据集。为了确定来自先验数据集的相关数据，基于检索的方法最常计算先验数据点到目标数据集中点的潜在空间中的最小距离。虽然基于检索的方法已经证明了使用这种度量进行数据选择的成功，但我们证明了它等价于目标数据分布的高斯核密度估计(KDE)的极限。这揭示了先前工作中使用的检索规则的两个缺点。首先，它依赖于易受噪声影响的高方差最近邻估计。其次，它在检索数据时没有考虑先验数据的分布。为了解决这些问题，我们引入了重要性加权检索(IWR)，它使用高斯KDE估计重要性权重，即目标数据和先验数据分布之间的比率，用于检索。通过考虑概率比率，IWR试图减轻先前选择规则的偏差，并且通过使用合理的建模参数，IWR有效地使用所有数据点来平滑估计。在模拟环境和Bridge数据集的真实评估中，我们发现我们的方法IWR始终提高了现有基于检索的方法的性能，尽管只需要进行少量修改。
🔬 方法详解
问题定义：现有基于检索的模仿学习方法在小样本情况下，依赖于计算先验数据点到目标数据集中点的最小距离来选择数据。这种方法等价于高斯核密度估计的极限，存在两个主要问题：一是最近邻估计方差高，容易受到噪声干扰；二是忽略了先验数据的分布，导致检索结果存在偏差。
核心思路：论文的核心思路是利用重要性加权来修正检索过程中的偏差。具体来说，通过估计目标数据分布和先验数据分布之间的比率（即重要性权重），来指导数据的选择。这样既能考虑到目标数据的特征，又能避免受到先验数据分布的影响，从而提高检索的准确性和鲁棒性。
技术框架：IWR方法的整体框架如下：1. 数据编码：将目标数据集和先验数据集中的数据点编码到潜在空间中。2. 密度估计：使用高斯核密度估计(KDE)分别估计目标数据分布和先验数据分布。3. 重要性权重计算：计算每个先验数据点的重要性权重，即目标数据分布在该点的值与先验数据分布在该点的值的比率。4. 数据检索：根据重要性权重选择最相关的先验数据点，用于扩充目标数据集。
关键创新：IWR的关键创新在于引入了重要性权重来指导数据检索。与传统的基于距离的检索方法相比，IWR考虑了目标数据和先验数据的分布差异，从而能够更准确地选择与目标任务相关的数据。此外，使用高斯核密度估计可以有效地平滑噪声，提高估计的鲁棒性。
关键设计：IWR的关键设计包括：1. 高斯核密度估计：选择合适的高斯核函数和带宽参数，以获得准确的密度估计。带宽参数的选择对性能至关重要，需要根据具体数据集进行调整。2. 重要性权重截断：为了避免重要性权重过大或过小，可以对权重进行截断，例如设置一个上限和下限。3. 数据编码方式：潜在空间的表示方式会影响检索效果，可以使用预训练的编码器或者针对特定任务进行微调。
🖼️ 关键图片
📊 实验亮点
实验结果表明，IWR方法在模拟环境和真实机器人数据集（Bridge dataset）上均取得了显著的性能提升。与现有的基于检索的方法相比，IWR能够更准确地选择与目标任务相关的数据，从而提高模仿学习的性能。具体提升幅度取决于数据集和任务的复杂程度，但总体而言，IWR能够稳定地提升性能。
🎯 应用场景
该研究成果可应用于机器人模仿学习、自动驾驶、医疗诊断等领域。在这些领域中，往往难以获取大量的标注数据，而利用IWR方法可以有效地从已有的先验数据集中检索相关数据，从而提高模型的泛化能力和性能。尤其是在机器人领域，可以利用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.