质量胜于数量: 利用重要性加权检索改进机器人学习

机器人领域目前正面临一场“数据饥渴”危机。虽然深度学习带来了能力的巨大飞跃，但这些模型需要海量数据。在计算机视觉或 NLP 中，抓取互联网能提供数十亿个示例。然而在机器人技术中，数据必须物理收集——这是一个缓慢、昂贵且劳动密集的过程。

为了解决这个问题，研究人员通常转向少样本模仿学习 (Few-Shot Imitation Learning) 。 目标简单但宏大: 仅通过少量演示 (“目标”数据) ，并通过从海量的现有数据集中提取相关片段 (“先验”数据) 进行补充，来教会机器人一项新任务。这个过程被称为检索 (Retrieval) 。

但问题在于: 你如何决定海量堆积的数据中哪些对你的特定新任务实际有用？

标准方法依赖于简单的几何距离，本质上是在问: “哪些旧数据看起来最像我的新数据？”在论文 “Data Retrieval with Importance Weights for Few-Shot Imitation Learning” 中，斯坦福大学的研究人员认为这种直觉在数学上是有缺陷的。他们提出了一种名为重要性加权检索 (Importance Weighted Retrieval, IWR) 的新概率方法，将数据选择不仅仅视为几何问题，而是视为概率密度估计问题。

在这篇文章中，我们将拆解为什么旧的检索方式既嘈杂又有偏差，以及 IWR 如何利用重要性采样在模拟和现实世界的机器人环境中大幅提高性能。

现状: 通过最近邻进行检索

要理解这篇论文的贡献，我们首先需要看看机器人目前是如何“检索”记忆的。

想象一下你想教机器人拿起一个红色马克杯。你给它 5 次演示 (\(D_{target}\)) 。你还有一个包含 100,000 次以前的机器人交互 (\(D_{prior}\)) 的巨大硬盘，范围从打开抽屉到捡香蕉。

大多数最先进的方法——如行为检索 (Behavior Retrieval, BR) 或光流检索 (Flow Retrieval) ——遵循一个标准配方:

  1. 编码 (Encode) : 将所有图像/动作数据压缩到一个低维度的“潜在空间” (通常使用变分自编码器，即 VAE) 。
  2. 测量 (Measure) : 对于海量先验数据集中的每一个点，找到其到目标数据集中最近点的距离。
  3. 选择 (Select) : 保留距离最小的点。

在数学上，标准选择规则如下所示:

The standard retrieval rule based on minimum L2 distance.

这里，\(f_{\phi}\) 是编码器。我们在先验数据集 (\(D_{prior}\)) 中寻找样本，其到最近目标样本的欧几里得距离平方 (\(L2\)) 低于特定阈值 \(\zeta\)。

“最近”的问题

虽然直观，但研究人员指出这种“最近邻”方法存在两个主要的数学弱点:

  1. 高方差 (噪声) : 依赖单个最近邻在统计上是脆弱的。如果你的目标演示稍微有点嘈杂，或者潜在空间不是非常平滑，你可能会检索到几何上看起来很近但在语义上无关的数据。
  2. 偏差 (Bias) : 这种方法完全忽略了先验数据的分布。它问“这看起来像目标吗？”，但没有问“这在先验数据集中有多常见？”这导致了训练数据的分布偏斜。

研究人员意识到，标准的 L2 距离规则实际上只是概率密度估计的一种粗略近似——具体来说，它是带宽趋近于零时的高斯核密度估计 (KDE) 的极限。通过承认这种联系，他们可以用适当的统计工具替换这种粗略的近似。

解决方案: 重要性加权检索 (IWR)

IWR 从根本上将视角从距离转移到了概率 。 IWR 不问“这个点有多近？”，而是问“与先验分布相比，这个点属于目标分布的可能性有多大？”

该方法包含三个主要阶段，如下图所示:

IWR consists of three main steps: (A) Learning a latent space, (B) Estimating importance weights, and (C) Co-training the policy.

让我们分解这两个关键创新: 通过 KDE 进行平滑处理和重要性采样。

1. 利用核密度估计 (KDE) 进行平滑

IWR 的第一个升级是用高斯核密度估计 (KDE) 替换“最近邻”检查。

想象你的 5 个目标演示是纸上的点。“最近邻”方法在每个点周围画一个小圆圈。如果先验数据点落在圆圈里，就保留它。

相反，KDE 在每个点上放置一个平滑的“小山丘” (高斯分布) 并将它们相加。这创造了一个连续的概率地形。先验数据点根据其在这个地形上的高度进行评估。

Comparison of L2 distance vs IWR smoothing.

如上图 Figure 2 所示，这种差异很重要。

  • L2 距离 (上图) : 先验点 \(y\) 在几何上比 \(x\) 更接近单个目标点 \(z\)。因此，标准方法选择了 \(y\)。
  • IWR (下图) : 点 \(x\) 位于许多目标点聚集的区域。即使它不是离任何单个点最近的，它在目标分布 \(p_t\) 下具有更高的概率密度。IWR 正确地识别出 \(x\) 是更好的候选者。

该密度估计的数学公式为:

The Gaussian KDE formula used to estimate probability density.

该公式通过对所有数据点的高斯贡献进行平均来计算概率 \(p^{KDE}(z)\)，并通过带宽参数 \(h\) 和协方差 \(\Sigma\) 进行平滑处理。这为“好”数据的位置创建了一个方差更低、更稳健的估计。

2. 重要性权重的力量

第二个，也许是更关键的创新，是重要性采样 (Importance Sampling) 的应用。

在模仿学习中，我们希望最小化目标分布 (\(p_t\)) 上的损失。然而，我们正在从先验分布 (\(p_{prior}\)) 中检索样本。如果我们简单地抓取看起来像 \(p_t\) 的数据，我们就忽略了一个事实，即先验数据集有其自身的偏差 (例如，它可能有 1,000 个打开抽屉的片段，但只有 5 个举起杯子的片段) 。

为了纠正这一点，我们需要对数据进行加权。我们要基于重要性权重比率对数据进行采样:

\[ w = \frac{p_t(z)}{p_{prior}(z)} \]
  • 分子 (\(p_t\)) : 这个数据与我的新任务相关吗？
  • 分母 (\(p_{prior}\)) : 这个数据在旧数据集中是否代表性过高？

研究人员寻求满足以下期望，确保我们要训练的数据在数学上代表目标任务:

The Importance Sampling expectation equality.

通过使用 KDE 估计 \(p_t\) 和 \(p_{prior}\)，IWR 为每一份先验数据分配一个分数。然后它检索权重最高的数据。

IWR 的最终检索规则如下所示:

The IWR retrieval rule using the log-sum-exp of Gaussian kernels.

这个不等式选择重要性权重 (密度比) 超过阈值 \(\zeta\) 的点。注意对数内的求和——这就是 KDE 的“平滑”效应在起作用，同时考虑所有目标点，而不仅仅是最近的一个。

效果如何？实验结果

研究人员在模拟环境和现实世界的机器人任务中，将 IWR 与标准基线 (行为克隆、行为检索、光流检索和 SAILOR) 进行了对比测试。

环境

评估涵盖了各种各样的任务:

  1. Robomimic Square: 一个精确的组装任务。
  2. LIBERO: 包含 10 个不同厨房风格任务的基准套件。
  3. 现实世界 Bridge V2: 在玩具水槽中操作玉米、胡萝卜和茄子等物体。

The evaluation environments: Robomimic, LIBERO, and Real World Bridge tasks.

性能提升

结果是一致且显著的。在几乎所有类别中，用 IWR 的基于概率的指标替换标准 L2 距离指标都提高了成功率。

Table of results showing IWR consistently outperforming baselines.

在 Table 1 中，注意“现实世界任务” (玉米、胡萝卜、茄子) 。

  • 标准行为检索 (BR) 表现挣扎，在玉米任务上仅取得 2/20 的成功。
  • IWR 跃升至 9/20 。
  • 在长程 茄子任务 上，IWR 取得了 11/20 的完全成功，而次优方法 (BR) 仅达到 3/20 。

这表明 IWR 不仅仅是理论上的改进；它转化为了稳健的、物理的机器人行为。

为什么它更有效？深入探究

为了理解 为什么 IWR 会胜出，作者分析了究竟检索到了什么数据。

以“马克杯-布丁”任务为例。目标是将白色马克杯放在盘子上，布丁放在左边。先验数据集包含令人困惑的“干扰”任务，比如放下巧克力布丁，或者把马克杯放在右边。

Comparison of retrieval distributions between BR and IWR.

Figure 4 揭示了关键差异:

  • 左图 (任务) : 标准方法 (BR) 检索了大量“有害”数据 (红色条) ——可能是涉及看起来相似但操作不正确的物体的任务。IWR (左下) 显著过滤了这些数据，检索到的主要是“相关”或“混合” (部分有用) 的数据。
  • 右图 (时间步) : 标准检索通常过度采样轨迹的开始部分 (此时还没发生什么) 。IWR 在整个任务时间线上检索了平衡的分布。

因为 IWR 对先验分布 (\(p_{prior}\)) 进行了建模，它隐含地理解“在剧集开始时静坐”在先验数据集中是非常常见的。重要性权重中的分母 (\(1/p_{prior}\)) 惩罚了这些常见的、信息量少的帧，让独特的、与任务相关的动作脱颖而出。

通用性

IWR 最强大的特性之一是它对底层表示无关 。 无论你使用的是 VAE 潜在变量 (如行为检索) 、光流 (光流检索) 还是基于技能的嵌入 (SAILOR) ，你都可以在其之上应用 IWR。

Table showing IWR improves performance when added to SAILOR (SR) and Flow Retrieval (FR).

Table 2 表明，将 IWR 应用于现有方法 (SR-IWR 和 FR-IWR) 通常会提升它们的性能，证明了该数学方法普遍适用于基于检索的学习。

结论

从“大数据”到“智能数据”的转变是机器人技术的关键一步。这篇论文强调，随着我们更多地依赖从海量的、未经筛选的数据集中检索数据，我们的选择标准必须成熟。

启发式方法如“最近邻”在早期为我们提供了良好的服务，但它们带有隐性偏差且容易受噪声影响。通过将检索形式化为一个概率问题——具体来说是带有高斯核密度估计的重要性采样——IWR 提供了一种有原则的数据选择方法。

对学生和从业者的启示很明确:

  1. 几何 \(\neq\) 概率: 在潜在空间中“近”并不总是意味着“可能”或“有用”。
  2. 上下文很重要: 你不能忽略你从中检索数据的来源数据集的分布。
  3. 平滑有助益: 聚合所有可用演示的信息 (通过 KDE) 比信任单个最近邻更稳健。

随着机器人数据集继续增长到数百万条轨迹的规模，像 IWR 这样的方法将成为必不可少的过滤器，确保机器人从信号而非噪声中学习。