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

要約

ロボット工学の基礎モデルのトレーニングを目的として、ますます大規模な模倣学習データセットが収集されています。
しかし、視覚および自然言語処理においてデータ選択が最も重要であるという事実にもかかわらず、ロボット工学における研究では、そのようなモデルが実際にどのようなデータに基づいてトレーニングされるべきなのかについて疑問を呈する研究はほとんどありませんでした。
この研究では、ロボット基礎モデルの事前トレーニングのためにロボット データセットのさまざまなサブセットまたは「ドメイン」を重み付けする方法を調査します。
具体的には、分散ロバスト最適化 (DRO) を使用して、考えられるすべての下流ドメインにわたって最悪の場合のパフォーマンスを最大化します。
私たちの手法である Re-Mix は、ロボット データセットに DRO を適用する際に生じる、アクション スペースの変動性やさまざまなデータセット間のダイナミクスなどの幅広い課題に対処します。
Re-Mix は、これらの問題に対処するために、早期停止、アクションの正規化、および離散化を採用しています。
最大のオープンソース ロボット操作データセットである Open X-Embodiment データセットでの広範な実験を通じて、データ キュレーションが下流のパフォーマンスに多大な影響を与える可能性があることを実証しました。
具体的には、Re-Mix によって学習されたドメインの重みは、既存のジェネラリスト ロボット ポリシー、特に RT-X モデルのトレーニングに使用されるデータセット上で均一の重みを平均 38\% 上回り、人間が選択した重みを 32\% 上回っています。

要約(オリジナル)

Increasingly large imitation learning datasets are being collected with the goal of training foundation models for robotics. However, despite the fact that data selection has been of utmost importance in vision and natural language processing, little work in robotics has questioned what data such models should actually be trained on. In this work we investigate how to weigh different subsets or “domains” of robotics datasets for robot foundation model pre-training. Concrete, we use distributionally robust optimization (DRO) to maximize worst-case performance across all possible downstream domains. Our method, Re-Mix, addresses the wide range of challenges that arise when applying DRO to robotics datasets including variability in action spaces and dynamics across different datasets. Re-Mix employs early stopping, action normalization, and discretization to counteract these issues. Through extensive experimentation on the largest open-source robot manipulation dataset, the Open X-Embodiment dataset, we demonstrate that data curation can have an outsized impact on downstream performance. Specifically, domain weights learned by Re-Mix outperform uniform weights by 38\% on average and outperform human-selected weights by 32\% on datasets used to train existing generalist robot policies, specifically the RT-X models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google