Re-Mix:大规模模仿学习数据混合的优化
机器人学
2025-02-24 v1 机器学习
摘要
越来越多的大规模模仿学习数据集正在收集,以训练用于机器人的基础模型。尽管数据选择在视觉和自然语言处理领域已被确立为关键因素,但机器人领域的工作仍很少质疑这些模型应实际训练于什么数据。本文调查如何为机器人基础模型的预训练对不同数据子集或“域”进行加权。具体而言,我们使用分布式鲁棒优化(DRO)来最大化跨所有可能下游域的最小化性能。我们的方法 Re-Mix 解决了在机器人数据集上应用 DRO 时出现的各种挑战,包括不同数据集中动作空间和动力学的差异。Re-Mix 采用提前停止、动作归一化和离散化来抵消这些问题。通过在最大开源机器人操作数据集(Open X-Embodiment 数据集)上的大量实验,我们展示数据精选对下游性能具有决定性影响。具体而言,Re-Mix 学习到的域权重在平均性能上显著优于均匀权重,提升 38%;在用于训练现有通用机器人策略模型(RT-X 模型)的数据集上,优于人工选择的权重提升 32%。
引用
@article{arxiv.2408.14037,
title = {Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning},
author = {Joey Hejna and Chethan Bhateja and Yichen Jiang and Karl Pertsch and Dorsa Sadigh},
journal= {arXiv preprint arXiv:2408.14037},
year = {2025}
}