24年8月来自斯坦福和伯克利分校的论文“Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning”， 被CoRL’24录用。

人们正在收集越来越多的大型模仿学习数据集，目的是训练机器人的基础模型。然而，尽管数据选择在视觉和自然语言处理中至关重要，但在机器人领域，很少有研究质疑应该使用哪些数据来训练这些模型。这项工作研究如何权衡机器人基础模型预训练中机器人数据集的不同子集或“域”。具体来说，用分布上鲁棒优化 (DRO) 来最大化所有可能的下游域最坏情况性能。提出的方法 Re-Mix （极小极大优化重加权机器人数据集混合）解决了将 DRO 应用于机器人数据集时出现的各种挑战，包括动作空间的变化和不同数据集之间的动态。Re-Mix 采用早期停止（early stopping）、动作规范化和离散化来解决这些问题。通过对最大的开源机器人操作数据集 Open X-Embodiment 数据集进行大量实验，证明数据管理会对下游性能产生巨大影响。具体来说，在用于训练现有通用机器人策略（特别是 RT-X 模型）的数据集上，Re-Mix 学习的领域权重平均比统一权重高出 38%，比人类选择的权重高出 32%。

机器学习的许多突破都可以归功于“互联网规模”数据集，从 CLIP [1] 等视觉模型的开发，到由 Common Crawl 数据集 [2] 提供支持的基于 Transformer 语言模型最新进展。为了利用这一趋势，机器人技术领域的近期一些努力集中在收集 [3–6] 或汇集 [7] 大规模机器人数据集，目的是训练更高效的模仿学习策略。然而，从这些数据中学习尤其具有挑战性：机器人数据集是使用不同的机器人、环境、状态空间、动作空间和动态 [8] 收集的。例如，常用的 Bridge V2 数据集 [4] 在小型 WidowX 机器人上使用第三人称摄像头和笛卡尔增量控制空间，而在流行且体型更大的 Franka Panda 机器人上收集的许多数据集 [9–12] 则使用腕式摄像头 [3] 或关节空间动作 [13]。虽然接受这种异质性会迅速扩大可用训练数据的数量 [7]，但它也放大了一个基本问题的重要性：如何将这些原始的、异构的数据源整理成通用机器人策略的有效训练数据集？

机器人技术中的数据问题。机器人技术领域的几项最新研究主要集中在收集大型演示数据集，用于模拟 [20, 32, 33] 和现实世界 [3, 7, 34– 38] 中的模仿学习，以训练大规模机器人策略 [6, 25, 39, 40]。一般来说，这些研究以及其他数据收集对组合泛化影响的研究 [21, 41, 42] 表明，数据集构建的各个方面（例如场景和任务多样性）对下游策略泛化有直接影响。尽管有几项研究重点关注如何通过特定硬件 [43]和收集程序 [11, 21, 44] 收集数据，或提供有关数据收集的理论见解 [22]，但机器人技术领域很少有研究解决事后数据集选择和分析问题。这一点尤为重要，因为机器人数据集的数量和多样性正在增加，而如何训练有效使用所有收集数据的策略，这个结论却不那么明确 [3, 7, 25]。Baker [45] 训练分类器来预测数据质量，但需要人工注释，而这在扩展时是不切实际的。也许最相关的是基于检索的方法，这些方法对数据集进行子集化 [12, 46]，但这样做是基于一个先验目标任务规范，因此不适用于训练通用策略。

计算机视觉中的数据管理。计算机视觉数据集最初是手工制作和手动标记的 [14, 47]。然而，在保持质量的同时将数据集扩展到手工管理范围之外，对于提高性能至关重要 [1, 48]。值得注意的是，基于元数据-计数的平衡 [49]、嵌入 [19]、光流 [50] 和聚类 [51] 等过滤技术已被证明可以大大提高下游性能，尽管过滤掉了大量数据。将其推向极端，核心集选择方法使用“主动学习” [52, 53] 选择视觉数据集的极小子集，但由于计算要求过高，只能选择小规模数据集 [54, 55]。虽然从演示中学习涉及视觉，但其核心是动作预测。计算机视觉的数据管理技术只能以与动作无关的方式过滤状态动作轨迹——可能会删除数据集的有用部分。

自然语言处理中的数据管理。在对大量现实世界的文本源进行训练时，语言建模流水线采用了许多特定于文本的预处理步骤，包括元数据过滤、语言过滤、重复数据删除和毒性降低 [16, 18, 56, 57]。更高级的数据选择方法考虑对数据进行子集化以最大化下游性能，但使​​用诸如对嵌入文本进行 k-均值聚类之类的技术 [58, 59]。虽然这种聚类技术在机器人技术中也可能具有视觉信息——类似于计算机视觉中的管理工作——但它们不提供有关动作的信息。混合技术，例如使用极小极大优化域重加权 (DoReMi) [28]，使用鲁棒优化来平衡文本域，并基于优先训练的思想 [60–62]。

虽然早期的视觉和语言模型是在高度整理的学术数据集（如 ImageNet [14]）上训练的，但围绕数据选择的问题已经塑造了使用互联网规模数据的现代训练流程 [15–17]。例如，大语言模型的训练涉及多个数据过滤阶段 [18]。同样，大型视觉数据集（例如 LAION [19]）使用预训练模型（如 CLIP [1]）评估每个数据点的质量。因此，随着机器人数据集的不断扩展，可以预期机器人数据管理将变得同样重要。不幸的是，简单的过滤技术在机器人技术中往往不够用；无法应用 n-gram 过滤器，并且视觉嵌入无法捕捉 episodic 机器人数据的顺序性。

尽管已证明演示数据（例如动作质量 [20] 和视觉多样性 [3, 4, 21]）对下游性能至关重要，但机器人数据管理方法仍然有限。在模仿学习中，数据选择问题仅在理论上 [22, 23] 或简单的小规模设置 [24] 中得到描述。因此在实践中只能使用临时解决方案。例如，尽管 Open-X-Embodiment 数据集 (OpenX) [7] 由 60 多个单独的数据集组成，总计超过 200 万条机器人轨迹，但与之一起发布的 RT-X 模型仅基于 12 个数据集的混合进行训练，并根据专家直觉进行加权。最近发布的 Octo [25] 和 OpenVLA [26] 通才策略同样是在 OpenX 的一个子集上进行训练的，其中作者根据主观的“兴趣度”概念选择以何种采样权重包含哪些数据集。虽然由此产生的数据组合在实践中效果良好，但它们的整理需要广泛的领域知识和手动数据检查。这种临时选择策略不太可能扩展到用于训练机器人策略的快速增长的数据集 [3, 5, 27]。

如何才能自动整理大规模机器人数据集，以最大限度地提高跨域通用模仿学习策略的性能？虽然许多过滤技术并不直接适用于机器人技术，但可以借鉴语言建模的思想，根据模型的性能系统地优化训练数据混合。具体来说，DoReMi [28] 使用组分布性鲁棒优化 [29] 来最大限度地提高给定数据集中所有“域”策略的性能。在机器人技术的背景下，这样的“域”可以对应于单个数据集内的不同场景，例如，当考虑 Bridge V2 数据集 [4] 中的数据混合时，可以对应于不同的玩具厨房，或者在多数据集混合（如 OpenX 数据集）的情况下可以参考完整的机器人数据集。然而，由于机器人数据集的异质性，单纯地应用这些技术是行不通的。分布型稳健优化方法，可以最大限度地减少最坏情况的损失。动作空间及其分布的差异可能导致损失幅度在各个域的不平衡，从而导致某些域的权重比应有的要大。此外，机器人数据集较小，容易出现过拟合。这两个问题都会导致模型性能估计不佳，从而导致混合权重不佳。

考虑一般的模仿学习问题，其中给定一个由状态-动作轨迹 τ = (s/1, a/1, …, s/Ti, a/Ti) 组成的演示数据集 D = {τ/1, …, τ/n}。目标是学习一个参数化策略π/θ，该策略学习从状态到动作的映射π/θ：S → A。在实践中，这通常通过标准模仿学习算法，例如行为克隆（BC）来实现，方法是最小化策略下动作的预期负对数似然。

但是，数据集通常包含的信息不仅仅是状态-动作对。假设整体数据集 D 可以拆分为 k 个异构域 D/1 ，…，D/k 。这是一个普遍假设：虽然这些域可以是更大的组，例如来自 Open X-Embodiment 数据集 [7] 的具有不同具身的不同数据集，但它们也可以小到只有单个轨迹。此外，k 个域中的每一个域在状态空间 S、动作空间 A、转换动态或它们的分布方面都可能有所不同。事实上，在学习大型行为模型时，这种异构性对于访问更多不同数据源是必不可少的。这项工作用 Bridge 数据集 [4]——以不同的环境作为域，以及 Open-X-Embodiment 数据集 [7]——以不同的机器人具身作为不同的域。

工作的目标是学习一个权重向量 α，该向量指定所有域的概率分布，使得任何模型在根据 α 加权的域混合上训练时，都能在所有域上获得最大性能。与旨在为特定目标任务整理数据的数据检索问题不同，本文目标是整理数据集以进行有效的预训练或联合训练，而无需任何有关目标任务的先验知识。

在对大量机器人数据进行预训练时，希望策略能够推广到新的设置和任务，而不是掌握特定的目标任务。考虑到这一点，希望优化数据混合，从而产生 i) 在每个领域都表现尽可能好的模型，但 ii) 不会以牺牲另一个领域为代价而过拟合任何一个领域。分布上鲁棒优化 (DRO) 技术旨在解决同样的问题：学习最小化最坏情况训练损失 [29]（在模仿学习的情况下为 BC 损失）的模型，跨领域 D/1，…，D/k。

在这一目标下，α 会增加损失值较高的域权重，强调最难的域。然而，在实践中，可能对仅拟合损失较高的域不感兴趣。例如，具有用于拧开瓶盖的复杂多模态旋转运动的机器人数据集，可能始终比简单的拾取数据集具有更高的 BC 损失。因此，标准的鲁棒优化技术最终可能会忽略后者。相反，与先前的工作 [28, 52, 63] 一样，考虑学习的策略 π/θ 和参考策略 π/ref 之间的损失差异，后者经过训练收敛于域权重的初始猜测，通常假设与每个域的大小成比例，即数据的均匀采样。此差异称为超额损失（excess loss），并将其用于鲁棒优化。与之前一样，这将降低策略适合域的权重，因为它可以实现与参考模型相似的损失。然而，由于超额损失的相对性，它至关重要地还降低了难拟合域的权重（即它们具有较高的策略损失和较高的参考损失）。例如，这可能发生在存在次优操作的情况下。因此，只有具有高超额损失的域（即策略可以改进以匹配参考模型）才会被增加权重，因为选择 α 是为了最大化超额总损失。

不幸的是，直接使用鲁棒优化学习的模型通常表现出更差的整体性能 [64, 65]，因为它们专注于最小化最坏情况损失而不是平均损失。或者，可以像 Xie 那样 [28] 将学习的 α 向量用于下游训练。这提供了一组可重复使用的权重，可用于训练不同的策略，而无需进行鲁棒优化。
Re-Mix 使用组分布鲁棒性来确定数据混合的权重 [28]，然后可用于策略训练，并结合以上关键设计考虑，解决不平衡损失、连续损失和过拟合的问题。Re-Mix 仅返回数据混合 α 的权重，而不是最终策略。这是为了将数据管理问题与策略训练问题分离。运行 Re-Mix 后，得到的权重可用于学习不同类型（即扩散）或更大规模的策略。

阶段 1：动作预处理。对具有不同动作空间和动态的每个域 D/k 分别应用高斯正则化，然后通过分箱将动作离散化。
阶段 2：参考模型训练。接下来，在域 D/1… , D/k 的均匀混合上训练离散参考模型 π/ref ，其中每个域的权重与其大小成比例。根据验证损失选择最终参考模型检查点，并使用它来估计每个域的超额损失。
阶段 3：组分布性鲁棒优化。通过一个离散策略 π/θ 的鲁棒优化来学习域权重 α。为了更新 α，按照 [29] 的方法，在每个训练步，对 α 执行一步指数梯度上升，然后对 θ 执行域加权梯度下降。得到的 α 值会增加仍然可以改进的域权重，同时降低那些不重要或不可能拟合域的权重。这意味着 Re-Mix 直接根据动作过滤数据，而不像视觉和语言中的其他技术那样仅根据嵌入进行过滤 [59, 69]。优化鲁棒方程，使其步骤数与参考模型相同。
阶段 4：策略训练的数据加权。在对超额损失进行鲁棒优化之后，在训练过程中取 α 的平均值，将其表示为 α ̄。然后，可以使用该 α ̄ 值重加权不同域，甚至子集数据集以进行策略训练。实际上，这意味着可以在不同配置下重复使用权重进行多次训练。例如，Re-Mix 使用离散动作，但使用扩散模型来训练最终策略，这在经验上已证明效果良好 [25, 30]。

旨在评估各种精选的预训练数据混合质量，以供下游策略学习使用。为此，在精选的数据集上联合训练通用目标条件策略。由于无法访问用于收集训练的数据集机器人设置，因此构建自己的 WidowX 和 Franka 机器人评估设置。不幸的是，仅在预训练数据上训练的策略无法零样本推广到分布外设置。为了解决这个问题，遵循先前的研究 [3、12、46、70]，并在少量领域内数据（3 个代表性任务各 25 个演示）上联合训练策略，并以 5% 的小权重添加到最终的训练混合中。然后，评估与联合训练数据相关的分布外任务策略以测试泛化。联合训练的结果是，所有策略都实现非零成功率。然而，域内数据集足够小，以至于预训练数据组合的质量仍然对评估结果有显著影响，为数据管理方法提供良好的试验台。所有模型都在现实世界中进行评估，每个任务进行 10 次试验，累计超过 500 次现实世界试验。

对于所有策略，用 ResNet 50 图像编码器 [71]。对于 Re-Mix 参考模型和 Group DRO 优化，用离散 MLP 动作头。对于所有最终策略，用 [4, 25, 72] 中的扩散头，Google Cloud TPU 研究资助提供的 TPU 对所有模型进行 400,000 个梯度步的训练。

架构。借用 [4] 中的架构，并做了一些小改动。策略将两个连续帧的历史记录和一个目标图像作为输入，并通过 DDPM [80] 输出一系列动作。
首先，对所有图像进行预处理，使其介于 -1 和 1 之间。然后，逐通道连接目标图像和包含 (x, y) 空间中每个像素位置的网格，该网格也标准化为 -1 和 1 之间。然后将图像输入到 ResNet 50 编码器，该编码器对输出采用全局平均池化，以获得每个图像的 512 维表示。然后将两个图像表示连接起来并输入到扩散动作预测头。
超参。用余弦衰减学习率进度，初始学习率为 0.0002。训练所有模型 400K 步并评估最终检查点，但 Bridge 10% 子集除外，在 200K 步后表现更好。Bridge 和 RTX 之间存在一些差异，这些差异是出于计算原因——在扩展到 RTX 之前，在 Bridge 数据集上进行更快的迭代。也确实为 Bridge 保留了纵横比，因此图像输入大小不同，但对于 RTX 并没有遵循 Octo Model Team [25] 的做法。在整个时间范围和目标图像上一致地对所有图像应用数据增强（这意味着目标图像和每个示例的所有过去图像都应用了相同的增强）。用随机调整大小裁剪、亮度、对比度和色调随机化。对于 Bridge 的 SSP 中 k 均值，将 k = 32，等于用于 Re-Mix 的域数。

如图所示根据不同方法对 RT-X 混合数据集进行加权的结果。对于 WidowX 机器人，考虑四个任务，测试对 1) 未见过目标的泛化能力：“胡萝卜到架子”、“OOD 杯”，2) 未见过的初始条件：“叉子到架子”，以及 3) 目标位置的干扰物“立方体到盘子”。

同样，对于 Franka Panda 机器人，考虑两个任务，测试对 1) 未见过的初始条件“笔在杯中”和 2) RT-X 数据中未见的动作“翻转碗”的泛化能力。此外，Panda 机器人使用 Robotiq 2F-85 夹持器，这在任何 RT-X 数据集中都不存在。请注意，对于 RT-X 混合，在 WidowX 和 Franka 数据上共同训练同一个模型。正如预期的那样，人类专家为 RT-X 模型选择的域权重平均比朴素的均匀采样基线高出 6%。更有趣的是，根据 Re-Mix 加权的数据集平均比均匀加权高出 38%，令人惊讶的是，平均比人类策划的权重高出 32%。Re-Mix 更好地推广到 Bridge 中未见过的物体和位置。Re-Mix 在“翻碗”任务中表现特别好，这可能是因为它增加 OpenX 中相对较小的“Toto”数据集的权重，其中包含与翻碗任务类似的倒水动作，尽管 Pen-in-cup 在所有模型中的表现都相似。假设这是因为它比其他任务需要更少的泛化，因此较少依赖于预训练混合。