lerobot/tests/policies/groot/utils/dump_original_n1_7.py

#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License").
"""Producer (run in the ORIGINAL gr00t env): dump original GR00T N1.7 outputs + inputs.

The original NVIDIA ``gr00t`` package pins ``transformers==4.57.3`` (py3.10) and its
model-config dataclasses are incompatible with the ``transformers==5.x`` that the
LeRobot GR00T N1.7 integration requires. The two implementations therefore cannot be
imported in the same Python process. To keep the parity comparison FAIR, we run the
original model in its native env here and serialize, PER EMBODIMENT TAG:

  * the RAW observation fed to the original processor (per-camera uint8 frames,
    per-key state vectors, the language instruction), so the LeRobot side can also
    run its OWN preprocessor on identical raw inputs and compare collated tensors,
  * the exact pre-processed/collated model inputs (so the LeRobot side consumes the
    byte-identical tensors -- same image preprocessing, tokenization, normalization),
  * the random seed used right before the flow-matching sampler,
  * the raw ``action_pred`` tensor returned by ``model.get_action`` (normalized space,
    before any per-implementation action decoding).

Inputs are built GENERICALLY from the checkpoint metadata (no per-tag hardcoding):
state keys + dims come from ``statistics.json``; video + language keys come from the
processor's per-embodiment modality configs. This lets us test many embodiment tags
from the SAME checkpoint and confirm the LeRobot integration is not overfit to
``libero_sim``.

The companion pytest (run in the LeRobot env) loads each .npz and asserts parity
twice: the collated inputs + seed are replayed through the LeRobot GR00T N1.7 model
(model parity), and the raw observation is replayed through LeRobot's own
preprocessor pipeline and compared against the collated inputs (preprocessor parity).

Usage:
    .venv-original/bin/python tests/policies/groot/utils/dump_original_n1_7.py \
        --ckpt <path-to-GR00T-N1.7-LIBERO/libero_10> \
        --out-dir tests/policies/groot/artifacts \
        [--tags libero_sim,oxe_droid_relative_eef_relative_joint,...] \
        [--device cuda] [--seed 42]

If --tags is omitted, every embodiment present in the checkpoint statistics is dumped.
"""

import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch

IMAGE_SIZE = 256
BATCH_SIZE = 2
PROMPT = "pick up the black bowl and place it on the plate"


def load_statistics(ckpt: str) -> dict:
    with open(os.path.join(ckpt, "statistics.json")) as f:
        return json.load(f)


def make_observation(seed: int, video_keys, lang_key, state_spec):
    """Build a dummy observation dict generically from the embodiment metadata."""
    rng = np.random.default_rng(seed)
    video = {
        k: rng.integers(0, 256, (BATCH_SIZE, 1, IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8)
        for k in video_keys
    }
    # One ndarray per state key, shape (B, T=1, key_dim); dim taken from statistics.
    # Keys with dim 0 (e.g. disabled eef on some embodiments) are still emitted as
    # present-but-empty so the processor's state transform finds every expected key.
    state = {k: rng.standard_normal((BATCH_SIZE, 1, dim)).astype(np.float32) for k, dim in state_spec}
    language = {lang_key: [[PROMPT] for _ in range(BATCH_SIZE)]}
    return {"video": video, "state": state, "language": language}


def dump_one_tag(policy, fair_model, tag, modality_cfg, state_spec, args, out_path):
    from gr00t.data.types import MessageType

    video_keys = modality_cfg["video"].modality_keys
    lang_key = modality_cfg["language"].modality_keys[0]
    observation = make_observation(args.seed, video_keys, lang_key, state_spec)

    # Snapshot the RAW observation exactly as fed to the original processor below. The
    # consumer's preprocessor-parity case replays it through LeRobot's own preprocessor
    # and compares the resulting collated tensors against the "in::" ones saved further
    # down. raw_state_keys records the checkpoint modality-key order, which is the
    # concatenation order of the flat LeRobot ``observation.state`` vector.
    spec_keys = [key for key, _ in state_spec]
    state_modality = modality_cfg.get("state")
    state_keys = [key for key in state_modality.modality_keys if key in spec_keys] if state_modality else []
    state_keys += [key for key in spec_keys if key not in state_keys]
    raw_language = [
        str(item[0]) if isinstance(item, (list, tuple)) else str(item)
        for item in observation["language"][lang_key]
    ]
    raw_flat = {f"raw::video.{key}": arr.copy() for key, arr in observation["video"].items()}
    raw_flat.update({f"raw::state.{key}": arr.copy() for key, arr in observation["state"].items()})
    raw_flat["raw::language"] = np.array(raw_language, dtype=object)
    raw_flat["raw_video_keys"] = np.array([str(key) for key in video_keys], dtype=object)
    raw_flat["raw_state_keys"] = np.array([str(key) for key in state_keys], dtype=object)

    # Point the policy preprocessing at this embodiment (mirrors Gr00tPolicy.__init__).
    policy.embodiment_tag = type(policy.embodiment_tag)(tag)
    policy.modality_configs = {
        k: v for k, v in policy.processor.get_modality_configs()[tag].items() if k != "rl_info"
    }
    policy.language_key = policy.modality_configs["language"].modality_keys[0]

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    unbatched = policy._unbatch_observation(observation)
    processed = []
    for obs in unbatched:
        vla = policy._to_vla_step_data(obs)
        processed.append(policy.processor([{"type": MessageType.EPISODE_STEP.value, "content": vla}]))
    collated = policy.collate_fn(processed)

    def to_dev(x):
        if isinstance(x, torch.Tensor) and torch.is_floating_point(x):
            return x.to(args.device, torch.float32)
        if isinstance(x, torch.Tensor):
            return x.to(args.device)
        if isinstance(x, dict):
            return {k: to_dev(v) for k, v in x.items()}
        return x

    collated = {k: to_dev(v) for k, v in collated.items()}

    torch.manual_seed(args.seed)
    with torch.inference_mode():
        out = fair_model.get_action(**collated)
    action_pred = out["action_pred"].float().cpu().numpy()

    flat, meta = {}, {}

    def flatten(prefix, obj):
        if isinstance(obj, torch.Tensor):
            arr = obj.float().cpu().numpy() if torch.is_floating_point(obj) else obj.cpu().numpy()
            flat[f"in::{prefix}"] = arr
            meta[f"in::{prefix}"] = str(obj.dtype)
        elif isinstance(obj, dict):
            for k, v in obj.items():
                flatten(f"{prefix}.{k}" if prefix else k, v)
        elif isinstance(obj, (list, tuple)):
            flat[f"in::{prefix}"] = np.array(obj, dtype=object)
        else:
            flat[f"in::{prefix}"] = np.array(obj)

    flatten("", collated)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez(
        out_path,
        action_pred=action_pred,
        seed=np.array(args.seed),
        device=np.array(args.device),
        embodiment_tag=np.array(tag),
        meta_keys=np.array(list(meta.keys()), dtype=object),
        meta_dtypes=np.array(list(meta.values()), dtype=object),
        **raw_flat,
        **flat,
    )
    print(f"[{tag}] action_pred {action_pred.shape} -> {out_path.name} ({os.path.getsize(out_path)} B)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--out-dir", required=True, help="directory for per-tag .npz files")
    ap.add_argument("--tags", default="", help="comma-separated embodiment tags (default: all in stats)")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    from gr00t.policy.gr00t_policy import Gr00tPolicy
    from transformers import AutoConfig, AutoModel

    stats = load_statistics(args.ckpt)
    requested = [t.strip() for t in args.tags.split(",") if t.strip()] or list(stats.keys())

    # Load the policy once (for its processor/preprocessing) on any valid tag.
    bootstrap_tag = "libero_sim" if "libero_sim" in stats else requested[0]
    policy = Gr00tPolicy(embodiment_tag=bootstrap_tag, model_path=args.ckpt, device=args.device)
    all_modality = policy.processor.get_modality_configs()

    # Load a FAIR model (SDPA + fp32) once and reuse across tags. Otherwise the
    # original checkpoint default (flash_attention_2 + bf16) introduces kernel/rounding
    # noise vs the LeRobot env (which has no flash_attn and runs SDPA).
    cfg = AutoConfig.from_pretrained(args.ckpt, trust_remote_code=True)
    cfg.use_flash_attention = False
    cfg.load_bf16 = False
    fair_model = AutoModel.from_pretrained(args.ckpt, config=cfg, trust_remote_code=True)
    fair_model.to(device=args.device, dtype=torch.float32)
    fair_model.eval()

    out_dir = Path(args.out_dir)
    done, skipped = [], []
    for tag in requested:
        if tag not in stats or tag not in all_modality:
            print(f"[skip] {tag}: not present in checkpoint statistics/modality configs")
            skipped.append(tag)
            continue
        state_spec = [(k, len(v["min"])) for k, v in stats[tag]["state"].items()]
        try:
            dump_one_tag(
                policy,
                fair_model,
                tag,
                all_modality[tag],
                state_spec,
                args,
                out_dir / f"original_n1_7_{tag}.npz",
            )
            done.append(tag)
        except Exception as exc:  # noqa: BLE001
            print(f"[fail] {tag}: {type(exc).__name__}: {exc}")
            skipped.append(tag)

    print(f"\nDumped {len(done)} tags: {done}")
    if skipped:
        print(f"Skipped/failed {len(skipped)} tags: {skipped}")


if __name__ == "__main__":
    main()