# π₀-FAST (Pi0-FAST) π₀-FASTは、**汎用ロボット制御のためのVision-Language-Actionモデル**であり、自己回帰的な次のトークン予測を用いて連続的なロボットアクションをモデル化します。 ## モデル概要 π₀-FASTは、Vision-Language Modelsの力を、**FAST（Frequency-space Action Sequence Tokenization）**と呼ばれる新しいアクションのトークン化手法と組み合わせます。これにより、標準的なバインニングベースの離散化では不可能な、高度な器用さを要求するタスクでの自己回帰型VLAの学習が可能になり、π₀などの拡散ベースのアプローチと比較して、学習速度が最大で**5倍速**になります。 An overview of Pi0-FAST ### なぜFASTなのか? ロボットのアクションのトークン化における標準的なアプローチは、次元ごと・タイムステップごとの単純なバインニング（区切り）スキームを使用します。単純な動作であれば問題ありませんが、精密さと高周波数制御を必要とする複雑で器用なスキルでは、この手法は急速に機能しなくなります。 FASTは、信号処理技術を用いてアクションシーケンスを圧縮することでこの問題を解決します。その結果、言語トークンと同様に自己回帰的に予測可能な、密なアクショントークンのシーケンスが得られます。 ### FASTトークナイゼーションの仕組み FASTトークナイザーは、以下の手順を通じてアクションシーケンスを圧縮します: 1. **Normalize**: ホライズンが`H`、アクション次元が`D`の、形状`(H, D)`の連続アクションチャンクを取得します。サポートされている正規化手法のいずれか（外れ値への対応にはQuantilesが推奨されます）を使用して正規化します。 2. **離散コサイン変換 (DCT)**: scipy を介して、各アクション次元に対して個別に DCT を適用します。DCT は、画像や音声コーデック (JPEG、MP3) でよく使用される圧縮アルゴリズムです。 3. **Quantization**: 各アクション次元について、無視できる係数を丸めて除去し、スパースな周波数行列を生成します。 4. **Flatten**: 低周波成分を先頭に配置して、行列を1次元ベクトルにフラット化します。 5. **Byte Pair Encoding (BPE)**: BPEトークナイザーを学習し、DCT係数を高密度なアクショントークンに圧縮します。従来のトークナイズ手法と比較して、通常**10倍の圧縮率**を実現します。 このアプローチにより、これらのFASTトークンを予測するように学習させることで、**既存のVLM**をVLAに変換できます。 ## インストール要件 1. LeRobotのインストールは、[インストールガイド](./installation)に従って行ってください。 2. …を実行して、π₀-FASTの依存関係をインストールします。 ```bash pip install -e ".[pi]" ``` ## カスタムFASTトークナイザーの学習 FASTトークナイザーには2つのオプションがあります: 1. **事前学習済みトークナイザーを使用する**: `lerobot/fast-action-tokenizer` トークナイザーは、100万本以上の実際のロボットのアクションシーケンスで学習されており、汎用トークナイザーとして機能します。 2. **独自のトークナイザーを学習する**: 特定のデータセットでの最大限のパフォーマンスを実現するために、ご自身のデータでトークナイザーをファインチューニングできます。 ### 独自のTokenizerの学習 ```bash lerobot-train-tokenizer \ --repo_id "user/my-lerobot-dataset" \ --action_horizon 10 \ --encoded_dims "0:6" \ --vocab_size 1024 \ --scale 10.0 \ --normalization_mode QUANTILES \ --output_dir "./my_fast_tokenizer" \ --push_to_hub true \ --hub_repo_id "username/my-action-tokenizer" ``` ### 主なトークナイザーのパラメータ | パラメータ | 説明 | デフォルト | | ---------------------- | --------------------------------------------------------------------------------- | ------------ | | `--repo_id` | LeRobot データセットのリポジトリ ID | 必須 | | `--action_horizon` | 各チャンク内の将来アクションの数 | `10` | | `--encoded_dims` | エンコードするカンマ区切りの次元範囲（例: `"0:6,7:23"`） | `"0:6,7:23"` | | `--vocab_size` | BPE語彙サイズ | `1024` | | `--scale` | 量子化のためのDCTスケーリング係数 | `10.0` | | `--normalization_mode` | 正規化モード(`MEAN_STD`、`MIN_MAX`、`QUANTILES`、`QUANTILE10`、`IDENTITY`) | `QUANTILES` | | `--sample_fraction` | エピソードごとにサンプリングするチャンクの割合 | `0.1` | ## 使用方法 LeRobotでπ₀-FASTを使用するには、ポリシータイプを以下のように指定します: ```python policy.type=pi0_fast ``` ## トレーニング π₀-FASTの学習には、LeRobotの学習スクリプトを使用できます： ```bash lerobot-train \ --dataset.repo_id=your_dataset \ --policy.type=pi0_fast \ --output_dir=./outputs/pi0fast_training \ --job_name=pi0fast_training \ --policy.pretrained_path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --policy.gradient_checkpointing=true \ --policy.chunk_size=10 \ --policy.n_action_steps=10 \ --policy.max_action_tokens=256 \ --steps=100000 \ --batch_size=4 \ --policy.device=cuda ``` ### 主要な学習パラメータ | パラメータ | 説明 | デフォルト | | -------------------------------------- | -------------------------------------------------- | ------------------------------- | | `--policy.gradient_checkpointing=true` | 学習中のメモリ使用量を大幅に削減します | `false` | | `--policy.dtype=bfloat16` | 効率化のために混合精度トレーニングを使用します | `float32` | | `--policy.chunk_size` | 予測するアクションステップ数（アクションホライズン） | `50` | | `--policy.n_action_steps` | 実行するアクションステップ数 | `50` | | `--policy.max_action_tokens` | アクションチャンクあたりの最大FASTトークン数 | `256` | | `--policy.action_tokenizer_name` | 使用するFASTトークナイザ | `lerobot/fast-action-tokenizer` | | `--policy.compile_model=true` | 高速なトレーニングのためにtorch.compileを有効にする | `false` | | `--policy.optimizer_fused=true` | fused AdamWを使用します。再開時には無視されます | `None` | ## 推論 ### 高速推論のためのKVキャッシュ π₀-FASTは、LLM推論で広く使用されている最適化手法である**KV-caching**に対応しています。これにより、アテンション機構からのキーとバリューのペアをキャッシュし、自己回帰デコーディング時の冗長な計算を回避します。 ```python # KV-caching is enabled by default policy.use_kv_cache=true ``` ### 推論の例 ```python from lerobot.policies.pi0_fast import PI0FastPolicy, PI0FastConfig # Load the policy policy = PI0FastPolicy.from_pretrained("your-model-path") # During inference actions = policy.predict_action_chunk(batch) ``` ## モデルアーキテクチャ π₀-FASTは、PaliGemmaベースのアーキテクチャを使用しています： - **Vision Encoder**: 画像理解のためのSigLIPビジョンタワー - **Language Model**: 言語指示の処理とアクショントークンの予測に使用される Gemma 2B モデルは画像、テキストによる指示、ロボットの状態を入力として受け取り、離散的なFASTトークンを出力します。このFASTトークンはデコードされて連続的なアクションに変換されます。 ## 設定オプション | パラメータ | 説明 | デフォルト | | -------------------- | ----------------------------------------------- | ---------- | | `paligemma_variant` | VLMバックボーンバリアント(`gemma_300m`、`gemma_2b`) | `gemma_2b` | | `max_state_dim` | 最大状態ベクトル次元（パディング済み） | `32` | | `max_action_dim` | 最大アクションベクトル次元数（パディング済み） | `32` | | `temperature` | サンプリング温度（貪欲法の場合は0.0） | `0.0` | | `max_decoding_steps` | 最大デコーディングステップ数 | `256` | | `use_kv_cache` | 高速な推論のためにKVキャッシュを有効にする | `true` | ## π₀との比較 | 機能 | π₀ | π₀-FAST | | --------------------- | ------------------------- | ---------------------------- | | アクションの表現 | Flow Matching（Diffusion） | 自己回帰トークン (FAST) | | トレーニング速度 | 1x | **5倍高速** | | 器用さ | 高 | 高 | | 推論方法 | 反復的なノイズ除去 | 自己回帰デコーディング | | KVキャッシュ | N/A | サポート | ## π₀Fastの結果の再現 LeRobotの実装を用いて、LIBEROベンチマークにおけるπ₀Fastの結果を再現します。LeRobotのPiFastベースモデル[lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base)を取得し、[HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero)を使用して、8枚のH100 GPU上でバッチサイズ256、bfloat16でさらに40kステップのファインチューニングを行います。 ファインチューニングされたモデルは、こちらで確認できます: - **π₀Fast LIBERO**: [lerobot/pi0fast-libero](https://huggingface.co/lerobot/pi0fast-libero) 以下のトレーニングコマンドで: ```bash lerobot-train \ --dataset.repo_id=lerobot/libero \ --output_dir=outputs/libero_pi0fast \ --job_name=libero_pi0fast \ --policy.path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --steps=100000 \ --save_freq=20000 \ --batch_size=4 \ --policy.device=cuda \ --policy.scheduler_warmup_steps=4000 \ --policy.scheduler_decay_steps=100000 \ --policy.scheduler_decay_lr=1e-5 \ --policy.gradient_checkpointing=true \ --policy.chunk_size=10 \ --policy.n_action_steps=10 \ --policy.max_action_tokens=256 \ --policy.empty_cameras=1 ``` その後、LeRobot の LIBERO 実装を使用して、次のコマンドを実行することで、ファインチューニングされたモデルを評価します。 ```bash tasks="libero_object,libero_spatial,libero_goal,libero_10" lerobot-eval \ --policy.path=lerobot/pi0fast-libero \ --policy.max_action_tokens=256 \ --env.type=libero \ --policy.gradient_checkpointing=false \ --env.task=${tasks} \ --eval.batch_size=1 \ --eval.n_episodes=1 \ --rename_map='{"observation.images.image":"observation.images.base_0_rgb","observation.images.image2":"observation.images.left_wrist_0_rgb"}' ``` **注:** 元のOpenPI実装と同様に、`n_action_steps=10`を設定しています。 ### 結果 LIBEROベンチマークで以下の結果が得られました: | モデル | LIBERO Spatial | LIBERO Object | LIBERO Goal | LIBERO 10 | 平均 | | ----------- | -------------- | ------------- | ----------- | --------- | -------- | | **π₀-fast** | 70.0 | 100.0 | 100.0 | 60.0 | **82.5** | 動画を含む完全な評価出力フォルダは、[こちら](https://drive.google.com/drive/folders/1HXpwPTRm4hx6g1sF2P7OOqGG0TwPU7LQ?usp=sharing)でご利用いただけます。 ## ライセンス このモデルは、元の[OpenPIリポジトリ](https://github.com/Physical-Intelligence/openpi)と一貫して、**Apache 2.0ライセンス**に従っています。 ## 参考文献 - [FAST: 効率的なロボットのアクショントークン化](https://www.physicalintelligence.company/research/fast) - Physical Intelligence ブログ - [OpenPI リポジトリ](https://github.com/Physical-Intelligence/openpi) - 元の実装 - [Hugging Face 上の FAST Tokenizer](https://huggingface.co/physical-intelligence/fast) - 事前学習済みトークナイザー