Open Source Robotics

高品質なオープンソースロボティクスプロジェクト、ツール、ソフトウェアの厳選コレクション。

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 はキュレーションを min-max 問題として定式化する。離散ポリシーが参照モデルに対する行動クローニングの超過損失を最小化する一方、ドメイン重みはそれを最大化するように指数勾配上昇法で更新され、ポリシーがまだ改善できるドメインに重みが割り当てられる。ロボット間で損失を比較可能にするため、行動はドメインごとにガウス正規化されてビンに離散化され、参照チェックポイントは、いずれかのドメインに過学習する前に早期打ち切りによって選ばれる。その後、学習全体で平均した重みを再利用して、別の拡散ポリシーを学習させたり、データをサブサンプリングしたりする。

主な知見

  • 生の損失ではなく参照ポリシーに対する超過損失を用いることで、すでに適合しているドメインと、適合が難しすぎるドメインの両方の重みが下がる。
  • 行動スケールの違いはロバスト最適化を歪める。境界による正規化を用いた簡易試験では、重みの 94.3% が純粋なランダム行動のドメインに割り当てられた。ドメインごとのガウス正規化では、実データが 84.2% を得た。
  • 連続的な L1/L2 の行動損失は外れ値や多峰的な行動に左右されるため、最終的なポリシーが連続的な拡散モデルであっても、行動を離散化したほうが信頼できる重みが得られる。
  • ロボットデータセット（約 10k-100k デモ）は参照モデルが過学習するほど小さい。早期打ち切りから 50K ステップ後のチェックポイントを使うと、性能は平均で 15% 以上低下した。

主な結果

  • RT-X のデータ配合において、ReMix の重みで学習したポリシーは、4 つの WidowX と 2 つの Franka の実世界タスクの平均で、一様な重みを 38%、RT-X の人間の専門家による重みを 32% 上回った。専門家の重みは一様な重みを 6% 上回った。
  • RT-X 配合の 25% のサブセットでは、ReMix は性能を維持した一方、人間の重みは 10% 以上低下した。ReMix は UR5 のデータの 72% と Kuka の 12% を残し、人間の重みではそれぞれ 30% と 24% であった。
  • ReMix は小規模な Toto データセットの比率を 3.42%（一様）から 16.3% に引き上げ、Kuka を 24.9% から 12.1% に引き下げた。
  • 完全な Bridge V2 データセットでは ReMix と一様な重みの間に目立った差はなかった。Bridge の 25% のサブセットでは、ReMix は一様な重みと SSP の画像埋め込みベースラインを上回った。
  • 妨害物体のある WidowX のキューブタスクでは、ReMix の重みは 7/10 で成功し、人間の重みと一様な重みはいずれも 0/10 であった。
  • シミュレーション（オペレータごとに分割した Robomimic NutAssemblySquare）では、50% のサブサンプリングで ReMix は 77/100、一様な重みは 53/100、25% では 59/100 対 39/100 であった。

制限事項

  • 著者らは、実世界評価が WidowX と Franka の 2 種類のアームのみを対象としていたと述べ、より多くの身体形態は今後の課題としている。
  • ReMix は Toto のような異常で多峰的な行動分布を持つデータセットの重みを上げるが、著者らはこれが必ずしも望ましいとは限らないと述べている。
  • 重みの計算には、参照モデル用とロバスト最適化用に、完全なデータで 2 回学習する必要がある。
  • 試験は Bridge V2 と OpenX の RT-X サブセットで行われ、2M を超えるエピソードを含む OpenX データセット全体では行われていない。Section 4.1 は Bridge V2 を 50k デモとしているが、Appendix B.2 はデータセット全体の 60K の軌道のうち約 45,000 本を使用したと述べている。
  • 事前学習データのみで学習したポリシーは著者らの環境にゼロショットで汎化しなかったため、すべての評価はタスクごとに 25 回のドメイン内デモを用いた共同学習に依存している。

分野への影響

CoRL 2024 Outstanding Paper Award のファイナリストであり、被引用数は 74 である。著者らは github.com/jhejna/remix でコードを公開している。論文はまた、RT-X 配合について学習したドメイン重み（Table 2）と、32 の Bridge V2 ドメインについての重み（Table 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

このウェブサイトにはJavaScriptが必要です。オープンソースロボティクスプロジェクトを閲覧するにはJavaScriptを有効にしてください。