pi0-mlx

Physical Intelligenceのロボット基盤モデル π0（PaliGemma 3B + アクションエキスパート・Vision-Language-Action）のApple Silicon（MLX）移植です。カメラ画像・言語指示・関節状態から50手先までのアクションチャンクをflow matchingで生成します。

実行系	1チャンク（50手）生成	ピークメモリ
本移植（MLX・bf16）	522ms	8.0GB
PyTorch MPS（参照実装）	660ms
PyTorch CPU（参照実装）	2,395ms

同一入力・同一ノイズでPyTorch参照実装（lerobot main / openpi直系）とコサイン類似度0.99993（bf16、fp32重みでは1.00000）の出力一致を検証済みです。プレフィル→エキスパートがキャッシュへ毎層アテンションする二相構造、Gemma固有の正規化（(1+w)RMSNorm）・GeGLU・言語埋め込みの√widthスケーリングまで参照実装を忠実に再現しています。

使い方

トークナイザは上流（lerobot）と同じくgoogle/paligemma-3b-pt-224を参照します。Hugging FaceでGemma利用規約に同意し、hf auth loginしてから実行してください。

pip install mlx-vlm pillow transformers
hf download tokimoa/pi0-mlx --local-dir pi0-mlx

from pi0_mlx import Pi0MLX

model = Pi0MLX.from_pretrained("pi0-mlx")
actions = model.predict(
    images=[cam0, cam1, cam2],            # HWC uint8（1〜3カメラ）
    instruction="pick up the cube",
    state=[0.1, -0.2, 0.3, 0.0, 0.5, 0.0],
)                                          # -> (50, len(state)) アクションチャンク

CLIでも動きます。

python pi0-mlx/pi0_mlx.py --images cam0.png cam1.png cam2.png \
  --instruction "pick up the cube" --state 0,0,0,0,0,0

位置づけ

π0はベースモデルであり、実タスクへの適用には手元のロボットでのファインチューニングが前提です。学習はLeRobotで行い、Mac上での推論・検証・デモに本移植を使う構成を想定しています。同一アーキテクチャのFT済み重みはmodel.safetensorsを差し替えれば動きます。前処理（224pxアスペクト維持リサイズ・言語トークナイズ・状態パディング）はランタイムに内蔵しています。

同シリーズ: smolvla-mlx（450M・軽量版のVLA移植）

ライセンス

Apache-2.0（ベースモデルlerobot/pi0のライセンスを継承）

---

Developed by tokimoa