Allow policies to declare their input image format

[acwrenn-nv]

Title

Allow policies to declare their input image format

Summary / Motivation

GR00T N1.7 dataset workers already produce compact uint8 images, but the training loop converted them to float32 [0, 1] before the GR00T preprocessor converted them back to uint8 and staged them through a CPU NumPy/HWC buffer. This change adds an explicit class-level raw-image input contract to policies so GR00T can preserve worker-produced uint8 tensors through packing and transfer. Existing policies inherit the historical float32 behavior without adding policy-specific state to serialized processor configs.

The GR00T video intermediate is now an ordered tuple of per-camera (B, T, C, H, W) tensors. The cv2/Albumentations path still converts individual frames to HWC at its point of use, preserving the reference transform and crop behavior while removing the full-batch staging copy.

Related issues

• Fixes / Closes: None.
• Related: None.

What changed

• Added a class-level ImageInputFormat contract to PreTrainedPolicy, with float32_0_1 as the inherited default and no processor-pipeline serialization.
• Updated training and validation image preparation to convert worker-produced uint8 images only when the selected policy expects float32 input; the value is captured before PEFT or Accelerate wraps the model.
• Declared GrootPolicy.input_image_format = UINT8_0_255 and removed image-format state from new and loaded GR00T preprocessors.
• Kept GR00T camera batches as ordered channels-first tensors through packing and VLM preparation, avoiding the full-batch NumPy/HWC staging buffer and preserving the original storage where possible.
• Retained direct/inference support for float observations, camera/time ordering, temporal padding, Albumentations geometry, and train-time crop replay across camera views.
• Added regression coverage for policy-level image conversion, the absence of processor serialization, GR00T's class override, zero-copy single- and multi-camera packing, VLM ordering/parity, and random-crop behavior.
• No user-facing breaking change or migration is required.

How was this tested (or how to run locally)

• Targeted touched-file suite:

  uv run pytest -q \
    tests/policies/groot/test_groot_n1_7.py \
    tests/policies/groot/test_groot_n1_7_oss_parity.py \
    tests/policies/groot/test_groot_train_random_crop.py \
    tests/processor/test_pipeline.py \
    tests/training/test_visual_validation.py

  CUDA run on the current PR head: 169 passed, 7 skipped.

• Independent real-hardware A/B on an RTX PRO 6000 with a 52k-frame, two-camera SO-101 dataset, the N1.7 3B checkpoint, use_albumentations=true, and batch size 32:

  • All 15 preprocessor output tensors were byte-exact in train and eval modes.
  • Forward loss was bit-identical; gradient deltas stayed within the control run's inherent backward nondeterminism.
  • Image preparation + preprocessing improved from 260.6 ms to 158.8 ms.
  • End-to-end step time improved from 664 ms to 567 ms; throughput improved from 48.2 to 56.5 examples/s (+17.2%).

• Full benchmark methodology and parity results: A/B verification and gradient follow-up.

W&B training/performance results

These performance results apply to the MR as a whole. The class-level ImageInputFormat declaration changes where the input contract lives without changing image values, preprocessing, or the training execution path.

Runs are in nvidia/lerobot-gr00t-17n-train. After a 10% warmup, examples/s is median train/samples_per_s and GPU utilization is mean system.gpu.*.gpu across GPUs with model memory allocated.

Four-suite performance comparison (mean across LIBERO-10, Spatial, Object, and Goal; global batch 640, 2×320)

Variant	Examples/s	Δ examples/s vs Pre-MR	GPU utilization	Δ GPU utilization vs Pre-MR
Pre-MR	142.2	—	27.1%	—
OSS	179.8	+26.5%	33.8%	+6.6 pp
Post-MR	182.0	+28.0% (+1.2% vs OSS)	33.9%	+6.7 pp (+0.1 pp vs OSS)

• The isolated real-hardware A/B improved end-to-end throughput from 48.2 to 56.5 examples/s (+17.2%) and reduced image preparation + preprocessing from 260.6 ms to 158.8 ms.
• Across the completed four-suite validation, median data time fell from 2.584 to 1.888 s/step (−26.9%).
• The 5,000-step single-GPU uint8-direct validation sustained 78.5 examples/s, 77.6% mean GPU, and 0.063 s median data time.

Checklist (required before merge)

• Linting/formatting run (pre-commit run -a) — targeted Ruff format/check passes on all changed files; full pre-commit was not run.
• All tests pass locally (pytest) — the targeted touched-file suite passes; the full suite was not independently verified.
• Documentation updated — not required; there is no user-facing CLI or configuration change.
• CI is green — no checks are currently reported for the current head.
• Community Review: I have reviewed another contributor's open PR and linked it here: not yet linked.

Reviewer notes

• Please focus on ImageInputFormat ownership on PreTrainedPolicy, the GR00T class override, capture before model wrappers, and the ordered per-camera tensor contract between GrootN17PackInputsStep and GrootN17VLMEncodeStep.
• The Albumentations/cv2 path intentionally remains CPU-bound to preserve bit-exact reference transforms. Prefetching/overlap and fused AdamW are separate follow-up opportunities and are out of scope here.
• Anyone in the community is free to review the PR.

[johnnynunez]

Ran an independent A/B verification of this PR on real hardware (RTX PRO 6000, torch 2.11 cu128) with a real 52k-frame SO-101 dataset (2 cams, 480×640), the N1.7 3B checkpoint, use_albumentations=true, batch 32 — i.e. the actual training config, not a synthetic micro-bench.

1. Numeric parity: byte-exact ✅

Same fixed batch, same seeds, through both branches (main @ 0530dd9 vs this PR). All 15 preprocessor output tensors (pixel_values, input_ids, attention_mask, action, state, masks, …) are bit-identical in both eval and train mode (max|Δ| = 0 on floats, 0 mismatches on ints).

Also worth noting: the old uint8 → float32/255 → ×255 → uint8 roundtrip is lossless for all 256 values (verified on CPU, CUDA fp32, and even bf16). So the concern about "floating point math on the GPU" doesn't apply to correctness here — the old path's problem was purely throughput, never precision. If train curves differ between runs, look at RNG/data ordering, not image dtype.

2. End-to-end training step (30 steps, medians, n=30)

Phase	main	PR46	Δ
image prep + preprocessor	260.6 ms	158.8 ms	−102 ms
forward	250.6 ms	255.0 ms	noise
backward	65.3 ms	65.4 ms	—
optimizer	87.7 ms	87.6 ms	—
step total	664 ms	567 ms	−97 ms
throughput	48.2 samples/s	56.5 samples/s	+17.2%
GPU util (NVML, 20 ms sampling)	59.9% mean	66.6% mean	+6.7 pt

Effect is ~7σ vs step-to-step noise (σ ≈ 11–16 ms). The PackInputsStep improvement alone is 17.0 → 0.8 ms (the zero-copy tuple contract works — data_ptr preserved, no numpy staging).

3. Where the remaining GPU headroom is (nsight systems, NVTX GPU-projected)

During the preprocessor range the GPU is 1.2% busy — it's a pure CPU-serial bubble. Inside GrootN17VLMEncodeStep (B=32, 64 frames, albumentations path):

Component	median	share
tensor→numpy HWC staging (this PR's code)	2.9 ms	~2%
cv2 resize/crop/resize (albumentations, bit-exact contract)	~63 ms	~42%
Qwen processor encode (patchify+normalize+tokenize, CPU)	~87 ms	~56%

GPU-bound floor (fwd+bwd+optim) is 408 ms → 78.4 samples/s ideal. PR46 lands at 56.5, so ~+39% is still on the table, all of it in that CPU bubble.

Suggested follow-up (separate PR)

Rather than accelerating the transforms (which risks the bit-exact albumentations contract), overlap them: prefetch one batch ahead through pack+VLM-encode in a background thread (or push it into the dataloader workers) with a double buffer. Then step ≈ max(GPU 408 ms, CPU 159 ms) — the CPU cost leaves the critical path entirely, cv2 math stays untouched, and utilization goes to ~95%. Cheap orthogonal win on top: AdamW(fused=True).

Verdict: LGTM. Clean opt-in contract, backward compatible (legacy configs default to float, old serialized GR00T pipelines get upgraded at load), byte-exact, +17% measured end-to-end.

_{Method note: the gated nvidia/Cosmos-Reason2-2B processor was substituted with Qwen/Qwen3-VL-2B-Instruct (same family) identically on both branches for the tokenizer/image-processor — this cancels out in the A/B; model weights were the real N1.7 3B checkpoint. All touched test files pass on both branches (169 passed, 7 skipped, cuda).}

[johnnynunez]

Follow-up: gradient-level verification (was asked whether the parity claim covers gradient calculation — it does now).

Gradient parity: main vs PR46 on the real 3B

Full forward+backward on a fixed real-data batch, per-tensor SHA256 over all 537 gradient tensors, TF32 disabled, same seeds. Crucially, with a run-vs-run determinism control (same code executed twice) to calibrate what "different" even means here.

Comparison	Loss (fp64 bits)	Global grad norm	Tensors w/ differing hashes
PR46 run1 vs PR46 run2 (control)	bit-identical	0.63022734 vs 0.63024698	500/537
main vs PR46	bit-identical (1.1334302425)	0.63020239 vs 0.63022734	500/537

The control matters: the 3B backward is inherently nondeterministic run-to-run (the same ~500 tensors — DiT action head blocks — change hashes between two identical runs; classic atomics nondeterminism). So raw hashes can't be the criterion. The criterion is whether main↔PR46 differs more than run↔run:

Per-tensor grad-norm relative delta	median	p90	max
run-vs-run (intrinsic noise)	1.28e-4	1.38e-3	4.3e-1
main-vs-PR46	1.53e-4	1.19e-3	2.9e-1

Median ratio = 1.20 → main-vs-PR46 sits inside the same noise envelope as re-running identical code. Combined with the byte-exact preprocessor outputs and bit-identical loss (forward is deterministic; only backward has atomics), the conclusion is: no gradient difference attributable to this PR.

AdamW fused=True (the suggested orthogonal win, now measured)

On the actual 1.62B trainable parameter set (537 tensors), identical synthetic grads on cloned params, 5 steps, fused=False vs fused=True:

• max |Δweight| after 5 updates: 2.98e-7 — fp32 epsilon level, numerically equivalent
• optimizer.step(): 55.0 ms → 31.8 ms = 1.73×
• End-to-end: ~567 → ~544 ms/step ≈ +4% throughput on top of this PR, one-line change

For reference, Isaac-GR00T already defaults to it: gr00t/configs/training/training_config.py:55 → optim: str = "adamw_torch_fused". In LeRobot the gap is in src/lerobot/optim/optimizers.py::AdamWConfig.build(), which calls torch.optim.AdamW(params, **kwargs) without fused. Worth a tiny separate PR.

_{Same setup as the previous comment: RTX PRO 6000, real N1.7 3B weights, vials dataset, gated Cosmos-Reason2-2B tokenizer/processor substituted with Qwen3-VL-2B identically on both branches.}