π₀-FAST (Pi0-FAST)

π₀-FASTは、汎用ロボット制御のためのVision-Language-Actionモデルであり、自己回帰的な次のトークン予測を用いて連続的なロボットアクションをモデル化します。

モデル概要

π₀-FASTは、Vision-Language Modelsの力を、FAST（Frequency-space Action Sequence Tokenization）と呼ばれる新しいアクションのトークン化手法と組み合わせます。これにより、標準的なバインニングベースの離散化では不可能な、高度な器用さを要求するタスクでの自己回帰型VLAの学習が可能になり、π₀などの拡散ベースのアプローチと比較して、学習速度が最大で5倍速になります。

なぜ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のインストールは、インストールガイドに従って行ってください。

2. …を実行して、π₀-FASTの依存関係をインストールします。

   コピーしました

   pip install -e ".[pi]"

カスタムFASTトークナイザーの学習

FASTトークナイザーには2つのオプションがあります:

1. 事前学習済みトークナイザーを使用する: lerobot/fast-action-tokenizer トークナイザーは、100万本以上の実際のロボットのアクションシーケンスで学習されており、汎用トークナイザーとして機能します。

2. 独自のトークナイザーを学習する: 特定のデータセットでの最大限のパフォーマンスを実現するために、ご自身のデータでトークナイザーをファインチューニングできます。

独自のTokenizerの学習

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"

主なトークナイザーのパラメータ

使用方法

LeRobotでπ₀-FASTを使用するには、ポリシータイプを以下のように指定します:

policy.type=pi0_fast

トレーニング

π₀-FASTの学習には、LeRobotの学習スクリプトを使用できます：

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

主要な学習パラメータ

推論

高速推論のためのKVキャッシュ

π₀-FASTは、LLM推論で広く使用されている最適化手法であるKV-cachingに対応しています。これにより、アテンション機構からのキーとバリューのペアをキャッシュし、自己回帰デコーディング時の冗長な計算を回避します。

# KV-caching is enabled by default
policy.use_kv_cache=true

推論の例

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トークンはデコードされて連続的なアクションに変換されます。

設定オプション

π₀との比較

π₀Fastの結果の再現

LeRobotの実装を用いて、LIBEROベンチマークにおけるπ₀Fastの結果を再現します。LeRobotのPiFastベースモデルlerobot/pi0fast-baseを取得し、HuggingFace LIBERO datasetを使用して、8枚のH100 GPU上でバッチサイズ256、bfloat16でさらに40kステップのファインチューニングを行います。

ファインチューニングされたモデルは、こちらで確認できます:

• π₀Fast LIBERO: lerobot/pi0fast-libero

以下のトレーニングコマンドで:

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 実装を使用して、次のコマンドを実行することで、ファインチューニングされたモデルを評価します。

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ベンチマークで以下の結果が得られました:

動画を含む完全な評価出力フォルダは、こちらでご利用いただけます。

ライセンス

このモデルは、元のOpenPIリポジトリと一貫して、Apache 2.0ライセンスに従っています。

参考文献

• FAST: 効率的なロボットのアクショントークン化 - Physical Intelligence ブログ
• OpenPI リポジトリ - 元の実装
• Hugging Face 上の FAST Tokenizer - 事前学習済みトークナイザー