admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 02:06:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(policies): add Gr00t N1.7 policy (#3922)

Browse Source 
* Add GR00T N1.7 support

Add GR00T N1.7 policy configuration, checkpoint compatibility, processor parity, LIBERO documentation, and focused tests.

Co-authored-by: Ryan Halabi <ryhalabi@nvidia.com>

* Move Groot processor compatibility into Groot loader

* Restore GR00T Flash Attention install guidance

* Allow Groot fake RTC chunk prefetch

* Fix GR00T N1.7 RTC action decoding

* Trim GR00T N1.7 RTC chunks to valid horizon

* Ignore padded GR00T N1.7 RTC prefix rows

* removed n1.5 dependency

* removed remaining N1.5 traces

* groot: auto-enable LIBERO gripper action transform for libero_sim

GR00T N1.7 emits gripper in [0,1] but LIBERO expects [-1,1]. The decode
transform existed but was never auto-enabled for embodiment_tag=libero_sim,
so the policy scored 0% on LIBERO eval. Auto-set it in __post_init__ (still
overridable). LIBERO Spatial eval: 0% -> 98%.

* Reconnect GR00T relative action processors

* groot: remove dead N1.5 code (eagle2_hg_model, flow_matching_action_head, action_encoder)

N1.7 backbone is nvidia/Cosmos-Reason2-2B via Qwen3VLForConditionalGeneration,
not Eagle2 — eagle2_hg_model/ had zero refs outside its own dir.

GR00TN17ActionHead (groot_n1_7.py) re-implements MultiEmbodimentActionEncoder +
CategorySpecificLinear + swish + SinusoidalPositionalEncoding locally, so
flow_matching_action_head.py (N1.5 FlowmatchingActionHead) and its sole
dependency action_encoder.py are dead. Verified: no src/ or tests/ reference.

Removed (~2037 LOC):
- eagle2_hg_model/ (4 files, ~1575 LOC)
- action_head/flow_matching_action_head.py (408 LOC)
- action_head/action_encoder.py (54 LOC)

cross_attention_dit.py KEPT (DiT/AlternateVLDiT/SelfAttentionTransformer live in N1.7).

* groot: reuse lerobot get_device_from_parameters instead of inline lookup

modeling_groot.py duplicated next(self.parameters()).device twice. LeRobot
ships get_device_from_parameters in policies/utils.py (used by diffusion,
vqbet, tdmpc, gaussian_actor). Reuse it for consistency with the framework.

* groot: fix stale Eagle VLM docstring in processor (N1.7 uses Qwen3-VL backbone)

Addresses checker nit: processor_groot.py docstring still described the N1.5
Eagle VLM path with eagle_content/eagle_* keys that no longer exist in the code.

* test(groot): add N1.7 original-vs-LeRobot output parity test

Verifies the LeRobot GR00T N1.7 integration produces equivalent raw
action_pred to NVIDIA Isaac-GR00T for the same checkpoint, inputs, seed,
precision (fp32) and attention kernel (SDPA): max|diff|=8.9e-7 on the
libero_sim embodiment (GR00T-N1.7-LIBERO/libero_10).

The two impls pin incompatible transformers majors (orig 4.57.3 vs
LeRobot 5.x) and cannot share a process, so the original outputs + exact
collated inputs are produced out-of-process and loaded from an .npz. The
test skips on CI / when the checkpoint or artifact are absent.

* test(groot): parametrize N1.7 parity across all checkpoint embodiments

Generalize the original-vs-LeRobot N1.7 output-parity test from a single
libero_sim case to every embodiment tag in the checkpoint (libero_sim, oxe_droid,
real_g1, the real_r1_pro_sharpa family, and the xdof family). Inputs are built
generically from checkpoint metadata; the test discovers per-tag .npz artifacts
and runs one parametrized case each, loading the LeRobot model once via a fixture.

All 9 embodiments match the original to fp32 epsilon (max|diff| < 3e-6), confirming
the integration is correct across the model's full embodiment space and not overfit
to libero_sim.

* 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(groot): drop WHY TWO ENVIRONMENTS block from parity test docstring

* test(groot): move parity producer into utils/ package

Mirror the tests/policies/pi0_pi05/utils convention: move dump_original_n1_7.py into
a tests/policies/groot/utils/ package (with __init__.py) and update all path
references in the test docstring/skip-message and the policy README.

* test(groot): adopt test_groot_lerobot for GR00T N1.7, drop N1.5

The test loaded MODEL_PATH='aractingi/bimanual-handover-groot-10k', an N1.5
checkpoint (config base_model_path=nvidia/GR00T-N1.5-3B, no model_version). On
load, model_version defaults to n1.7 while the base path infers n1.5, so the
version-consistency guard in GrootConfig.__post_init__ raised ValueError and both
test_lerobot_groot_inference and test_lerobot_groot_forward_pass failed. N1.5 is no
longer a supported model_version.

Adopt the test for N1.7:
- MODEL_PATH -> nvidia/GR00T-N1.7-3B (root-level sharded safetensors; loads via
  GrootPolicy.from_pretrained as a base N1.7 model).
- Embodiment tag 'gr1' (N1.5) -> 'gr1_unified' (valid N1.7 tag from the checkpoint
  embodiment_id.json), via a single EMBODIMENT_TAG constant.
- DUMMY_ACTION_HORIZON 16 -> 40 to match N1.7's native action-chunk size.
- Docstrings/labels updated to 'GR00T N1.7'.

Both tests run and pass on CUDA; full tests/policies/groot/ suite is
73 passed / 0 failed / 0 skipped.

* docs(groot): document the N1.5 removal and the N1.7 parity test

- groot.mdx: breaking-change warning and migration path (pin lerobot==0.5.1 to
  keep N1.5, or move to N1.7); the dead `huggingface-cli download` is replaced
  with `hf download`.
- policy_groot_README.md: N1.5 removal note, updated paper / model-card links,
  and the two-comparison (model parity + preprocessor parity) description of
  the original-vs-LeRobot test, including the raw-observation artifacts and
  recorded seed.

* fix(groot): N1.7 backbone loading and DiT parameter-count logging

- select_layer default tracks the N1.7-3B checkpoint value (16); real
  checkpoint loads still override it from config.json.
- get_backbone_cls recognizes Cosmos-Reason2 / Qwen3-VL backbones by name and
  warns (instead of silently assuming) when an unrecognized backbone is loaded
  only on the strength of backbone_model_type='qwen'.
- 'revision' pins the GR00T checkpoint repo only and is no longer forwarded
  into the unrelated backbone repo load; pin the backbone via
  transformers_loading_kwargs instead.
- DiT / SelfAttentionTransformer parameter counts go through logging.debug
  instead of print().

* fix(groot): N1.7 config defaults, N1.5 rejection, and processor/model runtime fixes

Covers the GR00T N1.7 source trio (configuration, processor, model wrapper).

Config:
- GrootConfig defaults are the N1.7 values; explicitly passed legacy N1.5-era
  values (chunk_size=50, max_state_dim=64, ...) are remapped with a warning
  instead of silently.
- action_decode_transform gains an 'auto' sentinel so an explicit 'none'
  opt-out wins over the libero_sim default and survives save/load round-trips.
- action_delta_indices is cached on the inputs that determine it.
- Legacy N1.5 checkpoints/configs (tokenizer_assets_repo, model_type/
  architectures/eagle backbone markers) are rejected with a single clear
  error pointing to lerobot==0.5.1.

Processor:
- GrootN17ActionDecodeStep handles the 2-D (B, D) actions delivered by sync
  select_action (relative eef/non-eef decode in eval/record flows).
- Postprocessor falls back to dataset stats when a raw checkpoint lacks the
  configured embodiment tag; raw-state cache is per-instance, not
  process-global; caller overrides (device, rename_map) are honored on the
  raw-checkpoint branch.
- Camera/modality-key mismatches warn (including the zero-match fallback);
  deprecated Qwen2VLImageProcessorFast replaced with Qwen2VLImageProcessor;
  removed N1.5 processor steps are stubbed to raise the removal guidance and
  the action-unpack step is re-registered as _v2.

Model:
- Flash-attention probe is diagnostic-only; forward raises on a missing loss;
  print() replaced with logging; N1.5 base-path mismatch includes the
  removal guidance.

* fix(groot): skip normalization overrides for training

* fix(groot): GPU/tensor N1.7 image preprocessing + resize to trained resolution

GR00T training was dataloader-bound (0->100->0 GPU-utilization sawtooth).
GrootN17VLMEncodeStep ran the Qwen3-VL image processor per frame on PIL images
on the single CPU main-loop thread, and that cost is timed inside dataloading_s
(preprocessor(batch) runs in the main process, not the dataloader workers), so
adding workers cannot hide it.

- Feed the torchvision-backed Qwen3-VL processor (C,H,W) uint8 tensors instead
  of a per-frame Image.fromarray PIL roundtrip, and run resize/normalize/patchify
  on config.device (GPU) when available. Bit-identical on CPU when no resize is
  configured; with a resize only the PIL->torchvision bicubic backend differs
  (<2/255 per pixel). The use_albumentations path stays PIL/cv2; reload on a box
  without the saved device falls back to CPU.

- Default image_target_size/crop to the N1.7 backbone's training geometry
  (256x256 / 230x230) when a checkpoint ships no image sizing (checkpoint_assets
  is None, e.g. finetuning nvidia/GR00T-N1.7-3B via repo-id with a new
  embodiment). Previously image_target_size=None disabled the resize, so
  full-resolution frames were patchified into ~4.7x more vision tokens than the
  model was trained on -- inflating dataloading_s (patchify) and update_s (VLM
  sequence) and skewing the input distribution. Checkpoints that pin their own
  sizing are honored; the default constants are shared with GR00T_N1_7_DEFAULTS.

Net: preprocessing leaves the CPU critical path and the VLM sees the resolution
it was trained on -- faster training/inference and a correct train/serve
distribution. Affects inference too (shared preprocessor); existing checkpoints
still load (backward compatible) but must be retrained to gain the benefits.

* refactor(groot): N1.7 style cleanup (utils, imports, flash-attn, config)

Mechanical refactor of the GR00T N1.7 policy to match the repo's architecture and
style standards. No change to policy algorithm/numerics; only UX/CLI and packaging
changes. Tests are intentionally left untouched (out of scope) and need updating
for the removed `model_version` field.

Cleanup & consolidation:
- Add `groot/utils.py` holding the pure, side-effect-free helpers (JSON I/O, value
  coercion, stat flattening, rot6d/SE3 math, language/batch prep) shared by the
  config and processor layers.
- Remove dead code: the unused `resolve_groot_n1_7_backbone_model` cache-resolver
  cluster, `GR00TN17Config.to_filtered_dict/json`, and the `_copy_default` wrapper.

Imports & execution guards:
- Hoist nested imports to module top; relative imports within the package, absolute
  for external modules. The version-gated Qwen3-VL classes import under the single
  `_transformers_available` guard (transformers is pinned >=5.4, which ships them).
- No import-time side effects: `_register_with_transformers()` now runs in
  `GR00TN17.__init__` (idempotent via `register(exist_ok=True)`), and the N1.5 step
  stubs register lazily before pipeline deserialization (idempotent via the
  registry, no run-once globals).
- Gate optional deps at the point of use with `require_package(..., extra="groot")`.

Dependencies & docs:
- Drop `flash-attn` (and its build-only dep `ninja`) from the `groot` extra; default
  to SDPA (numerically equivalent) with opt-in via `--policy.use_flash_attention`.
  Un-comment `lerobot[groot]` in the `all` extra and regenerate `uv.lock`.
- Rewrite the `groot.mdx` install section: flash-attn is a purely optional,
  user-managed optimization that LeRobot neither installs nor requires.

Config & CLI:
- Surface previously-frozen knobs on `GrootConfig` (plumbed into `GR00TN17Config`;
  no-ops at their defaults): inference — `num_inference_timesteps`, `rtc_ramp_rate`,
  `use_flash_attention`; fine-tuning — `tune_top_llm_layers` (partial-LLM tuning)
  and `tune_vlln` (previously hardwired to True).
- Convert the single-valued `model_version` and `n1_7_backbone_model` fields to
  internal constants.
- Keep `base_model_path`: it is NOT equivalent to `pretrained_path` (raw NVIDIA
  checkpoints have no LeRobot `type` field and load only via `base_model_path`) and
  is genuinely user-tunable.
- Keep the deprecated Isaac-GR00T/N1.5 fields (and the dead LoRA fields) as a
  back-compat block so a v0.5.1 N1.5 `config.json` still parses under draccus and is
  rejected with the friendly N1.5 removal message instead of an opaque decode error.

* Optimize GR00T N1.7 image preprocessing

* Remove PIL fallback from GR00T preprocessing

* Fix GROOT relative action training stats

* Address GROOT relative action review feedback

* Fix GROOT N1.7 relative action stats

* Fix GROOT relative action training stats

* Fix GROOT relative action padding and RTC leftovers

* Reset rollout state after robot episode end

* Revert "Reset rollout state after robot episode end"

This reverts commit 1322f45aec.

* Move GROOT relative stats out of train script

* Guard GR00T relative action stepwise decode

* Match GR00T N1.7 OSS preprocessing and relative actions

* Apply LIBERO action decode override after loading

* Format GR00T OSS parity changes

* chore(policies): add guards, warnings and comments + recover tests n1.5 check

* fix(style): pre-commit

* fix(ci): guard dependecy checks

* chore(groot): move cv2 to the top as its in the default install tag

* chore(policies): add explicit dataset dependecy to gr00t implementation

* fix(test): add guard

* fix(groot): make N1.7 letterbox opt-in

* feat(groot): activate checkpoint-configured N1.7 raw-state dropout during training

Isaac-GR00T applies dual state regularization during fine-tuning: raw-state
zeroing driven by the processor sidecar's state_dropout_prob (0.2 for the
inspected N1.7 checkpoint) plus encoded-feature dropout. Baseline LeRobot kept
the processor in deterministic mode, so the raw-state dropout never activated
(RCA Tier-2 contributor to the LeRobot-trained SO-101 failures).

- GrootN17PackInputsStep: runtime-only 'training' flag + state_dropout_prob;
  whole-sample state zeroing gated on torch.is_grad_enabled() so eval and
  no_grad validation paths are unaffected
- sidecar loader reads state_dropout_prob from processor_config.json
- state_dropout_prob serializes with the step; the training flag intentionally
  does not (reloaded pipelines default to eval, re-enabled only when processors
  are rebuilt with dataset_meta)
- _set_groot_preprocessor_training toggles any dataclass step exposing a
  'training' field on serialized-pipeline reloads

Verification: tests/policies/groot/test_groot_state_dropout.py (4 passed) on
RTX PRO 6000 / CUDA 13.3.

* fix(groot): align N1.7 fine-tuning optimizer/scheduler/precision with Isaac-GR00T

Evidence from the LeRobot-vs-OSS checkpoint comparison: the LeRobot/HF 8k
checkpoint's DiT moved only ~19% as far from base as the OSS-trained one
(0.0547 vs 0.285 relative L2) - undertrained because the scheduler decayed over
a hardcoded 10k steps regardless of --steps, on top of beta1/clip mismatches.

- AdamW betas (0.95, 0.999) -> (0.9, 0.999) and grad_clip_norm 10.0 -> 1.0
  (Isaac defaults)
- scheduler: hardcoded CosineDecayWithWarmup(10k decay, floor 10% peak) ->
  DiffuserSchedulerConfig HF cosine with ceil(max_steps * warmup_ratio) warmup,
  deriving num_training_steps from the outer --steps at runtime
- model_params_fp32 (default true): keep master weights in FP32 and compute
  under BF16 autocast like the native N1.7 recipe (fixes optimizer-update
  numerics vs pure-BF16 params)
- weight-decay grouping via transformers get_parameter_names: biases and norm
  parameters excluded from decay
- restore the TF4 lm_head/embedding weight tie so the unused Qwen LM head stays
  frozen and deduplicated in checkpoints
- action_mask kept in native dtype for the masked flow-matching loss
- drop_n_last_frames: exclude episode tails that cannot supply a complete
  action chunk (Isaac sampler behavior)

Verification: tests/policies/groot/test_groot_training_optim_contract.py
(7 passed) + remaining groot suite 11 passed/5 skipped on RTX PRO 6000 /
CUDA 13.3. Note: tests/policies/groot/test_groot_n1_7.py does not collect on
the base branch (pre-existing ImportError, fixed in PR #37).

* feat(groot): train-time random crop for N1.7 (eval keeps center crop)

Isaac-GR00T crops a random crop_fraction window during training and the
deterministic center window at eval, replaying the sampled window across all
camera views of a sample. This contract is unchanged since the N1.5 release
(gr00t/data/transform/video.py: "If mode is 'train', return a random crop
transform. If mode is 'eval', return a center crop transform.") and mirrors
LeRobot's own Diffusion/VQBeT crop_is_random pattern. The LeRobot N1.7 port
used the eval center crop for training too, so the fine-tuned projector/DiT
never sees frame borders and trains on a single fixed appearance point.

Scope: crop geometry ONLY - no color jitter, no new dependencies. The random
window is plain numpy slicing inside the existing cv2 eval transform:

- _transform_n1_7_image_for_vlm_albumentations gains crop_position=(y, x)
  fractions; None keeps the center crop byte-identical to before (verified
  by test)
- GrootN17VLMEncodeStep gains a runtime-only 'training' flag (never
  serialized; reloaded pipelines default to eval); training samples ONE
  window per sample and reuses it across (timestep, view) frames - Isaac's
  cross-view consistency
- gated on torch.is_grad_enabled() so no_grad validation and frozen-eval
  paths are unaffected
- wired via dataset_meta is not None in make_groot_pre_post_processors and
  the existing _set_groot_preprocessor_training on serialized reloads

Verification: tests/policies/groot/test_groot_train_random_crop.py (8 passed:
center-crop bit-exactness with crop_position=None, corner/center windows,
cross-view replay, train!=eval, no_grad gating, seed reproducibility,
serialization contract) + groot suite 23 passed / 5 skipped on RTX PRO 6000 /
CUDA 13.3.

* docs(groot): update Training & hardware Evaluation commands

Replace the multi-GPU accelerate-launch Training snippet with the current
single-command 'uv run lerobot-train' N1.7 recipe (relative actions excluding
gripper, bf16, flash attention, chunk/n_action_steps=16, bs64/20k steps).

Replace the bimanual 'Evaluate in your hardware setup' rollout example with the
SO-101 follower RTC 'uv run lerobot-rollout' command (strategy.type=base,
inference.type=rtc, wrist+front cameras, place-the-vial task).

Docs-only; no source/test changes.

* docs(groot): parameterize commands with env vars + fill LIBERO results

- Introduce BASE_MODEL / DATASET_ID / REPO_ID / JOB_NAME / OUTPUT_DIR env vars
  in the training command and reuse OUTPUT_DIR + BASE_MODEL in the rollout cmd.
- Fill the LIBERO benchmark table with GR00T-LeRobot success rates
  (Spatial 94%, Object 98%, Goal 93%, LIBERO 10/Long 90%; avg 93.75%),
  drop the OSS column and XX placeholders. LeRobot-focused.

* docs(groot): drop export block, reference env vars directly

Use $DATASET_ID / $BASE_MODEL / $REPO_ID / $OUTPUT_DIR / $JOB_NAME as
bare placeholders in the commands without concrete export assignments.

* docs(groot): keep BASE_MODEL export in training command

* docs(groot): use literal HF repo IDs for dataset/policy repo_id

Public-facing Hub references (--dataset.repo_id, --policy.repo_id) shown as
concrete IDs; local-only values ($OUTPUT_DIR, $JOB_NAME) stay as placeholders.

* docs(groot): add LIBERO training command example

* docs(groot): remove LIBERO checkpoints subdirectory section

* docs(groot): use $BASE_MODEL for base_model_path in LIBERO eval

* docs(groot): drop hf download step from LIBERO eval, fix intro

* docs(groot): restore suite checkpoint download intro sentence

* docs(groot): remove checkpoint download note above LIBERO eval

* docs(groot): update training and rollout commands with new parameters and dependencies

* Add sample so101 training command

* Remove sample so101 training command

* docs(groot): remove optional Flash Attention setup instructions and update base model path for evaluation

* docs(groot): update training command with  image transformation parameters

* docs(groot): add note on inference.queue_threshold value for stable inference

* chore(style): pre-commit gr00t

* docs(groot): update

* chore(policies): minor details

* fix(groot): license headers + test guards

* chore(policies): fix tests

* docs(groot): relative actions param doc

* chore(policy): address some of the AI review items

---------

Co-authored-by: Andrew Wrenn <awrenn@nvidia.com>
Co-authored-by: Ryan Halabi <ryhalabi@nvidia.com>
Co-authored-by: nv-sachdevkartik <ksachdev@nvidia.com>
Co-authored-by: groot-validation <groot-validation@localhost>
Co-authored-by: johnnynunez <johnnynuca14@gmail.com>
Co-authored-by: lbenhorin <lbenhorin@nvidia.com>
This commit is contained in:
Steven Palma
2026-07-03 21:15:09 +02:00
committed by GitHub
parent e275ea3960
commit 708fa1d189
30 changed files with 8593 additions and 3474 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/groot/test_groot_lerobot.py
+20 -13
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Test script for LeRobot's Groot policy forward and inference passes."""
"""Test script for LeRobot's GR00T N1.7 policy forward and inference passes."""
import gc
import os
@@ -25,6 +25,8 @@ import numpy as np
import pytest
import torch
pytest.importorskip("transformers", reason="groot requires the `groot` extra (transformers)")
from lerobot.policies.groot.configuration_groot import GrootConfig
from lerobot.policies.groot.modeling_groot import GrootPolicy
from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
@@ -33,21 +35,26 @@ from lerobot.types import PolicyAction
from lerobot.utils.device_utils import auto_select_torch_device
from tests.utils import require_cuda
pytest.importorskip("transformers")
pytestmark = pytest.mark.skipif(
os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
reason="This test requires local Groot installation and is not meant for CI",
)
# Define constants for dummy data
# Define constants for dummy data (GR00T N1.7 native conventions).
# N1.7 internally uses a 40-step action chunk, 132-dim state/action, and 256px images
# (see GrootConfig.__post_init__). Use a chunk-sized action horizon so the dummy batch
# matches the model's native action space.
DUMMY_STATE_DIM = 44
DUMMY_ACTION_DIM = 44
DUMMY_ACTION_HORIZON = 16
DUMMY_ACTION_HORIZON = 40
IMAGE_SIZE = 256
DEVICE = auto_select_torch_device()
MODEL_PATH = "aractingi/bimanual-handover-groot-10k"
# GR00T N1.7 checkpoint (N1.5 is no longer supported). The N1.7-3B base model loads
# via GrootPolicy.from_pretrained with root-level sharded safetensors.
MODEL_PATH = "nvidia/GR00T-N1.7-3B"
# Valid N1.7 embodiment tag carried by the checkpoint metadata.
EMBODIMENT_TAG = "gr1_unified"
def cleanup_memory():
@@ -88,13 +95,13 @@ def instantiate_lerobot_groot(
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
"""Instantiate LeRobot Groot policy with preprocessor and postprocessor."""
"""Instantiate LeRobot GR00T N1.7 policy with preprocessor and postprocessor."""
if from_pretrained:
policy = GrootPolicy.from_pretrained(
pretrained_name_or_path=model_path,
strict=False,
)
policy.config.embodiment_tag = "gr1"
policy.config.embodiment_tag = EMBODIMENT_TAG
else:
config = GrootConfig(
base_model_path=model_path,
@@ -102,7 +109,7 @@ def instantiate_lerobot_groot(
chunk_size=DUMMY_ACTION_HORIZON,
image_size=[IMAGE_SIZE, IMAGE_SIZE],
device=DEVICE,
embodiment_tag="gr1",
embodiment_tag=EMBODIMENT_TAG,
)
policy = GrootPolicy(config)
@@ -148,8 +155,8 @@ def create_dummy_data(device=DEVICE):
@require_cuda
def test_lerobot_groot_inference():
"""Test the inference pass (select_action) of LeRobot's Groot policy."""
print("Test: LeRobot Groot Inference Pass")
"""Test the inference pass (select_action) of LeRobot's GR00T N1.7 policy."""
print("Test: LeRobot GR00T N1.7 Inference Pass")
set_seed_all(42)
@@ -181,9 +188,9 @@ def test_lerobot_groot_inference():
@require_cuda
def test_lerobot_groot_forward_pass():
"""Test the forward pass of LeRobot's Groot policy."""
"""Test the forward pass of LeRobot's GR00T N1.7 policy."""
print("\n" + "=" * 50)
print("Test: LeRobot Groot Forward Pass (Training Mode)")
print("Test: LeRobot GR00T N1.7 Forward Pass (Training Mode)")
set_seed_all(42)
tests/policies/groot/test_groot_n1_7.py
+3025
View File
File diff suppressed because it is too large Load Diff
tests/policies/groot/test_groot_n1_7_oss_parity.py
+259
View File
@@ -0,0 +1,259 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import hashlib
import os
from pathlib import Path
import numpy as np
import pytest
import torch
from lerobot.policies.groot.action_head.cross_attention_dit import AlternateVLDiT
from lerobot.policies.groot.groot_n1_7 import GR00TN17
from lerobot.policies.groot.processor_groot import (
GrootN17ActionDecodeStep,
GrootN17PackInputsStep,
GrootN17VLMEncodeStep,
_transform_n1_7_image_for_vlm_albumentations,
)
from lerobot.types import TransitionKey
from lerobot.utils.constants import OBS_STATE
OSS_REFERENCE_COMMIT = "ab88b50c718f6528e1df9dcbaf75865d1b604760"
def _fixture_path(filename: str) -> Path:
fixture_dir = os.environ.get("GROOT_N17_OSS_PARITY_FIXTURE_DIR")
if fixture_dir is None:
pytest.skip("Set GROOT_N17_OSS_PARITY_FIXTURE_DIR to run external OSS parity fixtures.")
path = Path(fixture_dir) / filename
if not path.is_file():
pytest.skip(f"External OSS parity fixture not found: {path}")
return path
def test_groot_n1_7_eval_image_transform_matches_oss_reference():
"""Match the native N1.7 eval transform for a non-square SO-101 frame."""
y, x = np.indices((480, 640), dtype=np.uint16)
image = np.stack(
((x + 3 * y) % 256, (2 * x + y) % 256, (x + 5 * y) % 256),
axis=-1,
).astype(np.uint8)
actual = _transform_n1_7_image_for_vlm_albumentations(
image,
image_crop_size=[230, 230],
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
)
assert actual.shape == (256, 340, 3)
assert hashlib.sha256(actual.tobytes()).hexdigest() == (
"c17e47af68a812aa79db3bb7b64b549ddf10148ac1b204a9686095018561ae9e"
)
def test_groot_n1_7_vlm_chat_content_order_matches_oss_reference():
"""Native OSS places all image items before the language item."""
class RecordingProcessor:
def __init__(self):
self.content_types = None
def apply_chat_template(self, conversation, tokenize, add_generation_prompt):
assert tokenize is False
assert add_generation_prompt is False
self.content_types = [item["type"] for item in conversation[0]["content"]]
return "rendered"
def __call__(self, **kwargs):
return {}
processor = RecordingProcessor()
step = GrootN17VLMEncodeStep(
image_crop_size=[230, 230],
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
use_albumentations=True,
device="cpu",
)
step._proc = processor
transition = {
TransitionKey.OBSERVATION: {
"video": np.zeros((1, 1, 2, 480, 640, 3), dtype=np.uint8),
},
TransitionKey.COMPLEMENTARY_DATA: {"language": ["pick up the vial"]},
}
step(transition)
assert processor.content_types == ["image", "image", "text"]
def test_groot_n1_7_alternate_vl_dit_matches_oss_reference():
"""Run the LeRobot DiT with native OSS weights and identical inputs."""
pytest.importorskip("diffusers")
fixture = torch.load(_fixture_path("alternate_vl_dit_small.pt"), map_location="cpu", weights_only=True)
model = AlternateVLDiT(
output_dim=8,
num_attention_heads=2,
attention_head_dim=4,
num_layers=4,
dropout=0.0,
final_dropout=False,
max_num_positional_embeddings=16,
compute_dtype=torch.float32,
interleave_self_attention=True,
cross_attention_dim=6,
).eval()
model.load_state_dict(fixture["state_dict"], strict=True)
actual = model(
hidden_states=fixture["hidden_states"],
encoder_hidden_states=fixture["encoder_hidden_states"],
timestep=fixture["timestep"],
image_mask=fixture["image_mask"],
backbone_attention_mask=fixture["backbone_attention_mask"],
)
torch.testing.assert_close(actual, fixture["output"], atol=1e-6, rtol=1e-6)
def _state_decode_reference():
fixture = np.load(_fixture_path("state_and_action_decode.npz"))
raw_stats = {
"state": {
"single_arm": {"q01": fixture["state_single_arm_q01"], "q99": fixture["state_single_arm_q99"]},
"gripper": {"q01": fixture["state_gripper_q01"], "q99": fixture["state_gripper_q99"]},
},
"action": {
"single_arm": {"q01": fixture["action_single_arm_q01"], "q99": fixture["action_single_arm_q99"]},
"gripper": {"q01": fixture["action_gripper_q01"], "q99": fixture["action_gripper_q99"]},
},
"relative_action": {
"single_arm": {
"min": fixture["relative_single_arm_min"],
"max": fixture["relative_single_arm_max"],
},
},
}
for modality_stats in raw_stats.values():
for entry in modality_stats.values():
for key, value in entry.items():
if isinstance(value, np.ndarray):
entry[key] = value.tolist()
modality_config = {
"state": {"modality_keys": ["single_arm", "gripper"]},
"action": {
"delta_indices": list(range(16)),
"modality_keys": ["single_arm", "gripper"],
"action_configs": [
{"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
{"rep": "ABSOLUTE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
],
},
}
state_min = np.concatenate((fixture["state_single_arm_q01"], fixture["state_gripper_q01"]))
state_max = np.concatenate((fixture["state_single_arm_q99"], fixture["state_gripper_q99"]))
pack_step = GrootN17PackInputsStep(
normalize_min_max=True,
stats={OBS_STATE: {"min": state_min, "max": state_max}},
raw_stats=raw_stats,
modality_config=modality_config,
use_percentiles=True,
)
raw_state = np.concatenate((fixture["state_single_arm"], fixture["state_gripper"]), axis=-1)
transition = {
TransitionKey.OBSERVATION: {OBS_STATE: torch.from_numpy(raw_state)},
TransitionKey.COMPLEMENTARY_DATA: {},
}
packed = pack_step(transition)
return fixture, raw_stats, modality_config, pack_step, packed
def test_groot_n1_7_state_normalization_matches_oss_checkpoint_reference():
fixture, _raw_stats, _modality_config, _pack_step, packed = _state_decode_reference()
expected = np.concatenate(
(fixture["normalized_state_single_arm"], fixture["normalized_state_gripper"]), axis=-1
)
actual = packed[TransitionKey.OBSERVATION]["state"][:, 0, :6]
torch.testing.assert_close(actual, torch.from_numpy(expected), atol=1e-6, rtol=1e-6)
def test_groot_n1_7_relative_action_decode_matches_oss_checkpoint_reference():
fixture, raw_stats, modality_config, pack_step, _packed = _state_decode_reference()
decode_step = GrootN17ActionDecodeStep(
env_action_dim=6,
raw_stats=raw_stats,
modality_config=modality_config,
use_percentiles=True,
use_relative_action=True,
pack_step=pack_step,
)
decoded = decode_step({TransitionKey.ACTION: torch.from_numpy(fixture["normalized_action"])})[
TransitionKey.ACTION
]
expected = np.concatenate((fixture["decoded_single_arm"], fixture["decoded_gripper"]), axis=-1).astype(
np.float32
)
torch.testing.assert_close(decoded, torch.from_numpy(expected), atol=1e-5, rtol=1e-5)
def test_groot_n1_7_qwen_backbone_matches_oss_checkpoint_reference():
"""Compare the actual 3B checkpoint backbone when explicitly enabled."""
checkpoint = os.environ.get("GROOT_N17_PARITY_CHECKPOINT")
if checkpoint is None:
pytest.skip("Set GROOT_N17_PARITY_CHECKPOINT to run the 3B OSS Qwen parity test.")
if not torch.cuda.is_available():
pytest.skip("The 3B OSS Qwen parity test requires CUDA.")
pytest.importorskip("transformers")
from transformers.feature_extraction_utils import BatchFeature
fixture = torch.load(_fixture_path("qwen_backbone_so101.pt"), map_location="cpu", weights_only=True)
model = GR00TN17.from_pretrained(checkpoint).to(device="cuda", dtype=torch.bfloat16).eval()
backbone_input = BatchFeature(
data={
key.removeprefix("input."): value.to("cuda")
for key, value in fixture.items()
if key.startswith("input.")
}
)
with torch.inference_mode():
actual = model.backbone(backbone_input)
feature_error = (
actual.backbone_features.cpu().float() - fixture["output.backbone_features"].float()
).abs()
# Native OSS and LeRobot use different Torch/Transformers/Flash-Attention releases.
# Require the measured BF16 accumulation envelope while rejecting structural drift.
assert feature_error.mean().item() <= 0.04
assert feature_error.max().item() <= 2.0
torch.testing.assert_close(
actual.backbone_attention_mask.cpu(), fixture["output.backbone_attention_mask"]
)
torch.testing.assert_close(actual.image_mask.cpu(), fixture["output.image_mask"])
tests/policies/groot/test_groot_state_dropout.py
+100
View File
@@ -0,0 +1,100 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Isaac-GR00T N1.7 raw-state dropout training contract.
Isaac-GR00T zeroes the entire proprioceptive state of a sample with probability
``state_dropout_prob`` (configured in the checkpoint's processor sidecar) during
training only. Baseline LeRobot kept the processor deterministic, so this
regularization never activated. These tests pin the train/eval split.
"""
import torch
from lerobot.policies.groot.processor_groot import GrootN17PackInputsStep
from lerobot.types import TransitionKey
from lerobot.utils.constants import OBS_STATE
def _make_transition():
return {
TransitionKey.OBSERVATION: {OBS_STATE: torch.tensor([[1.0, 2.0], [3.0, 4.0]])},
TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move", "Move"]},
}
def test_groot_n1_7_training_applies_raw_state_dropout_before_encoder():
step = GrootN17PackInputsStep(
max_state_dim=4,
max_action_dim=4,
normalize_min_max=False,
training=True,
state_dropout_prob=1.0,
)
output = step(_make_transition())
expected = torch.zeros(2, 1, 4)
torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected)
def test_groot_n1_7_training_state_dropout_is_disabled_under_no_grad():
step = GrootN17PackInputsStep(
max_state_dim=4,
max_action_dim=4,
normalize_min_max=False,
training=True,
state_dropout_prob=1.0,
)
with torch.no_grad():
output = step(_make_transition())
expected = torch.tensor([[[1.0, 2.0, 0.0, 0.0]], [[3.0, 4.0, 0.0, 0.0]]])
torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected)
def test_groot_n1_7_eval_mode_state_dropout_is_inactive():
step = GrootN17PackInputsStep(
max_state_dim=4,
max_action_dim=4,
normalize_min_max=False,
training=False,
state_dropout_prob=1.0,
)
output = step(_make_transition())
expected = torch.tensor([[[1.0, 2.0, 0.0, 0.0]], [[3.0, 4.0, 0.0, 0.0]]])
torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected)
def test_groot_n1_7_pack_step_serializes_dropout_prob_but_not_training_mode():
step = GrootN17PackInputsStep(
max_state_dim=4,
max_action_dim=4,
normalize_min_max=False,
training=True,
state_dropout_prob=0.2,
)
serialized = step.get_config()
restored = GrootN17PackInputsStep(**serialized)
assert "training" not in serialized
assert serialized["state_dropout_prob"] == 0.2
assert restored.training is False
assert restored.state_dropout_prob == 0.2
tests/policies/groot/test_groot_train_random_crop.py
+169
View File
@@ -0,0 +1,169 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Isaac-GR00T N1.7 train-time random crop contract (crop geometry only).
Isaac-GR00T crops a random ``crop_fraction`` window during training and the
deterministic center window at eval, replaying the sampled window across all
camera views of a sample (gr00t/data/transform/video.py, n1.5-release onward:
"If mode is 'train', return a random crop transform. If mode is 'eval', return
a center crop transform."). This mirrors LeRobot's own Diffusion/VQBeT
``crop_is_random`` pattern. Color jitter is intentionally out of scope here.
"""
import random
import numpy as np
import torch
from lerobot.policies.groot.processor_groot import (
GrootN17VLMEncodeStep,
_transform_n1_7_image_for_vlm_albumentations,
)
def _structured_image(h=480, w=640):
yy, xx = np.mgrid[0:h, 0:w]
return np.stack([(xx * 255 / w), (yy * 255 / h), ((xx + yy) * 255 / (h + w))], axis=-1).astype(np.uint8)
def test_crop_position_none_is_bitexact_center_crop():
"""crop_position=None must remain byte-identical to the pre-change eval path."""
img = _structured_image()
ref = _transform_n1_7_image_for_vlm_albumentations(
img,
image_crop_size=None,
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
)
out = _transform_n1_7_image_for_vlm_albumentations(
img,
image_crop_size=None,
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
crop_position=None,
)
np.testing.assert_array_equal(ref, out)
def test_crop_position_center_matches_center_crop():
img = _structured_image()
center = _transform_n1_7_image_for_vlm_albumentations(
img,
image_crop_size=None,
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
crop_position=None,
)
explicit = _transform_n1_7_image_for_vlm_albumentations(
img,
image_crop_size=None,
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
crop_position=(0.5, 0.5),
)
# int-floor center vs rounded positional center may differ by <=1 px of grid
assert center.shape == explicit.shape
diff = np.abs(center.astype(np.int16) - explicit.astype(np.int16))
assert diff.mean() < 3.0
def test_crop_position_corners_differ_from_center():
img = _structured_image()
def crop_at(position):
return _transform_n1_7_image_for_vlm_albumentations(
img,
image_crop_size=None,
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
crop_position=position,
)
center = crop_at(None)
tl = crop_at((0.0, 0.0))
br = crop_at((1.0, 1.0))
assert not np.array_equal(center, tl)
assert not np.array_equal(tl, br)
def _video(img, views=2):
return np.stack([img] * views, axis=0).reshape(1, 1, views, *img.shape)
def _step(training):
return GrootN17VLMEncodeStep(
image_target_size=[256, 256],
shortest_image_edge=256,
crop_fraction=0.95,
use_albumentations=True,
training=training,
)
def test_training_crop_replays_one_window_across_views():
video = _video(_structured_image())
frames = _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0]
np.testing.assert_array_equal(np.asarray(frames[0]), np.asarray(frames[1]))
def test_training_crop_differs_from_eval_center_crop():
video = _video(_structured_image())
random.seed(3) # a draw that is not the exact center
train_frame = np.asarray(
_step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
)
eval_frame = np.asarray(
_step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
)
assert not np.array_equal(train_frame, eval_frame)
def test_training_crop_is_disabled_under_no_grad():
video = _video(_structured_image())
with torch.no_grad():
no_grad_frame = np.asarray(
_step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
)
eval_frame = np.asarray(
_step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
)
np.testing.assert_array_equal(no_grad_frame, eval_frame)
def test_training_mode_is_not_serialized():
step = _step(training=True)
serialized = step.get_config()
assert "training" not in serialized
restored = GrootN17VLMEncodeStep(**serialized)
assert restored.training is False
def test_training_crop_respects_global_seed():
video = _video(_structured_image())
def draw():
random.seed(11)
return np.asarray(
_step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
)
np.testing.assert_array_equal(draw(), draw())
tests/policies/groot/test_groot_training_optim_contract.py
+125
View File
@@ -0,0 +1,125 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Isaac-GR00T N1.7 optimizer/scheduler/precision training contract.
Pins the LeRobot GR00T fine-tuning recipe to the native Isaac-GR00T contract:
AdamW(lr=1e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-5, grad clip 1.0),
HF cosine schedule with ~5% warmup over the actual update count, FP32 master
parameters under BF16 autocast, transformers-style weight-decay grouping, the
frozen LM-head weight tie, and episode-tail exclusion for incomplete chunks.
"""
import pytest
import torch
from lerobot.optim.schedulers import DiffuserSchedulerConfig
from lerobot.policies.groot.configuration_groot import GrootConfig
from lerobot.policies.groot.groot_n1_7 import _tie_unused_qwen_lm_head
from lerobot.policies.groot.modeling_groot import GrootPolicy
def test_groot_n1_7_optimizer_matches_isaac_training_contract():
optimizer = GrootConfig().get_optimizer_preset()
assert optimizer.lr == pytest.approx(1e-4)
assert optimizer.betas == pytest.approx((0.9, 0.999))
assert optimizer.eps == pytest.approx(1e-8)
assert optimizer.weight_decay == pytest.approx(1e-5)
assert optimizer.grad_clip_norm == pytest.approx(1.0)
def test_groot_n1_7_sampler_excludes_incomplete_action_tails():
config = GrootConfig(chunk_size=16, n_action_steps=16)
assert len(config.action_delta_indices) == 16
assert config.drop_n_last_frames == 15
def test_groot_n1_7_scheduler_matches_isaac_hf_cosine_contract():
pytest.importorskip("diffusers", reason="the scheduler preset requires the `groot` extra (diffusers)")
config = GrootConfig(max_steps=20_000)
scheduler_config = config.get_scheduler_preset()
assert isinstance(scheduler_config, DiffuserSchedulerConfig)
assert scheduler_config.name == "cosine"
assert scheduler_config.num_warmup_steps == 1_000
parameter = torch.nn.Parameter(torch.ones(()))
optimizer = torch.optim.AdamW([parameter], lr=config.optimizer_lr)
scheduler = scheduler_config.build(optimizer, num_training_steps=20_000)
lr_factor = scheduler.lr_lambdas[0]
assert lr_factor(0) == pytest.approx(0.0)
assert lr_factor(1_000) == pytest.approx(1.0)
assert lr_factor(10_500) == pytest.approx(0.5)
assert lr_factor(20_000) == pytest.approx(0.0, abs=1e-12)
def test_groot_n1_7_scheduler_rounds_fractional_warmup_up_like_transformers():
scheduler_config = GrootConfig(max_steps=777).get_scheduler_preset()
assert scheduler_config.num_warmup_steps == 39
def test_groot_n1_7_model_parameters_use_fp32_checkpoint_and_optimizer_precision():
module = torch.nn.Module()
module.trainable = torch.nn.Parameter(torch.ones(3, dtype=torch.bfloat16))
module.frozen = torch.nn.Parameter(torch.ones(3, dtype=torch.bfloat16), requires_grad=False)
GrootPolicy._cast_model_parameters_to_fp32(module)
assert module.trainable.dtype == torch.float32
assert module.frozen.dtype == torch.float32
def test_groot_n1_7_ties_unused_qwen_lm_head_to_frozen_input_embeddings():
class DummyQwen(torch.nn.Module):
def __init__(self):
super().__init__()
self.embed_tokens = torch.nn.Embedding(7, 3)
self.lm_head = torch.nn.Linear(3, 7, bias=False)
def get_input_embeddings(self):
return self.embed_tokens
model = DummyQwen()
_tie_unused_qwen_lm_head(model)
assert model.lm_head.weight is model.embed_tokens.weight
assert len(list(model.parameters())) == 1
def test_groot_n1_7_optimizer_groups_match_transformers_weight_decay_rules():
pytest.importorskip(
"transformers", reason="weight-decay grouping requires the `groot` extra (transformers)"
)
module = torch.nn.Module()
module.linear = torch.nn.Linear(3, 2)
module.norm = torch.nn.LayerNorm(2)
module.frozen = torch.nn.Parameter(torch.ones(1), requires_grad=False)
groups = GrootPolicy._build_weight_decay_parameter_groups(module)
assert len(groups) == 2
assert "weight_decay" not in groups[0]
assert groups[1]["weight_decay"] == 0.0
assert groups[0]["params"] == [module.linear.weight]
assert {id(parameter) for parameter in groups[1]["params"]} == {
id(module.linear.bias),
id(module.norm.weight),
id(module.norm.bias),
}
tests/policies/groot/test_groot_vs_original.py
+160 -397
View File
@@ -1,6 +1,6 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
@@ -14,431 +14,194 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Test script to verify Groot policy integration with LeRobot vs the original implementation, only meant to be run locally!"""
"""Parity test: original NVIDIA GR00T N1.7 vs the GR00T N1.7 integration in LeRobot.
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.
To keep the comparison fair, the original outputs + the exact collated inputs are
produced once per embodiment in the original ``gr00t`` env via the companion script
``utils/dump_original_n1_7.py`` (in the ``utils`` package 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. 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 gc
import os
from copy import deepcopy
from typing import Any
from pathlib import Path
import numpy as np
import pytest
import torch
from lerobot.policies.groot.configuration_groot import GrootConfig
from lerobot.policies.groot.modeling_groot import GrootPolicy
from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
from lerobot.processor import PolicyProcessorPipeline
from lerobot.types import PolicyAction
pytest.importorskip("gr00t")
pytest.importorskip("transformers")
pytestmark = pytest.mark.skipif(
os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
reason="This test requires local Groot installation and is not meant for CI",
reason="Requires a local GR00T N1.7 checkpoint + pre-generated artifacts; not for CI.",
)
from lerobot.policies.groot.configuration_groot import GROOT_N1_7 # noqa: E402,F401
from gr00t.data.dataset import ModalityConfig # noqa: E402
from gr00t.data.embodiment_tags import EmbodimentTag # noqa: E402
from gr00t.data.transform.base import ComposedModalityTransform # noqa: E402
from gr00t.model.policy import Gr00tPolicy # noqa: E402
SEED = 42
DEVICE = os.environ.get("GROOT_PARITY_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
ATOL = float(os.environ.get("GROOT_PARITY_ATOL", "1e-3"))
RTOL = float(os.environ.get("GROOT_PARITY_RTOL", "1e-3"))
# GR1 humanoid dimensions (from pretrained model metadata)
# The actual GR1 robot has 44 dimensions for both state and action
# GR00TTransform will pad state to 64 and truncate action to 32
DUMMY_STATE_DIM = 44
DUMMY_ACTION_DIM = 44
DUMMY_ACTION_HORIZON = 16
IMAGE_SIZE = 256
DEVICE = "cpu"
MODEL_PATH = "nvidia/GR00T-N1.5-3B"
GR1_BODY_PARTS = {
"left_arm": 7,
"left_hand": 6,
"left_leg": 6,
"neck": 3,
"right_arm": 7,
"right_hand": 6,
"right_leg": 6,
"waist": 3,
}
# Artifact filenames are original_n1_7_<embodiment_tag>.npz
_ARTIFACT_PREFIX = "original_n1_7_"
_ARTIFACT_SUFFIX = ".npz"
def cleanup_memory():
"""Clean up GPU/MPS memory to prevent OOM errors between tests."""
print("\nCleaning up memory...")
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
if torch.backends.mps.is_available():
torch.mps.empty_cache()
print("Memory cleanup complete.")
def _artifact_dir() -> Path:
"""Directory holding the per-embodiment .npz artifacts.
def set_seed_all(seed: int):
"""Set random seed for all RNG sources to ensure reproducibility."""
import random
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# Set deterministic behavior
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.use_deterministic_algorithms(True, warn_only=True)
def instantiate_lerobot_groot(
from_pretrained: bool = False,
model_path: str = MODEL_PATH,
) -> tuple[
GrootPolicy,
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
"""Instantiate LeRobot Groot policy with preprocessor and postprocessor."""
if from_pretrained:
policy = GrootPolicy.from_pretrained(
pretrained_name_or_path=model_path,
strict=False,
)
policy.config.embodiment_tag = "gr1"
else:
config = GrootConfig(
base_model_path=model_path,
n_action_steps=DUMMY_ACTION_HORIZON,
chunk_size=DUMMY_ACTION_HORIZON,
image_size=[IMAGE_SIZE, IMAGE_SIZE],
device=DEVICE,
embodiment_tag="gr1",
)
policy = GrootPolicy(config)
policy.to(DEVICE)
policy.config.device = DEVICE
preprocessor, postprocessor = make_groot_pre_post_processors(
config=policy.config,
dataset_stats=None, # Pass None for dataset_stats to disable normalization (original GR00T doesn't normalize)
)
return (policy, preprocessor, postprocessor)
def instantiate_original_groot(
from_pretrained: bool = False,
model_path: str = MODEL_PATH,
):
"""Instantiate original Groot policy from NVIDIA's implementation."""
from gr00t.data.transform.concat import ConcatTransform
from gr00t.data.transform.state_action import StateActionToTensor
from gr00t.data.transform.video import VideoToNumpy, VideoToTensor
from gr00t.model.transforms import GR00TTransform
video_keys = ["video.ego_view"]
state_keys = [
"state"
] # Important: Use single concatenated "state" key (not split body parts) to match preprocessing
action_keys = [
"action.left_arm",
"action.right_arm",
"action.left_hand",
"action.right_hand",
"action.left_leg",
"action.right_leg",
"action.neck",
"action.waist",
]
language_keys = ["annotation.human.action.task_description"]
modality_config = {
"video": ModalityConfig(
delta_indices=[0], # Current frame only
modality_keys=video_keys,
),
"state": ModalityConfig(
delta_indices=[0],
modality_keys=state_keys,
),
"action": ModalityConfig(
delta_indices=list(range(DUMMY_ACTION_HORIZON)),
modality_keys=action_keys,
),
"language": ModalityConfig(
delta_indices=[0],
modality_keys=language_keys,
),
}
modality_transform = ComposedModalityTransform(
transforms=[
VideoToTensor(apply_to=video_keys),
VideoToNumpy(apply_to=video_keys), # Convert to numpy (GR00TTransform expects numpy arrays)
# State is already a single concatenated key, so no StateActionToTensor needed
# Convert action from numpy to tensor
StateActionToTensor(apply_to=action_keys),
# Concatenate only video and actions (state is already single key)
ConcatTransform(
video_concat_order=video_keys,
state_concat_order=[], # Empty:state is already single key
action_concat_order=action_keys,
),
GR00TTransform(
max_state_dim=64,
max_action_dim=32,
state_horizon=1,
action_horizon=DUMMY_ACTION_HORIZON,
training=False,
),
]
)
policy = Gr00tPolicy(
model_path=model_path,
embodiment_tag=EmbodimentTag.GR1,
modality_config=modality_config,
modality_transform=modality_transform,
device=DEVICE,
)
return policy, modality_config, modality_transform
def create_dummy_data(device=DEVICE):
"""Create dummy data for testing both implementations."""
batch_size = 2
prompt = "Pick up the red cube and place it in the bin"
state = torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=device)
batch = {
"observation.state": state,
"action": torch.randn(
batch_size,
DUMMY_ACTION_HORIZON,
DUMMY_ACTION_DIM,
dtype=torch.float32,
device=device, # Action ground truth (for training)
),
"observation.images.ego_view": torch.rand(
batch_size,
3,
IMAGE_SIZE,
IMAGE_SIZE,
dtype=torch.float32,
device=device, # Images in [0, 1] range as expected by LeRobot
),
"task": [prompt for _ in range(batch_size)],
}
return batch
def convert_lerobot_to_original_format(batch, modality_config):
"""Convert LeRobot batch format to original Groot format.
The original Groot expects observations in this format:
{
"video.<camera_name>": np.ndarray (T, H, W, C) or (B, T, H, W, C)
"state.<state_component>": np.ndarray (T, D) or (B, T, D)
"action.<action_component>": np.ndarray (T, D) or (B, T, D)
"annotation.<annotation_type>": str or list[str]
}
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 ``utils/dump_original_n1_7.py``
run in the original gr00t environment; the test never creates it.
"""
# Original Groot expects (T, H, W, C) format for images
# LeRobot has (B, C, H, W) format, so we need to convert
observation = {}
for img_key in ["ego_view"]:
lerobot_key = f"observation.images.{img_key}"
if lerobot_key in batch:
img = batch[lerobot_key]
# Convert from (B, C, H, W) to (B, T=1, H, W, C)
img_np = img.permute(0, 2, 3, 1).unsqueeze(1).cpu().numpy()
# Convert [0, 1] to [0, 255] uint8 as expected by original
img_np = (img_np * 255).astype(np.uint8)
observation[f"video.{img_key}"] = img_np
# Important: The Original's GR00TTransform expects "state" as (B, T, D), not split body parts
if "observation.state" in batch:
state = batch["observation.state"]
state_np = state.unsqueeze(1).cpu().numpy() # (B, 1, D)
observation["state"] = state_np
if "action" in batch:
action = batch["action"]
action_np = action.cpu().numpy()
start_idx = 0
for part_name, part_dim in GR1_BODY_PARTS.items():
end_idx = start_idx + part_dim
observation[f"action.{part_name}"] = action_np[:, :, start_idx:end_idx]
start_idx = end_idx
if "task" in batch:
task_list = batch["task"]
# GR00TTransform expects language with (B, T) shape for batched data
# Create a (B, T=1) array where each element is the string directly
bsz = len(task_list)
task_array = np.empty((bsz, 1), dtype=object)
for i in range(bsz):
task_array[i, 0] = task_list[i] # Assign string directly to each (i, 0) position
observation["annotation.human.action.task_description"] = task_array
return observation
env = os.environ.get("GROOT_N1_7_PARITY_DIR")
if env:
return Path(env)
return Path(__file__).resolve().parent / "artifacts"
def test_groot_original_vs_lerobot_pretrained():
"""Test Groot original implementation vs LeRobot implementation with pretrained weights."""
print("Test: Groot Original vs LeRobot with Pretrained Weights (Inference)")
def _discover_artifacts() -> list[tuple[str, Path]]:
"""Return [(embodiment_tag, npz_path), ...] for every dumped artifact."""
d = _artifact_dir()
if not d.is_dir():
return []
out = []
for p in sorted(d.glob(f"{_ARTIFACT_PREFIX}*{_ARTIFACT_SUFFIX}")):
tag = p.name[len(_ARTIFACT_PREFIX) : -len(_ARTIFACT_SUFFIX)]
out.append((tag, p))
return out
set_seed_all(42)
lerobot_policy, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_groot(
from_pretrained=True
def _resolve_checkpoint() -> str:
env = os.environ.get("GROOT_N1_7_LIBERO_CKPT")
if env:
if not Path(env).exists():
pytest.skip(f"GROOT_N1_7_LIBERO_CKPT={env} does not exist")
return env
try:
from huggingface_hub import snapshot_download
root = snapshot_download(
"nvidia/GR00T-N1.7-LIBERO",
local_files_only=True,
allow_patterns=["libero_10/*"],
)
except Exception as exc: # noqa: BLE001
pytest.skip(f"GR00T N1.7 LIBERO checkpoint not available locally: {exc}")
ckpt = Path(root) / "libero_10"
if not (ckpt / "config.json").exists():
pytest.skip(f"GR00T N1.7 LIBERO checkpoint incomplete at {ckpt}")
return str(ckpt)
def _load_artifact(path: Path):
data = np.load(path, allow_pickle=True)
original_action = torch.from_numpy(data["action_pred"]).float()
dtypes = dict(zip(data["meta_keys"].tolist(), data["meta_dtypes"].tolist(), strict=False))
inputs = {}
for key in data.files:
if not key.startswith("in::"):
continue
name = key[4:]
arr = data[key]
t = torch.from_numpy(np.asarray(arr))
declared = dtypes.get(key, "")
if "int" in declared or "long" in declared:
t = t.long()
inputs[name] = t
return original_action, inputs
def _unflatten(inputs: dict[str, torch.Tensor]) -> dict:
"""Rebuild the nested model-input dict from dot-prefixed flat keys."""
nested: dict = {}
for dotted, value in inputs.items():
parts = dotted.split(".")
cur = nested
for p in parts[:-1]:
cur = cur.setdefault(p, {})
cur[parts[-1]] = value
return nested.get("inputs", nested)
@pytest.fixture(scope="module")
def lerobot_model():
"""Load the LeRobot GR00T N1.7 model once (fp32 + SDPA) and reuse across tags."""
ckpt = _resolve_checkpoint()
from lerobot.policies.groot.groot_n1_7 import GR00TN17
model = GR00TN17.from_pretrained(
ckpt,
tune_llm=False,
tune_visual=False,
tune_projector=False,
tune_diffusion_model=False,
tune_vlln=False,
transformers_loading_kwargs={"trust_remote_code": True},
)
original_policy, modality_config, modality_transform = instantiate_original_groot(from_pretrained=True)
# fp32 + SDPA on both sides: bf16 + differing attention kernels otherwise introduce
# ~1e-2 numerical noise unrelated to the implementations.
model.compute_dtype = "float32"
model.config.compute_dtype = model.compute_dtype
model.to(device=DEVICE, dtype=torch.float32)
model.eval()
return model
batch = create_dummy_data()
batch_lerobot = deepcopy(batch)
print("\n[LeRobot] Running inference...")
lerobot_policy.eval()
batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
_ARTIFACTS = _discover_artifacts()
# Important: Reset seed immediately before inference to ensure identical RNG state
torch.manual_seed(42)
with torch.no_grad():
lerobot_actions = lerobot_policy.select_action(batch_lerobot_processed)
@pytest.mark.skipif(
not _ARTIFACTS,
reason=(
"No GR00T N1.7 parity artifacts found. Generate them first in the original gr00t "
"env:\n .venv-original/bin/python tests/policies/groot/utils/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_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)
print("\n[Original] Running inference...")
original_policy.model.eval()
observation = convert_lerobot_to_original_format(batch, modality_config)
original_obs_transformed = modality_transform(deepcopy(observation))
# Align the flow-matching RNG exactly as the producer did (seed right before sampling).
torch.manual_seed(SEED)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(SEED)
with torch.inference_mode():
out = lerobot_model.get_action(model_inputs)
lerobot_action = out["action_pred"].float().cpu()
# Important: Reset seed immediately before inference to ensure identical RNG state
torch.manual_seed(42)
t = min(original_action.shape[1], lerobot_action.shape[1])
d = min(original_action.shape[2], lerobot_action.shape[2])
original_action = original_action[:, :t, :d]
lerobot_action = lerobot_action[:, :t, :d]
with torch.no_grad():
original_model_output = original_policy.model.get_action(original_obs_transformed)
original_actions_raw = original_model_output["action_pred"] # [2, 16, 32]
# Take first timestep
original_actions = original_actions_raw[:, 0, :].to(lerobot_actions.device).to(lerobot_actions.dtype)
print("Action Comparison:")
diff = lerobot_actions - original_actions
abs_diff = torch.abs(diff)
for batch_idx in range(lerobot_actions.shape[0]):
print(f"\n{'=' * 60}")
print(f"Batch {batch_idx}")
print(f"{'=' * 60}")
print(f"{'Idx':<5} {'LeRobot':<14} {'Original':<14} {'Difference':<14}")
print("-" * 60)
for action_idx in range(lerobot_actions.shape[1]):
lr_val = lerobot_actions[batch_idx, action_idx].item()
orig_val = original_actions[batch_idx, action_idx].item()
diff_val = abs(lr_val - orig_val)
sign = "+" if (lr_val - orig_val) > 0 else "-"
print(f"{action_idx:<5} {lr_val:>13.6f} {orig_val:>13.6f} {sign}{diff_val:>12.6f}")
max_diff = abs_diff.max().item()
tolerance = 0.001
assert torch.allclose(lerobot_actions, original_actions, atol=tolerance), (
f"Actions differ by more than tolerance ({tolerance}): max diff = {max_diff:.6f}"
diff = torch.abs(lerobot_action - original_action)
max_diff = diff.max().item()
print(
f"\n[{embodiment_tag}] shapes lerobot={tuple(lerobot_action.shape)} "
f"original={tuple(original_action.shape)} "
f"max|diff|={max_diff:.6e} mean|diff|={diff.mean().item():.6e}"
)
print(f"\nSuccess: Actions match within tolerance ({tolerance})!")
del lerobot_policy, lerobot_preprocessor, lerobot_postprocessor
del original_policy, modality_config, modality_transform
del batch, batch_lerobot, observation
cleanup_memory()
def test_groot_forward_pass_comparison():
"""Test forward pass comparison between LeRobot and Original Groot implementations."""
print("Test: Forward Pass Comparison (Training Mode)")
set_seed_all(42)
lerobot_policy, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_groot(
from_pretrained=True
assert torch.allclose(lerobot_action, original_action, atol=ATOL, rtol=RTOL), (
f"GR00T N1.7 raw action_pred differs for embodiment '{embodiment_tag}' beyond "
f"atol={ATOL}, rtol={RTOL}: max|diff|={max_diff:.6e}"
)
original_policy, modality_config, modality_transform = instantiate_original_groot(from_pretrained=True)
batch = create_dummy_data()
lerobot_policy.eval()
original_policy.model.eval()
print("\n[LeRobot] Running forward pass...")
batch_lerobot = deepcopy(batch)
batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
set_seed_all(42)
with torch.no_grad():
lerobot_loss, lerobot_metrics = lerobot_policy.forward(batch_lerobot_processed)
print(f" Loss: {lerobot_loss.item():.6f}")
print("\n[Original] Running forward pass...")
observation = convert_lerobot_to_original_format(batch, modality_config)
transformed_obs = modality_transform(observation)
if "action" not in transformed_obs:
action_for_forward = batch_lerobot_processed["action"]
action_mask_for_forward = batch_lerobot_processed["action_mask"]
# Match action horizon if needed
if action_for_forward.shape[1] != original_policy.model.action_horizon:
if action_for_forward.shape[1] < original_policy.model.action_horizon:
pad_size = original_policy.model.action_horizon - action_for_forward.shape[1]
last_action = action_for_forward[:, -1:, :]
padding = last_action.repeat(1, pad_size, 1)
action_for_forward = torch.cat([action_for_forward, padding], dim=1)
mask_padding = torch.zeros(
action_mask_for_forward.shape[0],
pad_size,
action_mask_for_forward.shape[2],
dtype=action_mask_for_forward.dtype,
device=action_mask_for_forward.device,
)
action_mask_for_forward = torch.cat([action_mask_for_forward, mask_padding], dim=1)
else:
action_for_forward = action_for_forward[:, : original_policy.model.action_horizon, :]
action_mask_for_forward = action_mask_for_forward[
:, : original_policy.model.action_horizon, :
]
transformed_obs["action"] = action_for_forward
transformed_obs["action_mask"] = action_mask_for_forward
set_seed_all(42)
with torch.no_grad():
original_outputs = original_policy.model.forward(transformed_obs)
original_loss = original_outputs["loss"]
print(f" Loss: {original_loss.item():.6f}")
loss_diff = abs(lerobot_loss.item() - original_loss.item())
loss_rel_diff = loss_diff / (abs(original_loss.item()) + 1e-8) * 100
print("\nLoss Values:")
print(f" LeRobot: {lerobot_loss.item():.6f}")
print(f" Original: {original_loss.item():.6f}")
print(f" Absolute difference: {loss_diff:.6f}")
print(f" Relative difference: {loss_rel_diff:.2f}%")
del lerobot_policy, lerobot_preprocessor, lerobot_postprocessor
del original_policy, modality_config, modality_transform
del batch, batch_lerobot, observation, transformed_obs
cleanup_memory()
tests/policies/groot/utils/dump_original_n1_7.py
+212
View File
@@ -0,0 +1,212 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under 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/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)
# 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()