ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

作者: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

许可协议

• 开放获取
• 论文: CC BY 4.0

以下内容均依据论文全文，按其 CC BY 4.0 许可协议使用。

一句话总结

ReMix 对策略的超额损失使用组分布鲁棒优化，学习每个机器人数据集应采样多少，在真实机器人上胜过均匀权重和 RT-X 专家权重。

为何重要

在汇集的机器人数据上进行预训练的团队，可以用根据数据计算出的权重取代人工调整的采样权重，并且在报告的测试中可以舍弃 75% 的数据而性能损失有限。学到的权重还表明应该降低哪些数据集的权重及其原因，例如自主收集的 Kuka 数据，这有助于决定收集或保留哪些数据。

问题

Open X-Embodiment 等大型机器人数据集汇集了来自不同机器人、相机、动作空间和控制器的数据，但用于训练 RT-X、Octo 和 OpenVLA 等通用策略的数据配比都是人工选择的。来自视觉和语言领域的数据筛选方法迁移效果不佳，因为 n-gram 并不适用，而视觉嵌入忽略了动作以及回合的序列特性。

提出的方法

ReMix 将数据整理表述为一个最小-最大问题：离散策略最小化其相对于参考模型的超额行为克隆损失，而领域权重通过指数梯度上升进行更新以使该损失最大化，从而把权重分给策略仍有提升空间的领域。为使不同机器人之间的损失可比，动作按领域进行高斯归一化并离散化到若干区间中，参考检查点则通过早停在其对任何领域过拟合之前选出。随后，训练过程中取平均的权重被重新用于训练单独的扩散策略或对数据进行子采样。

关键见解

• 使用相对于参考策略的超额损失而不是原始损失，可以同时降低已经拟合好的领域和难以拟合的领域的权重。
• 动作尺度的差异会使鲁棒优化产生偏差：在一个使用边界归一化的玩具测试中，94.3% 的权重被分给了一个纯随机动作的领域；使用逐领域的高斯归一化时，真实数据获得了 84.2% 的权重。
• 连续的 L1/L2 动作损失会受到离群值和多模态动作的影响，因此即使最终策略是连续扩散模型，离散化动作也能得到更可靠的权重。
• 机器人数据集（约 10k-100k 次演示）规模较小，足以使参考模型过拟合；使用早停点之后 50K 步的检查点，会使性能平均下降超过 15%。

主要结果

• 在 RT-X 数据配比上，使用 ReMix 权重训练的策略在四项 WidowX 和两项 Franka 真实世界任务上平均比均匀权重高 38%，比 RT-X 人类专家权重高 32%；专家权重比均匀权重高 6%。
• 在 RT-X 配比的 25% 子集上，ReMix 保持了性能，而人工权重下降超过 10%；ReMix 保留了 72% 的 UR5 数据和 12% 的 Kuka 数据，而人工权重分别为 30% 和 24%。
• ReMix 将小型 Toto 数据集的占比从 3.42%（均匀）提高到 16.3%，并将 Kuka 从 24.9% 降至 12.1%。
• 在完整的 Bridge V2 数据集上，ReMix 与均匀权重之间没有显著差异；在 25% 的 Bridge 子集上，ReMix 优于均匀权重和 SSP 图像嵌入基线。
• 在带有干扰物体的 WidowX 方块任务上，ReMix 权重成功 7/10 次，而人工权重和均匀权重均为 0/10。
• 在仿真中（按操作员划分的 Robomimic NutAssemblySquare），在 50% 子采样下 ReMix 得分 77/100，均匀权重为 53/100；在 25% 下为 59/100 对 39/100。

局限

• 作者指出，真实世界评估只涵盖 WidowX 和 Franka 两种机械臂，更多本体形态留待未来工作。
• ReMix 会提高动作分布异常、多模态的数据集（例如 Toto）的权重，作者表示这并不一定是理想的。
• 计算权重需要在完整数据上训练两次，一次用于参考模型，一次用于鲁棒优化。
• 仅在 Bridge V2 和 OpenX 的 RT-X 子集上测试，未在超过 2M 个回合的完整 OpenX 数据集上测试。第 4.1 节称 Bridge V2 有 50k 次演示，而附录 B.2 称从完整数据集的 60K 条轨迹中使用了约 45,000 条。
• 只在预训练数据上训练的策略无法零样本泛化到作者的实验设置，因此每项评估都依赖于每项任务使用 25 次领域内演示进行协同训练。

对领域的影响

入围 CoRL 2024 杰出论文奖决选，被引用 74 次。作者在 github.com/jhejna/remix 发布了代码。论文还公布了为 RT-X 配比学到的领域权重（表 2）以及 32 个 Bridge V2 领域的权重（表 5）。

摘要

为了训练机器人基础模型，人们正在收集规模越来越大的模仿学习数据集。然而，尽管数据选择在视觉和自然语言处理中极为重要，机器人领域却很少有工作探讨这类模型究竟应该在什么数据上训练。在这项工作中，我们研究了在机器人基础模型预训练中，应如何为机器人数据集的不同子集或「领域」分配权重。具体而言，我们使用分布鲁棒优化（DRO）来最大化所有可能下游领域中的最差情况性能。我们的方法 Re-Mix 解决了将 DRO 应用于机器人数据集时出现的各种挑战，包括不同数据集之间动作空间和动力学的差异。Re-Mix 采用早停、动作归一化和离散化来应对这些问题。通过在最大的开源机器人操作数据集 Open X-Embodiment 上进行大量实验，我们证明数据整理会对下游性能产生巨大影响。具体而言，在用于训练现有通用机器人策略（即 RT-X 模型）的数据集上，Re-Mix 学到的领域权重平均比均匀权重高 38%，比人工选择的权重高 32%。

链接

• arXiv 2408.14037: https://arxiv.org/abs/2408.14037
• DOI 10.48550/arXiv.2408.14037: https://doi.org/10.48550/arXiv.2408.14037
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/hejna25a/hejna25a.pdf
• 代码: https://github.com/jhejna/remix