pi0-mlx: MLX port of pi0 (parity cos 0.99993 bf16, 522ms/chunk, 8GB peak) + runtime + card

README.md

@@ -0,0 +1,63 @@
+---
+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)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af62ab4f208e0263a1cd4b5e5e166cd08e68c56f3d936d4153085428496c6cee
+size 7827762546

pi0_mlx.py

@@ -0,0 +1,268 @@
+"""pi0 (Physical Intelligence / lerobot) inference runtime for Apple Silicon / MLX.
+
+Faithful port of the lerobot (openpi-derived) PI0Pytorch reference:
+- prefill: PaliGemma LM (18 layers, bidirectional prefix) with RoPE-applied KV cache
+- 10 Euler flow-matching steps; the 300M Gemma expert attends to the cached
+  prefix KV at every layer (head_dim 256 / 1 KV head on both sides)
+- Gemma specifics: (1+w) RMSNorm, GeGLU (tanh), rope base 10000, language
+  embeddings scaled by sqrt(width); attention/softmax in fp32
+- verified against the PyTorch reference: cosine 1.00000 on full action chunks
+
+Usage:
+    from pi0_mlx import Pi0MLX
+    m = Pi0MLX.from_pretrained(".")
+    actions = m.predict(images=[cam0, cam1, cam2],
+                        instruction="pick up the cube", state=[...])
+
+Note: the tokenizer is loaded from google/paligemma-3b-pt-224 (gated; accept
+the Gemma terms on Hugging Face and login first), matching upstream lerobot.
+"""
+import json
+import math
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from mlx_vlm.models.paligemma.vision import VisionModel
+from mlx_vlm.models.paligemma import VisionConfig
+
+CHUNK = 50
+MAX_DIM = 32
+NUM_STEPS = 10
+MIN_PERIOD, MAX_PERIOD = 4e-3, 4.0
+N_LAYERS = 18
+HID = 2048
+EXP_HID = 1024
+N_HEADS = 8
+N_KV = 1
+HEAD_DIM = 256
+EPS = 1e-6
+EMBED_SCALE = HID ** 0.5
+IMG_SIZE = 224
+VIS_CFG = {"model_type": "siglip_vision_model", "hidden_size": 1152, "num_hidden_layers": 27,
+           "intermediate_size": 4304, "num_attention_heads": 16, "image_size": 224,
+           "patch_size": 14, "num_channels": 3, "layer_norm_eps": 1e-6}
+
+
+def rms_norm_gemma(x, w):
+    x32 = x.astype(mx.float32)
+    out = x32 * mx.rsqrt(mx.mean(x32 * x32, axis=-1, keepdims=True) + EPS)
+    return (out * (1.0 + w.astype(mx.float32))).astype(x.dtype)
+
+
+def apply_rope(x, positions, base=10000.0):
+    d_half = x.shape[-1] // 2
+    dtype = x.dtype
+    x = x.astype(mx.float32)
+    freq_exp = (2.0 / x.shape[-1]) * mx.arange(d_half, dtype=mx.float32)
+    timescale = mx.power(base, freq_exp)
+    radians = positions[..., None].astype(mx.float32) / timescale[None, None, :]
+    radians = radians[..., None, :]
+    sin, cos = mx.sin(radians), mx.cos(radians)
+    x1, x2 = x[..., :d_half], x[..., d_half:]
+    return mx.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1).astype(dtype)
+
+
+def make_att_2d_masks(pad_masks, att_masks):
+    cs = mx.cumsum(att_masks.astype(mx.int32), axis=1)
+    return (cs[:, None, :] <= cs[:, :, None]) & pad_masks[:, None, :]
+
+
+def sinusoidal_time_emb(time, dim):
+    frac = np.linspace(0.0, 1.0, dim // 2, dtype=np.float64)
+    period = MIN_PERIOD * (MAX_PERIOD / MIN_PERIOD) ** frac
+    scale = 1.0 / period * 2 * np.pi
+    sin_in = scale[None, :] * np.asarray(time, dtype=np.float64)[:, None]
+    return mx.array(np.concatenate([np.sin(sin_in), np.cos(sin_in)], axis=1).astype(np.float32))
+
+
+def resize_with_pad(img_chw, size=IMG_SIZE):
+    """参照実装同等: アスペクト維持バイリニア縮小＋左・上ゼロパディング。"""
+    from PIL import Image
+
+    c, h, w = img_chw.shape
+    ratio = max(w / size, h / size)
+    rh, rw = int(h / ratio), int(w / ratio)
+    pil = Image.fromarray((np.transpose(img_chw, (1, 2, 0)) * 255).clip(0, 255).astype(np.uint8))
+    pil = pil.resize((rw, rh), Image.BILINEAR)
+    arr = np.asarray(pil).astype(np.float32) / 255.0
+    out = np.zeros((size, size, 3), dtype=np.float32)
+    out[size - rh:, size - rw:, :] = arr
+    return out
+
+
+class Pi0MLX:
+    P_LM = "model.paligemma_with_expert.paligemma.model.language_model."
+    P_EXP = "model.paligemma_with_expert.gemma_expert.model."
+    P_VIS = "model.paligemma_with_expert.paligemma.model.vision_tower."
+    P_PROJ = "model.paligemma_with_expert.paligemma.model.multi_modal_projector."
+    P_EMB = "model.paligemma_with_expert.paligemma.lm_head.weight"  # tied weights
+
+    def __init__(self, weights):
+        self.w = weights
+        self.vision = VisionModel(VisionConfig(**VIS_CFG))
+        vis_w = {}
+        for k, v in weights.items():
+            if k.startswith(self.P_VIS):
+                kk = k[len(self.P_VIS):]
+                if kk.endswith("patch_embedding.weight") and v.shape[-1] != 3:
+                    v = v.transpose(0, 2, 3, 1)
+                vis_w[kk] = v
+        self.vision.load_weights(list(vis_w.items()), strict=False)
+        mx.eval(self.vision.parameters())
+
+    @classmethod
+    def from_pretrained(cls, path):
+        path = Path(path)
+        m = cls(mx.load(str(path / "model.safetensors")))
+        from transformers import AutoTokenizer
+
+        m.tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
+        return m
+
+    def lm(self, i, name):
+        return self.w[f"{self.P_LM}layers.{i}.{name}"]
+
+    def exp(self, i, name):
+        return self.w[f"{self.P_EXP}layers.{i}.{name}"]
+
+    def _attn(self, mask2d, q, k, v):
+        B, Lk = k.shape[0], k.shape[1]
+        groups = N_HEADS // N_KV
+        k = mx.repeat(k[:, :, :, None, :], groups, axis=3).reshape(B, Lk, N_HEADS, HEAD_DIM)
+        v = mx.repeat(v[:, :, :, None, :], groups, axis=3).reshape(B, Lk, N_HEADS, HEAD_DIM)
+        q32 = q.astype(mx.float32).transpose(0, 2, 1, 3)
+        k32 = k.astype(mx.float32).transpose(0, 2, 1, 3)
+        att = (q32 @ k32.transpose(0, 1, 3, 2)) * (HEAD_DIM ** -0.5)
+        att = mx.where(mask2d[:, None, :, :], att, mx.finfo(mx.float32).min)
+        probs = mx.softmax(att, axis=-1).astype(v.dtype)
+        out = probs @ v.transpose(0, 2, 1, 3)
+        return out.transpose(0, 2, 1, 3).reshape(B, -1, N_HEADS * HEAD_DIM)
+
+    def _layer(self, get, i, h, mask2d, pos, cache=None, fill=False):
+        hn = rms_norm_gemma(h, get(i, "input_layernorm.weight"))
+        B, L = hn.shape[:2]
+        q = (hn @ get(i, "self_attn.q_proj.weight").T).reshape(B, L, -1, HEAD_DIM)
+        k = (hn @ get(i, "self_attn.k_proj.weight").T).reshape(B, L, -1, HEAD_DIM)
+        v = (hn @ get(i, "self_attn.v_proj.weight").T).reshape(B, L, -1, HEAD_DIM)
+        q = apply_rope(q, pos)
+        k = apply_rope(k, pos)
+        if fill:
+            cache[i] = (k, v)
+        elif cache is not None:
+            k = mx.concatenate([cache[i][0], k], axis=1)
+            v = mx.concatenate([cache[i][1], v], axis=1)
+        att = self._attn(mask2d, q, k, v)
+        out = att @ get(i, "self_attn.o_proj.weight").T + h
+        res = out
+        on = rms_norm_gemma(out, get(i, "post_attention_layernorm.weight"))
+        gate = on @ get(i, "mlp.gate_proj.weight").T
+        up = on @ get(i, "mlp.up_proj.weight").T
+        return (nn.gelu_approx(gate) * up) @ get(i, "mlp.down_proj.weight").T + res
+
+    def embed_prefix(self, imgs224, tokens, lang_mask):
+        embs, pads, atts = [], [], []
+        vdtype = self.vision.vision_model.embeddings.patch_embedding.weight.dtype
+        for img in imgs224:  # [1,224,224,3] in [-1,1]
+            feat = self.vision(img.astype(vdtype))
+            feat = feat[0] if isinstance(feat, tuple) else feat
+            if feat.ndim == 2:
+                feat = feat[None]
+            feat = feat @ self.w[self.P_PROJ + "linear.weight"].T + self.w[self.P_PROJ + "linear.bias"]
+            embs.append(feat.astype(mx.bfloat16))
+            pads.append(mx.ones(feat.shape[:2], dtype=mx.bool_))
+            atts += [0] * feat.shape[1]
+        lang = (self.w[self.P_EMB][tokens].astype(mx.float32) * EMBED_SCALE).astype(mx.bfloat16)
+        embs.append(lang)
+        pads.append(lang_mask.astype(mx.bool_))
+        atts += [0] * lang.shape[1]
+        embs = mx.concatenate(embs, axis=1)
+        pads = mx.concatenate(pads, axis=1)
+        atts = mx.array(atts, dtype=mx.int32)[None, :]
+        return embs, pads, mx.broadcast_to(atts, (embs.shape[0], atts.shape[1]))
+
+    def embed_suffix(self, state32, x_t, time):
+        state_emb = state32 @ self.w["model.state_proj.weight"].T + self.w["model.state_proj.bias"]
+        act = x_t @ self.w["model.action_in_proj.weight"].T + self.w["model.action_in_proj.bias"]
+        t_emb = sinusoidal_time_emb(time, EXP_HID).astype(act.dtype)
+        at = mx.concatenate([act, mx.broadcast_to(t_emb[:, None, :], act.shape)], axis=2)
+        at = at @ self.w["model.action_time_mlp_in.weight"].T + self.w["model.action_time_mlp_in.bias"]
+        at = nn.silu(at)
+        at = at @ self.w["model.action_time_mlp_out.weight"].T + self.w["model.action_time_mlp_out.bias"]
+        embs = mx.concatenate([state_emb[:, None, :], at], axis=1).astype(mx.bfloat16)
+        pads = mx.ones(embs.shape[:2], dtype=mx.bool_)
+        atts = mx.array([1, 1] + [0] * (CHUNK - 1), dtype=mx.int32)[None, :]
+        return embs, pads, mx.broadcast_to(atts, (embs.shape[0], atts.shape[1]))
+
+    def sample_actions(self, imgs224, tokens, lang_mask, state32, noise=None):
+        if noise is None:
+            noise = mx.random.normal((state32.shape[0], CHUNK, MAX_DIM))
+        prefix, pads, atts = self.embed_prefix(imgs224, tokens, lang_mask)
+        mask2d = make_att_2d_masks(pads, atts)
+        pos = mx.cumsum(pads.astype(mx.int32), axis=1) - 1
+        cache = {}
+        h = prefix
+        for i in range(N_LAYERS):
+            h = self._layer(self.lm, i, h, mask2d, pos, cache=cache, fill=True)
+        P = pads.shape[1]
+        offset = mx.sum(pads.astype(mx.int32), axis=-1)[:, None]
+        x_t = noise
+        dt = -1.0 / NUM_STEPS
+        for step in range(NUM_STEPS):
+            t = 1.0 + step * dt
+            suffix, s_pads, s_atts = self.embed_suffix(state32, x_t, [t] * x_t.shape[0])
+            L = s_pads.shape[1]
+            mask_full = mx.concatenate(
+                [mx.broadcast_to(pads[:, None, :], (s_pads.shape[0], L, P)),
+                 make_att_2d_masks(s_pads, s_atts)], axis=2)
+            pos_s = offset + mx.cumsum(s_pads.astype(mx.int32), axis=1) - 1
+            h = suffix
+            for i in range(N_LAYERS):
+                h = self._layer(self.exp, i, h, mask_full, pos_s, cache=cache, fill=False)
+            h = rms_norm_gemma(h, self.w[self.P_EXP + "norm.weight"])
+            out = h[:, -CHUNK:].astype(mx.float32)
+            v_t = out @ self.w["model.action_out_proj.weight"].T + self.w["model.action_out_proj.bias"]
+            x_t = x_t + dt * v_t
+            mx.eval(x_t)
+        return x_t
+
+    def predict(self, images, instruction, state, action_dim=None, noise=None):
+        imgs224 = []
+        for im in images:
+            arr = np.asarray(im).astype(np.float32)
+            if arr.max() > 1.5:
+                arr = arr / 255.0
+            hwc = resize_with_pad(np.transpose(arr, (2, 0, 1))) * 2.0 - 1.0
+            imgs224.append(mx.array(hwc[None]))
+        enc = self.tokenizer(instruction.rstrip("\n") + "\n", padding="max_length",
+                             max_length=48, return_tensors="np")
+        state = np.asarray(state, dtype=np.float32)[None]
+        action_dim = action_dim or state.shape[1]
+        state32 = np.zeros((1, MAX_DIM), dtype=np.float32)
+        state32[:, : state.shape[1]] = state
+        chunk = self.sample_actions(imgs224, mx.array(enc["input_ids"]),
+                                    mx.array(enc["attention_mask"]).astype(mx.bool_),
+                                    mx.array(state32), noise=noise)
+        return np.array(chunk[0, :, :action_dim])
+
+
+if __name__ == "__main__":
+    import argparse
+
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--images", nargs="+", required=True)
+    ap.add_argument("--instruction", required=True)
+    ap.add_argument("--state", default="0,0,0,0,0,0")
+    ap.add_argument("--out", default="actions.npy")
+    args = ap.parse_args()
+    from PIL import Image
+
+    model = Pi0MLX.from_pretrained(Path(__file__).parent)
+    imgs = [np.asarray(Image.open(p).convert("RGB")) for p in args.images]
+    state = [float(x) for x in args.state.split(",")]
+    actions = model.predict(imgs, args.instruction, state)
+    np.save(args.out, actions)
+    print(f"action chunk {actions.shape} -> {args.out}")
+    print("first action:", actions[0].round(4))