[Bugfix][Model] GR00T-N1.7: honor processor_config image pipeline (use_albumentations / crop_fraction / letter_box_transform) to match Isaac-GR00T

[liangmenghuang]

Purpose

Fixes #7081.

What broke

Gr00tN1d7Processor ignores the image-pipeline keys of the checkpoint's processor_config.json: use_albumentations and letter_box_transform are swallowed by **kwargs, and the eval transform is always LetterBox -> Resize(256, bilinear) -> CenterCrop(230) -> Resize(256). Both public N1.7 checkpoints (nvidia/GR00T-N1.7-3B, nvidia/GR00T-N1.7-DROID) ship use_albumentations: true, crop_fraction: 0.95, shortest_image_edge: 256, image_crop_size: [230, 230], letter_box_transform: false, which in Isaac-GR00T means SmallestMaxSize(256, INTER_AREA) -> FractionalCenterCrop(0.95) -> SmallestMaxSize(256, INTER_AREA): a 243 px crop (not 230; 230/256 = 0.898 vs 0.95), area resampling (not bilinear), no letterbox. The model gets out-of-distribution pixels.

Measured on GB200 with the DROID checkpoint against Isaac-GR00T 51d4c89, on one non-zero frame per view and with the reference's own initial denoising noise injected so the ODE starts from the same point: pixel_values max abs diff 1.945, backbone_features 16.78, decoded actions 0.107-0.489, while every other model input (input_ids, attention_mask, image_grid_thw, state, embodiment_id) is already bitwise equal — which places the divergence in the image pipeline. Full numbers, including the arm of the same experiment run on this branch, are under Test Result.

Repro (main, no Isaac-GR00T needed; cv2 ships with vLLM)

import cv2, numpy as np, torch
from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import _build_eval_image_transform

frame = (np.indices((256, 256)).sum(0) % 256).astype(np.uint8)[..., None].repeat(3, -1)  # any HWC uint8 frame
crop = frame[6:249, 6:249]                                     # Isaac: int(256 * 0.95) = 243, centred
isaac = torch.from_numpy(cv2.resize(crop, (256, 256), interpolation=cv2.INTER_AREA)).permute(2, 0, 1)

# What the checkpoint's processor_config actually asks for -- only reachable on this branch:
fixed = _build_eval_image_transform(
    [256, 256], [230, 230], use_albumentations=True, shortest_image_edge=256, crop_fraction=0.95
)(frame)
print((fixed.int() - isaac.int()).abs().max())   # this PR: tensor(0)

# The recipe main always runs (and this branch still runs when the flag is off):
ours = _build_eval_image_transform([256, 256], [230, 230])(frame)
print((ours.int() - isaac.int()).abs().max())    # tensor(230) -- out of 255

_build_eval_image_transform keeps its positional signature, so the second call behaves the same on main and on this branch; the flags are keyword-only and come from processor_config.json.

Root cause

The port re-implemented Isaac-GR00T's build_image_transformations (torchvision recipe) with the letterbox hard-wired on, and never ported build_image_transformations_albumentations, the recipe selected by use_albumentations=True. The processor __init__ dropped both selector keys as "deprecated training-only keys".

Fix

• Gr00tN1d7Processor.__init__ takes use_albumentations and letter_box_transform explicitly (defaults False, same as Isaac-GR00T) instead of dropping them; save_pretrained persists them so save/load round-trips.
• _build_eval_image_transform selects the recipe like Isaac's processor does: use_albumentations=True -> new AlbumentationsEvalImageTransform, a cv2/numpy transcription of Isaac-GR00T's eval branch and albumentations 1.4.18 (SmallestMaxSize rounding, FractionalCenterCrop truncation, LetterBoxPad), including the crop_fraction/shortest_image_edge None fallbacks. No albumentations dependency; cv2 is already there via vLLM's opencv-python-headless, imported lazily inside the helpers like cosmos3/transfer.py.
• The torchvision recipe now gates LetterBoxTransform on letter_box_transform, as Isaac-GR00T does (no-op for square frames).
• The per-frame contract at _get_vlm_inputs (HWC uint8 in, (C, H, W) uint8 tensor out) is unchanged, so the call site is untouched.

Behavior change

• Checkpoints with use_albumentations: true (both public N1.7 checkpoints) now get the pixels they were trained on; outputs change (correctness fix).
• Checkpoints with use_albumentations: false or without the key keep the torchvision recipe; the only difference is that letterboxing now follows letter_box_transform (absent -> False, Isaac's default), which is a no-op for square frames.
• state_dropout_prob / extra_augmentation_config remain training-only and are still ignored.

Not touched: #6568 / #6450 (pipeline_gr00t.py). #5028 changes the same file, but its AutoProcessor.register(Gr00tN1d7Config, ...) fix is already on main via #5042.

Test Plan

New CPU test tests/diffusion/models/gr00t/test_processing_eval_image_transform.py (core_model + cpu, runs in "Simple · Diffusion Test"):

• recipe selection from the flags, None fallbacks, geometry/rounding rules, contract errors, letterbox gating on both recipes, _get_vlm_inputs end-to-end with the new transform, save_pretrained -> from_pretrained round-trip, legacy configs without the flags;
• bit-exact goldens: sha256 of the (C, H, W) uint8 output of Isaac-GR00T's own eval transform (build_image_transformations_albumentations(...)[1] through apply_with_replay, Isaac-GR00T 51d4c89 / albumentations 1.4.18 / opencv 5.0.0.93) on 4 deterministic integer-generated frames x letterbox on/off. The non-square frame is 512x768 so the INTER_AREA downscales are integer factors (exact integer sums in OpenCV): the goldens do not depend on the CPU architecture. Generating them also cross-checked Isaac-GR00T, this branch's code and an independent transcription for bitwise agreement on every case, and the torchvision branch against Isaac's torchvision eval transform. Regenerating a digest needs Isaac-GR00T installed, whose transformers pin conflicts with this repo's, so it cannot run in CI; the test module's docstring states the procedure.

pytest -sv tests/diffusion/models/gr00t -m 'core_model and cpu'
pre-commit run --all-files

Processor-level parity (Isaac-GR00T processor vs patched processor, transformers 4.57.6 environment, DROID checkpoint, 2 seeds x batch 1/2): pixel_values, input_ids, image_grid_thw, attention_mask, state, backbone_features, absolute_action all bitwise equal, also through the real serving path.

vLLM Version: 0.28.0

vLLM-Omni Commit: 10701e1 (upstream/main, branched here)

Test Result

$ pytest -sv tests/diffusion/models/gr00t -m 'core_model and cpu'
21 passed, 22 warnings in 3.68s

(17 new cases plus the 4 pre-existing pipeline tests. On main the new module fails to import
AlbumentationsEvalImageTransform; against main's recipe the pinned Isaac-GR00T goldens differ by
up to 255 gray levels on the checkerboard frame, 247 on the noise frame, and the non-square frame
even differs in shape, 256x256 instead of 256x383.)

End to end on GB200

Isaac-GR00T pins transformers==4.57.3 / torch==2.9.0 while vLLM-Omni main needs transformers 5.x, so the two cannot share an interpreter. They were run as two processes on one GB200 node against nvidia/GR00T-N1.7-DROID, exchanging tensors through disk: the same observation file (one non-zero 256x256 frame per view — a zero frame is a fixed point of every resize/crop recipe and would show nothing), and the reference's own initial denoising noise captured and injected into the vLLM-Omni run. The serving side drives Gr00tPolicy.get_action and imports nothing but vllm_omni. Both arms use the same reference dump and differ only by git checkout:

tensor	main	this branch
pixel_values	1.945	0 (bitwise)
input_ids, attention_mask, image_grid_thw, state, embodiment_id, state_features	0	0
backbone_features	16.78	1.375
action_pred (normalized)	0.979	0.031
decoded eef_9d / gripper_position / joint_position	0.148 / 0.489 / 0.107	0.004 / 0.006 / 0.009

The processor reports its recipe as Compose on main and AlbumentationsEvalImageTransform on this branch, so the switch is observable rather than inferred.

pixel_values — the tensor this PR is about — goes to bitwise equality. The downstream rows do not, and should not be read as if they could: the two sides necessarily run different torch/transformers builds, and that difference alone is worth roughly the residual in the right-hand column. Both arms carry that same floor, so the honest claim is the contrast between the columns, not the absolute value of either. A same-process end-to-end bitwise check is not reachable while the dependency pins conflict.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/diffusion_model_integration.md.

Module owners: @fhfuih @Bounty-hunter @wtomin @RuixiangMa

@liangmenghuang, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit eaeaae804252 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[liangmenghuang]

Self-review.

What changed: Gr00tN1d7Processor now reads use_albumentations / letter_box_transform from processor_config.json instead of dropping them into **kwargs, and _build_eval_image_transform selects the recipe the way Isaac-GR00T's processor does. Two files: the processor and a new CPU test module.

What I checked:

• Behavior gating — checkpoints without the keys keep the previous torchvision recipe (defaults False, matching Isaac-GR00T); only the letterbox now follows letter_box_transform, a no-op for square frames. Both public N1.7 checkpoints ship use_albumentations: true, so they change behavior — that is the correctness fix.
• Call-site contract — _get_vlm_inputs still gets HWC uint8 in / (C, H, W) uint8 tensor out, so the single call site is untouched.
• Round-trip — both flags are persisted by save_pretrained; there is a save_pretrained -> from_pretrained test.
• No new dependency — cv2 comes from vLLM's opencv-python-headless and is imported lazily inside the helpers.
• Reachability of the None fallbacks — image_target_size / image_crop_size default to None in the processor __init__, so the guards in _build_eval_image_transform are live, not defensive dead code.
• Tests — pytest tests/diffusion/models/gr00t -m 'core_model and cpu' -> 21 passed (17 new + 4 pre-existing). Goldens are sha256 of Isaac-GR00T's own eval transform output; the non-square frame is 512x768 so every INTER_AREA downscale is an integer factor and the goldens are architecture-independent.
• Pre-commit — full hook set passes locally (SPDX, forbidden imports, no new torch.cuda, mypy-3.10, test CI marks, ruff, typos).
• End to end — two-process comparison against Isaac-GR00T on GB200, numbers in the description. pixel_values reaches bitwise equality; the downstream residual is the different torch/transformers builds the two sides must run, and it is present in both arms.

What I did not check: whether this changes any rollout benchmark outcome. The gripper_position observation in #7081 is flagged there as unverified.

[princepride]

@tzhouam PTAL

[timzsu]

This PR adds support for two transform options that my initial porting missed. The change is easy to read through and reproduction plus end-to-end precision numbers are also provided. LGTM overall.

I have left a minor comment on the docstring which seems to reference a non-existent file. PTAL when you have time.

[princepride]

Lgtm