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

@@ -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 |

tests/policies/groot/artifacts/.gitignore

@@ -0,0 +1,2 @@
+# Local-only parity artifacts (regenerated via dump_original_n1_7.py); never committed.
+*.npz

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 <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)
+
+    # 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()

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
+``<this dir>/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/<this file> -> 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 <ckpt> --out-dir groot_vs_lerobot/artifacts --device cuda"
+        "env:\n  .venv-original/bin/python tests/policies/groot/dump_original_n1_7.py "
+        "--ckpt <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)