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