groot_mlx.py

"""GR00T N1.7 (nvidia/GR00T-N1.7-3B) inference runtime for Apple Silicon / MLX.

NON-COMMERCIAL USE ONLY: the base model is released under the NVIDIA license
included in this repository (Section 3.3: research or evaluation purposes only).

Faithful port of the LeRobot reference (which is parity-tested against NVIDIA's
original gr00t package):
- backbone: Cosmos-Reason2-2B (Qwen3-VL) truncated to 16 layers, pre-final-norm
  features, fp32 execution (matching the reference runtime's fp32 upcast)
- action head: 32-layer AlternateVL-DiT (AdaLN conditioning, alternating
  image/text cross-attention) + per-embodiment encoders, 4-step flow matching
- verified end-to-end vs PyTorch reference: cosine 1.000000 / max diff 0.00027

Preprocessing (Qwen3-VL chat template, image packing, state normalization) is
produced with the LeRobot pipeline — see preprocess_lerobot.py.

Usage:
    from groot_mlx import GrootMLX
    m = GrootMLX.from_pretrained(".")
    chunk = m.sample_actions_from_processed("processed.pt")  # -> (B, 40, 132)
"""
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.qwen3_vl import Model, ModelConfig
from mlx_vlm.models.qwen3_vl.language import create_attention_mask

N_BB_LAYERS = 16


# ---------------- バックボーン（Qwen3-VL 16層・pre-norm出力・fp32） ----------------

def build_backbone(ckpt_dir):
    ckpt_dir = Path(ckpt_dir)
    cfg_dict = json.load(open(ckpt_dir / "cosmos_config.json"))
    cfg_dict["text_config"]["num_hidden_layers"] = N_BB_LAYERS
    cfg = ModelConfig.from_dict(cfg_dict)
    cfg.text_config = type(cfg.text_config).from_dict(cfg_dict["text_config"])
    cfg.vision_config = type(cfg.vision_config).from_dict(cfg_dict["vision_config"])
    model = Model(cfg)

    weights = {}
    for f in ckpt_dir.glob("model-*.safetensors"):
        weights.update(mx.load(str(f)))
    pref = "backbone.model."
    bb = {k[len(pref):]: v for k, v in weights.items() if k.startswith(pref)}
    bb = {k: v for k, v in bb.items() if ".layers." not in k or
          int(k.split(".layers.")[1].split(".")[0]) < N_BB_LAYERS or "visual" in k}
    # torch参照はbf16チェックポイントをfp32昇格して実行する。ここを揃えないと一致しない
    bb = {k: (v.astype(mx.float32) if v.dtype == mx.bfloat16 else v) for k, v in bb.items()}
    bb = model.sanitize(bb)
    if hasattr(model.vision_tower, "sanitize"):
        bb = model.vision_tower.sanitize(bb)
    model.load_weights(list(bb.items()), strict=False)
    mx.eval(model.parameters())
    return model


def backbone_features(model, input_ids, pixel_values, image_grid_thw):
    feats = model.get_input_embeddings(input_ids, pixel_values, image_grid_thw=image_grid_thw)
    lm = model.language_model.model
    h = feats.inputs_embeds
    mask = create_attention_mask(h, [None] * len(lm.layers))
    position_ids = feats.position_ids
    position_embeddings = None
    if position_ids is not None and not lm.layers[0].self_attn.rotary_emb.fused_apply:
        position_embeddings = lm.layers[0].self_attn.rotary_emb(h, position_ids)
    dse = feats.deepstack_visual_embeds
    vpm = feats.visual_pos_masks
    for layer_idx, layer in enumerate(lm.layers):
        h = layer(h, mask, None, position_ids, position_embeddings)
        if dse is not None and layer_idx in range(len(dse)):
            h = lm._deepstack_process(h, vpm, dse[layer_idx])
    return h  # GR00Tはpre-final-norm出力を使う


# ---------------- アクションヘッド（DiT + per-embodimentエンコーダ） ----------------

def layer_norm(x, w=None, b=None, eps=1e-5):
    x32 = x.astype(mx.float32)
    mu = mx.mean(x32, axis=-1, keepdims=True)
    var = mx.var(x32, axis=-1, keepdims=True)
    out = (x32 - mu) * mx.rsqrt(var + eps)
    if w is not None:
        out = out * w.astype(mx.float32) + b.astype(mx.float32)
    return out.astype(x.dtype)


def timesteps_embed(t_vals, num_channels=256, max_period=10000.0, shift=1.0):
    half = num_channels // 2
    exponent = -math.log(max_period) * np.arange(half, dtype=np.float64) / (half - shift)
    emb = np.exp(exponent)[None, :] * np.asarray(t_vals, dtype=np.float64)[:, None]
    return mx.array(np.concatenate([np.cos(emb), np.sin(emb)], axis=1).astype(np.float32))


def sinusoidal_time_action(timesteps_2d, dim):
    half = dim // 2
    exponent = -np.arange(half, dtype=np.float32) * (math.log(10000.0) / half)
    freqs = np.asarray(timesteps_2d, dtype=np.float32)[..., None] * np.exp(exponent)[None, None, :]
    return mx.array(np.concatenate([np.sin(freqs), np.cos(freqs)], axis=-1))


class GrootHead:
    def __init__(self, weights, cfg):
        self.w = {k[len("action_head."):]: v for k, v in weights.items() if k.startswith("action_head.")}
        d = cfg["diffusion_model_cfg"]
        self.N_DIT = d["num_layers"]
        self.N_HEADS = d["num_attention_heads"]
        self.HEAD_DIM = d["attention_head_dim"]
        self.ATTEND_TEXT_N = cfg.get("attend_text_every_n_blocks") or 2
        self.VL_CFG = cfg["vl_self_attention_cfg"]
        self.HORIZON = cfg["action_horizon"]
        self.NUM_STEPS = cfg["num_inference_timesteps"]
        self.BUCKETS = cfg["num_timestep_buckets"]

    def cs_linear(self, prefix, x, cat_id):
        return x @ self.w[prefix + ".W"][cat_id] + self.w[prefix + ".b"][cat_id][:, None, :]

    def cs_mlp(self, prefix, x, cat_id):
        h = mx.maximum(self.cs_linear(prefix + ".layer1", x, cat_id), 0)
        return self.cs_linear(prefix + ".layer2", h, cat_id)

    def attention(self, prefix, q_in, kv_in, n_heads, head_dim, mask=None):
        B, Lq = q_in.shape[:2]
        Lk = kv_in.shape[1]
        q = q_in @ self.w[prefix + ".to_q.weight"].T + self.w[prefix + ".to_q.bias"]
        k = kv_in @ self.w[prefix + ".to_k.weight"].T + self.w[prefix + ".to_k.bias"]
        v = kv_in @ self.w[prefix + ".to_v.weight"].T + self.w[prefix + ".to_v.bias"]
        q = q.reshape(B, Lq, n_heads, head_dim).transpose(0, 2, 1, 3)
        k = k.reshape(B, Lk, n_heads, head_dim).transpose(0, 2, 1, 3)
        v = v.reshape(B, Lk, n_heads, head_dim).transpose(0, 2, 1, 3)
        att = (q @ k.transpose(0, 1, 3, 2)) * (head_dim ** -0.5)
        if mask is not None:
            att = mx.where(mask[:, None, None, :], att, mx.finfo(mx.float32).min)
        probs = mx.softmax(att.astype(mx.float32), axis=-1).astype(v.dtype)
        out = (probs @ v).transpose(0, 2, 1, 3).reshape(B, Lq, n_heads * head_dim)
        return out @ self.w[prefix + ".to_out.0.weight"].T + self.w[prefix + ".to_out.0.bias"]

    def ff(self, prefix, x):
        h = x @ self.w[prefix + ".net.0.proj.weight"].T + self.w[prefix + ".net.0.proj.bias"]
        h = nn.gelu_approx(h)
        return h @ self.w[prefix + ".net.2.weight"].T + self.w[prefix + ".net.2.bias"]

    def vl_self_attention(self, feats):
        n, hd = self.VL_CFG["num_attention_heads"], self.VL_CFG["attention_head_dim"]
        h = feats
        for i in range(self.VL_CFG["num_layers"]):
            p = f"vl_self_attention.transformer_blocks.{i}"
            hn = layer_norm(h, self.w[p + ".norm1.weight"], self.w[p + ".norm1.bias"])
            h = self.attention(p + ".attn1", hn, hn, n, hd) + h
            hn = layer_norm(h, self.w[p + ".norm3.weight"], self.w[p + ".norm3.bias"])
            h = self.ff(p + ".ff", hn) + h
        return h

    def ada_norm(self, prefix, x, temb):
        t = nn.silu(temb) @ self.w[prefix + ".linear.weight"].T + self.w[prefix + ".linear.bias"]
        scale, shift = mx.split(t, 2, axis=1)
        return layer_norm(x) * (1 + scale[:, None]) + shift[:, None]

    def dit(self, sa_embs, vl_embs, temb, image_mask, bb_att_mask):
        image_att = image_mask & bb_att_mask
        text_att = (~image_mask) & bb_att_mask
        h = sa_embs
        for i in range(self.N_DIT):
            p = f"model.transformer_blocks.{i}"
            hn = self.ada_norm(p + ".norm1", h, temb)
            if i % 2 == 1:
                h = self.attention(p + ".attn1", hn, hn, self.N_HEADS, self.HEAD_DIM) + h
            else:
                m = text_att if i % (2 * self.ATTEND_TEXT_N) == 0 else image_att
                h = self.attention(p + ".attn1", hn, vl_embs, self.N_HEADS, self.HEAD_DIM, mask=m) + h
            h = self.ff(p + ".ff", layer_norm(h)) + h
        t = nn.silu(temb) @ self.w["model.proj_out_1.weight"].T + self.w["model.proj_out_1.bias"]
        shift, scale = mx.split(t, 2, axis=1)
        h = layer_norm(h, eps=1e-6) * (1 + scale[:, None]) + shift[:, None]
        return h @ self.w["model.proj_out_2.weight"].T + self.w["model.proj_out_2.bias"]

    def timestep_encoder(self, t_disc):
        e = timesteps_embed(t_disc)
        e = e @ self.w["model.timestep_encoder.timestep_embedder.linear_1.weight"].T \
            + self.w["model.timestep_encoder.timestep_embedder.linear_1.bias"]
        e = nn.silu(e)
        return e @ self.w["model.timestep_encoder.timestep_embedder.linear_2.weight"].T \
            + self.w["model.timestep_encoder.timestep_embedder.linear_2.bias"]

    def action_encoder(self, actions, t_disc, cat_id):
        B, T, _ = actions.shape
        a = self.cs_linear("action_encoder.W1", actions, cat_id)
        tt = np.broadcast_to(np.asarray(t_disc, dtype=np.float32)[:, None], (B, T))
        te = sinusoidal_time_action(tt, a.shape[-1]).astype(a.dtype)
        x = self.cs_linear("action_encoder.W2", mx.concatenate([a, te], axis=-1), cat_id)
        x = x * mx.sigmoid(x)
        return self.cs_linear("action_encoder.W3", x, cat_id)

    def get_action(self, bb_raw, image_mask, bb_att_mask, state, embodiment_id, noise):
        feats = layer_norm(bb_raw, self.w["vlln.weight"], self.w["vlln.bias"])
        vl_embs = self.vl_self_attention(feats)
        state_feats = self.cs_mlp("state_encoder", state.reshape(state.shape[0], 1, -1), embodiment_id)
        x_t = noise
        dt = 1.0 / self.NUM_STEPS
        for step in range(self.NUM_STEPS):
            t_disc = [int(step / float(self.NUM_STEPS) * self.BUCKETS)] * x_t.shape[0]
            temb = self.timestep_encoder(t_disc)
            act = self.action_encoder(x_t, t_disc, embodiment_id)
            act = act + self.w["position_embedding.weight"][mx.arange(act.shape[1])][None]
            pred = self.dit(mx.concatenate([state_feats, act], axis=1), vl_embs, temb,
                            image_mask, bb_att_mask)
            pred = self.cs_mlp("action_decoder", pred, embodiment_id)
            x_t = x_t + dt * pred[:, -self.HORIZON:]
            mx.eval(x_t)
        return x_t


# ---------------- 統合ランタイム ----------------

class GrootMLX:
    def __init__(self, path="."):
        self.dir = Path(path)
        self.cfg = json.load(open(self.dir / "config.json"))
        self.backbone = build_backbone(self.dir)
        weights = {}
        for f in self.dir.glob("model-*.safetensors"):
            weights.update(mx.load(str(f)))
        self.head = GrootHead(weights, self.cfg)
        self.image_token_id = json.load(open(self.dir / "cosmos_config.json"))["image_token_id"]

    @classmethod
    def from_pretrained(cls, path):
        return cls(path)

    def sample_actions(self, input_ids, pixel_values, image_grid_thw, attention_mask,
                       state, embodiment_id, noise=None, seed=42):
        bb = backbone_features(self.backbone, input_ids, pixel_values, image_grid_thw).astype(mx.float32)
        image_mask = input_ids == self.image_token_id
        att = attention_mask.astype(mx.bool_)
        if noise is None:
            np.random.seed(seed)
            noise = mx.array(np.random.standard_normal(
                (state.shape[0], self.head.HORIZON, 132)).astype("float32"))
        return np.array(self.head.get_action(bb, image_mask, att, state, embodiment_id, noise))

    def sample_actions_from_processed(self, processed_pt, noise=None, seed=42):
        """preprocess_lerobot.pyの出力（.pt）から行動チャンクを生成する。"""
        import torch

        ref = torch.load(processed_pt, weights_only=False)
        proc = ref.get("proc", ref)

        def get(k):
            return proc[k] if k in proc else ref[k]

        if noise is None and "noise" in ref:
            noise = mx.array(ref["noise"].cpu().float().numpy())
        return self.sample_actions(
            mx.array(proc["input_ids"].cpu().numpy()),
            mx.array(proc["pixel_values"].cpu().float().numpy()),
            mx.array(proc["image_grid_thw"].cpu().numpy()),
            mx.array(proc["attention_mask"].cpu().numpy()),
            mx.array(get("state").cpu().float().numpy()),
            mx.array(get("embodiment_id").cpu().numpy()),
            noise=noise, seed=seed,
        )


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser()
    ap.add_argument("--processed", required=True, help="preprocess_lerobot.pyの出力.pt")
    ap.add_argument("--out", default="actions.npy")
    args = ap.parse_args()
    m = GrootMLX.from_pretrained(Path(__file__).parent)
    chunk = m.sample_actions_from_processed(args.processed)
    np.save(args.out, chunk)
    print(f"action chunk {chunk.shape} -> {args.out}")