テクニカルサマリー：DataMIL – Datamodelsを用いたロボット模倣学習のためのデータ選択

1. 問題提起

ロボティクス・コミュニティは近年、大規模かつ多様なデータセット（例：Open-X Embodiment）を用いて汎用的な基盤ポリシーを訓練する方向へとシフトしています。これらのポリシーは、多様なタスクに対して高い平均性能を達成しますが、個別の専門的なタスクにおいては性能が不足することが頻繁にあり、新たに取得したタスク固有のデータによるファインチューニングが必要となります。

本研究が取り組む核心的な課題は、データの選択です。すなわち、大規模な既存データセットのどのサブセットを、限られたタスク固有のデモンストレーションと組み合わせることで、最適な専門化ポリシーが得られるかを特定することです。

• 既存手法の限界: セマンティックな類似性、視覚的な類似性、あるいは状態-行動の近接性といった、人間的な感覚に基づくナイーブな選択は、しばしば失敗します。これらのヒューリスティックは、「最も『似ている』データが最も『有用』である」という仮定に基づきますが、これは実際のダウンストリーム・ポリシーの性能への影響を無視しています。
• 最適化の障壁: 理想的には、サブセットを繰り返し訓練して結果となるポリシーを評価することで、網羅的にテストしてデータを選択すべきです。しかし、ロボティクスにおけるポリシーの評価には実世界でのロールアウトが必要であり、これは時間がかかり、危険であり、大規模な探索を行うには計算量的に不可能です。
• ギャップ: NLPやコンピュータビジョンでは成功しているdatamodels（モデルの性能を訓練データから予測するもの）のような既存のデータ属性フレームワークは、ロボティクスにおいては、実世界のポリシー評価が非微分的かつ高コストであるため、直接適用することが困難でした。

2. 手法：DataMIL

著者らは、datamodelのパラダイムをロボティクスへと拡張したフレームワークであるDataMIL（Datamodels for Imitation Learning）を提案します。DataMILは、静的なヒューリスティックに頼るのではなく、ポリシー自体を用いて性能を向上させるデータポイントを特定することで、エンドツーエンドでデータの選択を推論します。

2.1 コアフレームワーク

DataMILは、ポリシー学習アルゴリズムを「ブラックボックス」として扱い、特定のデータサブセットで訓練されたポリシーの性能を予測する推定器（datamodel）を訓練します。

• 目的: 学習アルゴリズム $A$ と性能指標 $M$ に対して、ターゲット指標 $M(A(D'))$ を最大化するサブセット $D' \subset D$ を見つけること。
• Datamodelの定式化: ポリシーを訓練する代わりに、DataMILは関数 $\hat{f}(D') \approx M(A(D'))$ を学習します。線形datamodelの場合、これは各データ点 $z_i$ に対してスカラー値のスコア $\tau(z_i)$ を割り当てることに帰着し、予測される性能は選択されたサブセットのスコアの総和となります。

2.2 ロボティクスへの重要な適応

Datamodelsをロボティクスにおいて扱いやすくするために、著者らは3つの重要な修正を導入しています。

1. サロゲート損失関数（ロールアウトの回避）:
   実世界のロールアウトは非微分的で高コストです。DataMILは、真の成功率メトリックを、保持されたターゲット・デモンストレーションに対するプロキシ・メトリック $\tilde{M}$ に置き換えます。これは以下の通り定義されます。
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   このメトリックは完全微分可能であり、環境との相互作用を必要としないため、勾配ベースの推定手法の使用を可能にします。

2. Datamodel推定器:
   論文では、影響度スコア $\tau(z_i)$ を推定するために2つの手法を採用しています。

   • 回帰推定器: ランダムなデータサブセットをサンプリングし、ポリシーを訓練し、プロキシ・メトリックを評価し、データの存在に基づいて性能を予測する線形モデルを適合させます。これは正確ですが、計算負荷が高い手法です。
   • Metagradientベースの推定器: 影響関数（influence functions）とメタグラディエントを用いて、データ重みに対するターゲット・メトリックの微分を計算します。この手法は、大規模なモデル（Octoなど）に対してより効率的であり、すべてのサブセットに対してポリシーを再訓練する必要を回避します。

3. クラスタリングと分布シフトの緩和:

   • クラスタリング: 影響度推定におけるノイズを減らすため、個々の状態-行動ペアを時間的なクラスター（部分軌跡または完全な軌道）にグループ化します。最適なクラスターの粒度はデータセットのサイズに依存します。大規模なデータセット（OXEなど）では、軌道レベルの集約が有効です。
   • 分布シフト: 既存データとターゲットドメイン間のシフトによる不正確な推定を防ぐため、学習プロセスをターゲットドメインに適合させる目的で、少量のターゲットタスクデータをdatamodel推定フェーズに含めます。

2.3 訓練パイプライン

影響度スコアが推定された後：

1. 最も高い正の影響を持つ上位 $x\%$ の既存データを選択し、$D_{sel}$ を形成します。
2. 最終的なポリシーは、ターゲットデータ $D_{target}$ と選択されたデータ $D_{sel}$ を用いた**共同訓練（co-training）**を通じて訓練されます。各ステップにおいて、確率 $\alpha$ で $D_{target}$ から、確率 $1-\alpha$ で $D_{sel}$ からバッチがサンプリングされます。

3. 主な貢献

• Datamodelsのロボティクスへの拡張: DataMILは、微分可能なサロゲート損失を用いることで、高コストな実世界ロールアウトという障壁を克服し、datamodelsをロボットの模倣学習に初めて成功裏に適用したフレームワークです。
• エンドツーエンドの性能認識: 視覚的、運動的、あるいは状態の類似性といったヒューリスティックな手法とは異なり、DataMILは性能への影響を直接最適化します。これにより、見た目が視覚的に似ていても、性能を向上させるデータを選択し、性能を低下させるデータを排除することが可能になります。
• スケーラブルな推定器: メタグラディエントベースの推定器の導入により、回帰ベースのアプローチでは計算量的に不可能であった大規模で複雑なポリシー（Octoなど）や大規模なデータセット（OXEなど）へのスケールが可能になりました。
• クロスエンボディメントおよびマルチタスク選択: 既存のデータセットに存在しないタスクやロボットのエンボディメントに対して有用なデータを選択する能力、および複数のダウンストリーム目標を同時にサポートするデータセットを精選する能力を示しています。

4. 実験結果

著者らは、シミュレーションおよび実世界の環境における60以上のタスクを用いてDataMILを検証しました。

• MetaWorld (50タスク): DataMILは、最先端のベースライン（Behavior Retrieval, Flow Retrieval, STRAP）に対して10%の性能向上を達成しました。自律探索データに含まれるノイズの多い、あるいは不適切なデモンストレーションを効果的にフィルタリングできました。一方で、類似性ベースのベースラインは、無関係または低品質なデータを取得してしまう傾向がありました。
• LIBERO (10タスク): Octoトランスフォーマーポリシーを使用した場合、DataMILは多様な長期間タスクにおいて一貫してベースラインを上回りました。視覚的類似性ベースの手法が特定のタスクで良好な性能を示すこともありましたが、DataMILは最も堅牢な平均性能を提供しました。
• Open-X Embodiment (OXE) (実世界):
  • Franka-Ball & Franka-Pouch: DataMILは、ランダム選択やヒューリスティックなベースラインと比較して、成功率を大幅に向上させるデータを選択しました。
  • Tiago-Sink: Tiagoのデータが含まれていないデータセットからTiagoロボット用のデータを選択するという困難なクロスエンボディメント設定において、DataMILは、他のロボット（例：RT-1, BC-Z）から、タスクの「本質」（テーブルトップ操作）を共有する関連性の高いデモンストレーションを特定することに成功しましたが、ベースラインは失敗しました。
  • Droid-Multitask: DataMILは、3つの異なるタスクに対して同時に性能を向上させる単一のデータセットを精選し、平均成功率においてベースラインを上回りました。

定性的知見:

• DataMILは、単一のソースに過度に依存する（ベースラインはしばしば1つのデータセットから80%以上を回収する）のではなく、多様なデータセットの混合を選択します。
• 見た目が視覚的に同一であっても（例：同じ状態だが行動分布が異なる）、役に立つデータと有害なデータを効果的に区別できます。これは、類似性ベースの検索においてよく見られる失敗モードです。

5. 重要性と主張

本論文は、エンドツーエンドの性能を考慮したデータ選択が、ロボティクスにおける大規模な既存データセットの潜在能力を引き出すために極めて重要であると主張しています。著者らは、基盤モデルが強力なベースを提供する一方で、その専門化はファインチューニングデータの質に大きく依存すると述べています。

• スケーラビリティに関する控えめな主張: 著者らは、datamodelの推定には依然として計算コストがかかること（全データに対するポリシー訓練の数倍のコスト）を認めていますが、メタグラディエントによるアプローチがこれを軽減しています。また、ハイパーパラメータのチューニング（例：クラスターサイズ、選択比率）には、現在、強力な理論的直観が欠けていることも指摘しています。
• 範囲: 本研究は主にシングルタスクの専門化に焦点を当てており、Droid-Multitaskの設定はマルチタスク精選への予備的なステップとして位置付けられています。
• 結論: DataMILは、人間が定義したヒューリスティックを超えて、データ駆動型かつモデル認識型の選択を行うことが、従来の検索手法が破綻する複雑でヘテロジニアスな実世界のシナリオにおいて、一貫した成功率の向上をもたらすことを示しています。