[論文レビュー] Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

Lirui Wang, Xinlei Chen+2 more|arXiv (Cornell University)|Sep 30, 2024

Image Retrieval and Classification Techniques被引用数 5

ひとこと要約

Heterogeneous Pre-trained Transformers (HPT) は embodiment-specific stems と task-specific heads を用いて多様なロボット実装間で共有ポリシートランクを事前学習し、新しい実装とタスクへの転移を可能にする。52データセットと1Bパラメータを超える規模で性能とスケーラビリティを改善。

ABSTRACT

One of the roadblocks for training generalist robotic models today is heterogeneity. Previous robot learning methods often collect data to train with one specific embodiment for one task, which is expensive and prone to overfitting. This work studies the problem of learning policy representations through heterogeneous pre-training on robot data across different embodiments and tasks at scale. We propose Heterogeneous Pre-trained Transformers (HPT), which pre-train a large, shareable trunk of a policy neural network to learn a task and embodiment agnostic shared representation. This general architecture aligns the specific proprioception and vision inputs from distinct embodiments to a short sequence of tokens and then processes such tokens to map to control robots for different tasks. Leveraging the recent large-scale multi-embodiment real-world robotic datasets as well as simulation, deployed robots, and human video datasets, we investigate pre-training policies across heterogeneity. We conduct experiments to investigate the scaling behaviors of training objectives, to the extent of 52 datasets. HPTs outperform several baselines and enhance the fine-tuned policy performance by over 20% on unseen tasks in multiple simulator benchmarks and real-world settings. See the project website (https://liruiw.github.io/hpt/) for code and videos.

研究の動機と目的

• 多様な実装とタスクに跨る、スケーラブルで一般化可能なロボットポリシー学習を動機づける。
• 異なるロボットからの固有受容覚と視覚を共有表現へ整列させるモジュラーアーキテクチャ（stem、trunk、heads）を提案する。
• 実データ、シミュレーション、ヒューマンビデオデータセットを跨るデータ量・モデルサイズ・計算量によるスケーリング挙動をデモンストレーションする。
• 監視付き事前学習と微調整を通じて、未見の実装・タスク・現実世界シナリオへの転移性能を示す。

提案手法

• 異種入力をモダリティごとに固定セットのトークン（例: 16）へ写像するステム（proprioception tokenizer と vision tokenizer）を導入する。
• 結合されたトークンを処理し、共通の潜在表現へと変換する共有トランスフォーマー・トランクを採用する。
• タスク固有のヘッドを用いて、トランク出力を各 embodiment-task ペアの行動へ写像する。
• K 個の異種データセットを跨ぐ正規化された行動に対してHuber損失を用いた挙動模倣（behavior cloning）目的で訓練し、データセットごとにステム/ヘッドを更新し、トランクを jointly に更新する。
• 最大52データセットを跨る事前学習で1Bパラメータ超を達成し、新しい embodiment への転移を可能にするためにステム/ヘッドを再初期化し、トランクを凍結する。

実験結果

リサーチクエスチョン

• RQ1 現実ロボット、シミュレーション、ヒューマン動画を跨るデータ量と多様性で、ヘテロジニアス前訓練はどのようにスケールするか？
• RQ2 多様な実装から学習された単一のトランクは、未知の実装とタスクへの効果的な転移を最小限の適応で実現できるか？
• RQ3 モデルサイズとバッチ規模が事前学習の収束と下流転移性能に及ぼす影響は？
• RQ4 事前学習された HPT 表現は、シミュレーションベンチマークと現実世界のロボットタスクの両方へどれだけ転移するか？
• RQ5 異種データを取り入れることは、実装と環境を跨る頑健性と一般化を向上させるか？

主な発見

• HPT はより大きなモデル、より多くのデータ、より高い計算量からのスケーリングの利益を享受し、データ/モードが増えるにつれて検証損失が改善される。
• 事前学習によりより多くの実装を含めると、トランクの一般化とタスク間転移性能が向上する。
• 最大1Bパラメータ（HPT-Huge）までの事前学習を大きなバッチサイズで行うと、プラトーに至るまで継続的な改善が見られ、深さ対幅のスケーリングによる利益は最小限。
• 合成シミュレーションデータとインターネットのヒューマンビデオを使った事前学習は実現可能で、補完的な embodiment データを提供し、転移効果を維持する。
• シミュレーションベンチマークへの転移では、HPT はゼロからの訓練やトランクなしよりタスク成功率を改善し、ファインチューニングされた HPT 変種はベースラインより高い性能（例: HPT-XL）に達する。
• 現実世界のテストでは、事前学習済みポリシーは、視点設定と物体多様性の異なるベースラインと比較して頑健性と一般化が向上する。