[BUG]Images.from_list() zero-pads mixed-size images to batch-max — black borders + wrong aspect corrupt source/condition images #169

Summary

Images.from_list() (unirl/types/primitives.py:156-177) builds the batched Images.pixels: [B, C, H, W] tensor. When the per-sample images in a batch have different H/W, it zero-pads each smaller image at the bottom/right up to the batch-max (max_h, max_w):

padded_p = torch.nn.functional.pad(p, (0, pad_w, 0, pad_h), mode="constant", value=0)

Consequences:

• Every non-max image gets black borders baked in, and its effective aspect ratio becomes max_h : max_w (wrong).
• The original per-sample size is lost — nothing downstream can recover it.
• to_list() / to_pils() return the padded per-sample images, so any consumer that unstacks (VLM processor, VAE/ViT encode) inherits the corruption.

This is also inconsistent with the sibling Videos.from_list() (primitives.py:236-246), which deliberately resizes (interpolates), not zero-pads, exactly to avoid black borders on mixed-resolution inputs.

How other frameworks handle this (they don't hit it)

The root cause is forcing heterogeneous images into a single rectangular [B, C, H, W] tensor at the type boundary. Mainstream VLM-RL stacks avoid this entirely:

• verl / EasyR1 / TRL / ms-swift: keep multimodal images as a per-sample list of PIL images (multi_modal_data = {"image": [PIL, PIL, ...]}), never stacked. The Qwen2-VL / 2.5-VL image processor smart-resizes each image independently and flatten+concatenates patches into pixel_values [sum(grid_t*grid_h*grid_w), ...] with image_grid_thw recording each image's grid. No zero-pad, no black borders, no batch-max coupling. (cf. verl PR [feat] Initial support for VLMs, add Qwen2.5VL GRPO example verl-project/verl#386; transformers Qwen2VLImageProcessor.)
• diffusers-style I2V / editing pipelines: resize each reference image to the target generation resolution per-sample; they never pad to a batch-max canvas.

So this looks like a UniRL-specific artifact of the Images rectangular-tensor design, not something inherited from the upstream references.

Impact — which training is affected

Triggers only when a batch mixes different-size source/condition images (req.primitives["image"]). _validate_homogeneous_images (unirl/data/data_source.py:128) only checks presence (all-or-none), not that sizes match, so heterogeneous batches pass straight into the padding branch. _load_condition_images loads raw PILs at native size (no resize), so any dataset with non-uniform images hits it.

A. Unstack → per-sample processing (padding is pure, avoidable damage):

• Qwen-VL VLM understanding — qwen_vl/pipeline.py:76 to_pils() → processor; rollout side sglang/adapters/vlm.py:85.
• Bagel i2t / it2t / it2i — bagel/pipeline.py:197 to_list() → to_pil() → resize_transform → VAE + ViT.

B. VAE-encode → "source image latent" (black borders baked into the latent):

• WAN 2.1 / 2.2 I2V — wan21/image_encode.py:80-87 interpolates the padded canvas to target; CLIP-vision clip_vision_encode.py:69.
• Flux2 Klein (Kontext edit) — flux2_klein/vae.py:148-152.
• HunyuanImage3 it2i / i2t — hunyuan_image3/vae.py:110 (encodes p.pixels directly, no resize).

Not affected: pure T2I / T2V / text-only; LTX2 I2V (raises NotImplementedError before use); V2V (uses Videos, already resize-based).

Concrete exposure today:

• geo3k_mc is currently safe — all 200 images happen to be uniform 224×224, so the padding branch never fires.
• arxivqa_mc is a live heterogeneous case — prepare_arxivqa_mc.py:154 uses pil.thumbnail((max_edge, max_edge)) (downscale-only, aspect-preserving → non-uniform sizes). examples/ar/bagel_grpo_arxivqa_mc_2x8_lora.yaml feeds exactly such a batch.

Notes:

• Rollout and trainside-replay consume the same padded Images, so the corruption is consistent across both → the GRPO importance ratio is not broken; the damage is to training / reward quality (the model learns on black-bordered, aspect-distorted images).
• Secondary risk: different DP shards can pad to different (max_h, max_w); any cross-shard concat of Images (FieldKind.CONCAT) then shape-mismatches on H/W — the exact failure Videos avoided by resizing.

Options

A. Preserve per-sample native size (root fix).
For the unstack-then-process consumers (Qwen-VL, Bagel), the rectangular tensor is never actually needed — they immediately go back to per-sample PILs. Let Images hold a ragged / per-sample list when sizes differ (or have the data source not stack heterogeneous images), so to_list() / to_pils() return native-size images and each model resizes to its own target at encode time. No black borders, no aspect loss. Mirrors how verl / diffusers keep images per-sample.

B. Mirror Videos.from_list() — resize instead of zero-pad (local fix).
Replace the zero-pad branch with F.interpolate to (max_h, max_w). Removes black borders, and for B-type consumers that re-resize to a fixed target the result is acceptable. Downside: still distorts aspect to the batch-max ratio (mixed portrait/landscape still warped). Cheapest, single-function change, and also removes the DP-concat crash.

C. Fail loud / require uniform inputs (guardrail).
Make _validate_homogeneous_images (and/or Images.from_list) reject or warn on non-uniform sizes and require datasets to pre-resize condition images (as geo3k_mc effectively does). Doesn't fix the capability, but converts silent corruption into a visible signal.

Recommendation

A is the correct long-term fix — it matches how reference VLM-RL stacks handle multi-resolution images and loses no information. B is a good self-contained stopgap that at least kills the black borders (and the DP-concat crash) if the type change is too invasive now. At minimum, C should land so heterogeneous batches never silently corrupt training.