pi0 and pi0-FAST: 汎用ロボット制御のためのVision-Language-Actionモデル

Hugging Faceは、Physical Intelligenceによって開発された汎用的なVision-Language-Action (VLA) モデルである π0 and π0-FAST を LeRobot リポジトリに統合しました。これらのモデルは、大規模な事前学習と革新的なアクション表現技術を活用することで、異なるハードウェア構成においてロボットが複雑で器用な操作タスクを実行することを可能にします。

π0: Flow Matchingによる汎用ロボット制御

π0 (Pi-Zero) は、7つのロボットプラットフォームと68のユニークなタスクのデータでトレーニングされた、汎用ロボット制御用に設計されたVLAモデルです。洗濯物の折り畳み、食料品の袋詰め、箱の組み立て、物体の回収といった実世界のタスクを、ゼロショットおよびファインチューニング設定の両方で実行する能力を備えています。

標準的なロボットポリシーとは異なり、π0は flow matching を使用して、50Hzで滑らかなリアルタイムのアクション軌道を生成します。このプロセスはランダムなノイズから始まり、一連のモーターアクションへと漸進的に収束し、デプロイメント中のモデルの効率性と精度を保証します。

π0-FAST: 自己回帰的な効率性とFAST Tokenization

π0-FASTは、トレーニング速度とアクションの忠実度を向上させるために Frequency-space Action Sequence Tokenization (FAST) を利用した、π0の自己回帰バージョンです。

π0-FASTの主な利点

π0-FASTは、拡散ベースのVLAと比較していくつかの改善点を提供します：

トレーニング速度: 5倍速いトレーニングを実現します。
アクション表現: アクションシーケンス内の冗長性を削減し、表現力を向上させます。
汎用性: ロボットの形態や未知の環境に対して、より強力な汎用性を示します。

FAST Tokenizationの仕組み

FASTは、Discrete Cosine Transform (DCT) を使用して、連続的なアクションシーケンスを離散的なトークンに圧縮します。パイプラインは以下のステップに従います：

Normalization: 異なるシステム間での一貫性を確保するため、生のロボットアクションは範囲 [-1, 1] に量子化正規化されます。
DCT Transformation: 各アクション次元は、時間領域から周波数領域に変換されます。
Compression: scale-and-round操作により、重要でない係数が除去されます。
Flattening: 結果として得られるスパースなDCT係数行列は、低周波成分を優先して1次元の整数シーケンスにフラット化されます。
BPE Encoding: Byte Pair Encoding (BPE) が頻出パターンを結合し、固定サイズの語彙を維持します。

これらの操作は可逆であるため、アクションをロスレスで再構成できます。ユニバーサルバージョンである FAST+ は、モバイル、両腕、単腕ロボットからの100万のアクションシーケンスでトレーニングされており、Hugging FaceのAutoProcessorとして利用可能です。

技術アーキテクチャ: VLA Attention Mechanisms

VLAは、アクションと観測状態のトークンを組み込むことで、Vision-Language Models (VLMs) を拡張したものです。π0は、これらのマルチモーダルな入力を処理するために特定のAttention構造を採用しています：

Prefix Tokens: 全シーン（画像とテキスト）を表し、PaliGemmaと同様に互いに完全にAttentionを向けます。
State Tokens: ロボットの現在の環境状態（例：関節角度）を表します。これらは、三角形（因果的）な方法でprefix tokensと前のstate tokensにAttentionを向けます。
Action Tokens: モーターコマンドシーケンスを表します。これらのトークンは、パディングを除き、画像トークン、テキストトークン、state tokens、および他のaction tokensに対して完全な可視性を持ちます。

FlexAttentionによるAttentionの最適化

強いブロックスパース性を示す結果としての2D因果マスクを処理するために、実装には PyTorch FlexAttention が使用されています。単純な行列乗算や、不規則なブロックマスクの扱いに苦労するFlashAttention2を使用するのではなく、チームは因果マスクのインデックス作成によって作成されたブロックマスクを使用しています。このアプローチにより、不要な計算を効率的にスキップし、パフォーマンスを向上させています。

LeRobotへの実装

π0とπ0-FASTは現在、LeRobotリポジトリで利用可能です。ユーザーは、学習済みモデルで推論を実行したり、特定の環境に合わせてファインチューニングしたりすることができます。ファインチューニングは lerobot/scripts/train.py スクリプトを介して実行でき、ベースとなるπ0モデルを特定のロボットタスクや環境に適応させることが可能です。