feat(train) Allow policies to declare their input image format

[acwrenn-nv]

Title

Allow policies to declare their input image format

Summary / Motivation

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.

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.

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	36.6	—	43.5%	—
OSS Gr00t	81.9	+123.7%	84.3%	+40.8 pp
Post-MR	77.6	+112.1% (−5.2% vs OSS)	77.9%	+34.5 pp (−6.3 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)
• All tests pass locally (pytest)
  tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py::test_pi0_fast_action_generation
  tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py::test_pi0_fast_action_token_sampling
  ^ both tests are failing on this branch, but also failing on main.
• Documentation updated
• CI is green .
• 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.