README.md

---
license: apache-2.0
base_model: lerobot/pi0
pipeline_tag: robotics
tags:
  - mlx
  - robotics
  - vla
  - lerobot
  - pi0
---

# pi0-mlx

Physical Intelligenceのロボット基盤モデル [π0](https://huggingface.co/lerobot/pi0)（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`してから実行してください。

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

```python
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でも動きます。

```bash
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](https://github.com/huggingface/lerobot)で行い、Mac上での推論・検証・デモに本移植を使う構成を想定しています。同一アーキテクチャのFT済み重みは`model.safetensors`を差し替えれば動きます。前処理（224pxアスペクト維持リサイズ・言語トークナイズ・状態パディング）はランタイムに内蔵しています。

同シリーズ: [smolvla-mlx](https://huggingface.co/tokimoa/smolvla-mlx)（450M・軽量版のVLA移植）

## ライセンス

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

---
Developed by [tokimoa](https://tokimoa.jp)