From a9a78f72feff9474614f6095b6f89b89c5d873a8 Mon Sep 17 00:00:00 2001 From: nv-sachdevkartik Date: Fri, 12 Jun 2026 07:47:11 +0000 Subject: [PATCH] test(groot): self-contained parity test + in-repo producer + docs - Rename test_groot_n1_7_vs_original.py -> test_groot_vs_original.py - Make the test self-contained: producer script (dump_original_n1_7.py) now lives next to the test; default artifact dir is repo-relative (tests/policies/groot/artifacts/), overridable via GROOT_N1_7_PARITY_DIR. The test only reads artifacts and skips if absent -- it never creates external dirs. - Heavy .npz artifacts (~6-9MB each) are gitignored and regenerated by the producer; never committed. - Drop the verbose 'MULTIPLE EMBODIMENTS' docstring block (kept a one-line note). - Document the parity procedure in the groot policy README (docs/source/policy_groot_README.md). - Rename test fn test_groot_n1_7_get_action_parity -> test_groot_get_action_parity. 9/9 embodiments still pass (max|diff| < 3e-6, fp32 eps). --- docs/source/policy_groot_README.md | 74 +++++++ tests/policies/groot/artifacts/.gitignore | 2 + tests/policies/groot/dump_original_n1_7.py | 198 ++++++++++++++++++ ..._original.py => test_groot_vs_original.py} | 55 +++-- 4 files changed, 299 insertions(+), 30 deletions(-) create mode 100644 tests/policies/groot/artifacts/.gitignore create mode 100644 tests/policies/groot/dump_original_n1_7.py rename tests/policies/groot/{test_groot_n1_7_vs_original.py => test_groot_vs_original.py} (77%) diff --git a/docs/source/policy_groot_README.md b/docs/source/policy_groot_README.md index f843f005e..c5fe078c3 100644 --- a/docs/source/policy_groot_README.md +++ b/docs/source/policy_groot_README.md @@ -28,3 +28,77 @@ Hugging Face Models: - GR00T N1.7: https://huggingface.co/nvidia/GR00T-N1.7-3B - GR00T N1.7 LIBERO checkpoints: https://huggingface.co/nvidia/GR00T-N1.7-LIBERO + +## Original-vs-LeRobot parity test + +`tests/policies/groot/test_groot_vs_original.py` verifies that this LeRobot +reimplementation of GR00T N1.7 (Qwen3-VL backbone + flow-matching action head) +produces the **same raw model output** (`get_action(...)["action_pred"]`, the +normalized flow-matching prediction) as NVIDIA's original `gr00t` package, given +byte-identical pre-processed inputs and the same flow-matching seed. It is +parametrized over every embodiment tag present in the checkpoint. + +### Why two environments + +The original `gr00t` package pins `transformers==4.57.3` (Python 3.10); this +integration requires `transformers>=5.x` (Qwen3-VL). Under 5.x, `PretrainedConfig` +is itself a defaulted dataclass, so the original config dataclasses fail to import +(`non-default argument follows default argument`). The two implementations therefore +**cannot be imported in the same Python process**. + +So the test uses a **producer / consumer** split across two venvs: + +1. **Producer** — `tests/policies/groot/dump_original_n1_7.py`, run in the *original* + gr00t venv. For each embodiment it builds dummy inputs generically from the + checkpoint metadata (state dims from `statistics.json`; camera/language keys from + the processor modality configs), runs the original model, and saves the exact + collated inputs + raw `action_pred` to one `.npz` per tag. +2. **Consumer** — the pytest above, run in the *LeRobot* venv. It discovers every + `.npz`, replays the byte-identical inputs through the LeRobot model with the same + seed, and asserts the outputs match. + +### Fairness controls + +- **Same pre-processed inputs** — the original processor's `input_ids`, + `pixel_values`, `image_grid_thw`, `attention_mask`, `state`, `embodiment_id` are + fed verbatim to the LeRobot model (no re-tokenization / re-normalization). +- **Same precision + attention kernel** — both sides run **fp32 + SDPA**. The + original defaults to `use_flash_attention=True` (flash_attention_2 + bf16); the + producer forces SDPA + fp32. (With the defaults the gap is ~3e-2 — pure + kernel/rounding noise, not an implementation difference.) +- **Same flow-matching seed** — fixed (42) right before sampling on both sides. + +### How to run + +```bash +# Resolve a local checkpoint (GR00T-N1.7-LIBERO / libero_10) +CKPT=$(python - <<'PY' +import os +from huggingface_hub import snapshot_download +print(os.path.join(snapshot_download("nvidia/GR00T-N1.7-LIBERO", + allow_patterns=["libero_10/*"]), "libero_10")) +PY +) + +# 1) Produce the original-side artifacts for all embodiments (original gr00t venv, CUDA) +CUDA_VISIBLE_DEVICES=0 /path/to/Isaac-GR00T/.venv-original/bin/python \ + tests/policies/groot/dump_original_n1_7.py \ + --ckpt "$CKPT" --out-dir tests/policies/groot/artifacts --device cuda --seed 42 + +# 2) Run the parity test (LeRobot venv) — one parametrized case per embodiment +CUDA_VISIBLE_DEVICES=0 GROOT_PARITY_DEVICE=cuda \ + uv run pytest tests/policies/groot/test_groot_vs_original.py -v -s +``` + +The `.npz` artifacts are local-only (gitignored, ~6–9 MB each) and are regenerated by +the producer; they are never committed. The test **skips** (does not fail) on CI or +when the checkpoint / artifacts are absent. + +#### Env knobs (all optional) + +| Var | Default | Purpose | +|---|---|---| +| `GROOT_N1_7_PARITY_DIR` | `tests/policies/groot/artifacts` | directory of per-tag `.npz` artifacts | +| `GROOT_N1_7_LIBERO_CKPT` | auto (HF cache) | override checkpoint dir | +| `GROOT_PARITY_DEVICE` | `cuda` if available | `cpu` or `cuda` | +| `GROOT_PARITY_ATOL` / `GROOT_PARITY_RTOL` | `1e-3` | comparison tolerance | diff --git a/tests/policies/groot/artifacts/.gitignore b/tests/policies/groot/artifacts/.gitignore new file mode 100644 index 000000000..0d0f35a4d --- /dev/null +++ b/tests/policies/groot/artifacts/.gitignore @@ -0,0 +1,2 @@ +# Local-only parity artifacts (regenerated via dump_original_n1_7.py); never committed. +*.npz diff --git a/tests/policies/groot/dump_original_n1_7.py b/tests/policies/groot/dump_original_n1_7.py new file mode 100644 index 000000000..c5dd0888e --- /dev/null +++ b/tests/policies/groot/dump_original_n1_7.py @@ -0,0 +1,198 @@ +#!/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 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, replays the identical +inputs + seed through the LeRobot GR00T N1.7 model, and asserts the outputs match. + +Usage: + .venv-original/bin/python tests/policies/groot/dump_original_n1_7.py \ + --ckpt \ + --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) + + # 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), + **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() diff --git a/tests/policies/groot/test_groot_n1_7_vs_original.py b/tests/policies/groot/test_groot_vs_original.py similarity index 77% rename from tests/policies/groot/test_groot_n1_7_vs_original.py rename to tests/policies/groot/test_groot_vs_original.py index 1e9ae6451..f0bef142b 100644 --- a/tests/policies/groot/test_groot_n1_7_vs_original.py +++ b/tests/policies/groot/test_groot_vs_original.py @@ -16,22 +16,12 @@ """Parity test: original NVIDIA GR00T N1.7 vs the GR00T N1.7 integration in LeRobot. -This is the N1.7 analogue of ``test_groot_vs_original.py`` (which covers N1.5/GR1). -It verifies that the *self-contained* LeRobot reimplementation of the GR00T N1.7 -action head + Qwen3-VL backbone produces the SAME raw model output (``action_pred``, -the normalized flow-matching prediction before any action decoding) as NVIDIA's -original ``gr00t`` package, given byte-identical pre-processed inputs and the same -flow-matching seed. - -MULTIPLE EMBODIMENTS (anti-overfitting) ---------------------------------------- -The comparison is parametrized over EVERY embodiment tag present in the checkpoint -(``libero_sim`` plus the cross-embodiment tags it was trained with: oxe_droid, -real_g1, the real_r1_pro_sharpa family, and the xdof family). Inputs for each tag are -built generically from the checkpoint metadata (state dims from ``statistics.json``, -camera/language keys from the processor modality configs), so passing on all of them -shows the LeRobot integration is correct across the model's full embodiment space and -not merely tuned for ``libero_sim``. +Verifies that the self-contained LeRobot reimplementation of the GR00T N1.7 action +head + Qwen3-VL backbone produces the SAME raw model output (``action_pred``, the +normalized flow-matching prediction before any action decoding) as NVIDIA's original +``gr00t`` package, given byte-identical pre-processed inputs and the same +flow-matching seed. The comparison is parametrized over every embodiment tag present +in the checkpoint. WHY TWO ENVIRONMENTS -------------------- @@ -43,16 +33,16 @@ argument follows default argument"). The two implementations therefore CANNOT be imported in the same Python process. To keep the comparison fair, the original outputs + the exact collated inputs are -produced once per embodiment in the original ``gr00t`` env via -``groot_vs_lerobot/scripts/dump_original_n1_7.py`` and saved to per-tag ``.npz`` -files. This test discovers those artifacts, replays the identical inputs through the -LeRobot model, and compares. +produced once per embodiment in the original ``gr00t`` env via the companion script +``dump_original_n1_7.py`` (next to this file) and saved to per-tag ``.npz`` files. +This test discovers those artifacts, replays the identical inputs through the LeRobot +model, and compares. This test is LOCAL-only and skips on CI, when ``gr00t``-side prerequisites are not -present, or when no artifact has been generated. No hardcoded paths: the artifact -directory comes from ``GROOT_N1_7_PARITY_DIR`` (default: -``groot_vs_lerobot/artifacts`` alongside the repo root). See -``groot_vs_lerobot/README.md`` for the full run procedure. +present, or when no artifact has been generated. By default it looks for artifacts in +``/artifacts/``; override with ``GROOT_N1_7_PARITY_DIR``. See the +"Original-vs-LeRobot parity test" section of ``src/lerobot/policies/groot/README.md`` +for the full run procedure. """ import os @@ -80,12 +70,17 @@ _ARTIFACT_SUFFIX = ".npz" def _artifact_dir() -> Path: + """Directory holding the per-embodiment .npz artifacts. + + Self-contained by default: a sibling ``artifacts/`` directory next to this test. + Override with ``GROOT_N1_7_PARITY_DIR`` (e.g. to point at a scratch location). + The directory is read-only here -- it is populated by ``dump_original_n1_7.py`` + run in the original gr00t environment; the test never creates it. + """ env = os.environ.get("GROOT_N1_7_PARITY_DIR") if env: return Path(env) - # repo_root/tests/policies/groot/ -> repo_root.parent holds groot_vs_lerobot/ - repo_root = Path(__file__).resolve().parents[3] - return repo_root.parent / "groot_vs_lerobot" / "artifacts" + return Path(__file__).resolve().parent / "artifacts" def _discover_artifacts() -> list[tuple[str, Path]]: @@ -183,12 +178,12 @@ _ARTIFACTS = _discover_artifacts() not _ARTIFACTS, reason=( "No GR00T N1.7 parity artifacts found. Generate them first in the original gr00t " - "env:\n .venv-original/bin/python groot_vs_lerobot/scripts/dump_original_n1_7.py " - "--ckpt --out-dir groot_vs_lerobot/artifacts --device cuda" + "env:\n .venv-original/bin/python tests/policies/groot/dump_original_n1_7.py " + "--ckpt --out-dir tests/policies/groot/artifacts --device cuda" ), ) @pytest.mark.parametrize("embodiment_tag,artifact", _ARTIFACTS, ids=[t for t, _ in _ARTIFACTS]) -def test_groot_n1_7_get_action_parity(embodiment_tag, artifact, lerobot_model): +def test_groot_get_action_parity(embodiment_tag, artifact, lerobot_model): """Raw model.get_action(action_pred) parity per embodiment: original vs LeRobot.""" original_action, flat_inputs = _load_artifact(artifact) model_inputs = _unflatten(flat_inputs)