"""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}")