Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning

越来越多的模仿学习数据集被收集，旨在训练机器人基础模型。然而，尽管在视觉和自然语言处理中数据选择至关重要，但在机器人领域，很少有研究探讨这些模型应该实际训练哪些数据。在这项工作中，我们研究了如何对机器人数据集的不同子集或“领域”进行加权，以进行机器人基础模型的预训练。具体而言，我们使用分布鲁棒优化（DRO）来最大化所有可能的下游领域的最坏情况表现。我们的方法 Re-Mix 解决了将 DRO 应用于机器人数据集时出现的各种挑战，包括不同数据集之间行动空间和动态的变异性。Re-Mix 采用了早期停止、行动标准化和离散化等方法来应对这些问题。通过对最大的开源机器人操作数据集 Open X-Embodiment 数据集进行广泛的实验，我们证明了数据策划对下游性能的影响可能会超出预期。具体而言，Re-Mix 学习到的领域权重平均比均匀权重高出 38％，比人为选择的权重在用于训练现有通用机器人策略的数据集上高出 32％，特别是 RT-X 模型。

本文试图解决的问题是，在机器人学习中，如何选择最佳的数据集子集用于训练机器人基础模型的预训练。同时，本文还试图解决在应用分布式鲁棒优化（DRO）到机器人数据集中时所面临的挑战。

本文提出了一种名为Re-Mix的方法，使用分布式鲁棒优化来最大化在所有可能的下游领域中的最差情况下的性能。Re-Mix通过早期停止、动作归一化和离散化等技术来解决在机器人数据集中应用DRO时所面临的各种挑战。

本文在Open X-Embodiment数据集上进行了大量实验，证明了数据筛选对下游性能的影响。Re-Mix学习的领域权重比均匀权重平均提高了38％，比人工选择的权重平均提高了32％。本文提供了开源代码。

最近在这个领域中，还有一些相关的研究。例如，论文“Learning to Learn from Simulation: Faster Improvement with Transfer-Enhanced Reinforcement Learning”和“Learning to Adapt Reinforcement Learning Algorithms to New Environments”等。