opentau.utils.vision_utils

Vision-tower helpers reused across policies (pi06, pi07, …).

Residents:

  • bilinear_resample_pos_embed() adapts SigLIP / ViT-style learned position-embedding tables to a different patch count when the policy runs the vision tower at a non-published resolution (offline weight surgery).

  • patch_grid_hw() / pad_to_patch_multiple() support running the SigLIP tower at native (non-square, non-patch-multiple) input resolutions at forward time: the stride-patch_size conv patch embedding silently floor-crops any sub-patch remainder (e.g. 180×320 with patch 14 drops 12 pixel rows and 12 columns), so callers pad up to the next patch multiple and enable interpolate_pos_encoding instead.

Why this is its own module rather than inline in a policy: π0.6 uses Gemma 3-4B at 448×448 (1024 patches) but google/gemma-3-4b-pt ships at 896×896 (4096 patches), so any script that bootstraps an “untrained” π0.6 checkpoint from the public VLM weights must resample. π0.7 uses the same backbone family and will inherit the same need — sharing the helper here keeps the recipe identical across them and avoids drift between re-implementations.

Functions

bilinear_resample_pos_embed(old_pos, ...)

Bilinearly resample a learned (N_patches, dim) position-embedding table to a different patch count, preserving dtype.

pad_to_patch_multiple(img, patch_size[, ...])

Pad the trailing two (H, W) dims up to the next multiple of patch_size.

patch_grid_hw(height, width, patch_size)

Patch-grid shape (grid_h, grid_w) that fully covers an image.
opentau.utils.vision_utils.bilinear_resample_pos_embed(old_pos: Tensor, target_num_patches: int) → Tensor[source]

Bilinearly resample a learned (N_patches, dim) position-embedding table to a different patch count, preserving dtype.

This is the standard recipe for adapting SigLIP / ViT position embeddings when running a published VLM at a different image resolution. Both the source and target patch counts must be perfect squares (square grids). Computation is performed in fp32 (F.interpolate rejects bf16 on CPU) and cast back to the input dtype before return.

The transform is fully deterministic: two calls with bit-identical inputs return bit-identical outputs, with no RNG consumption. This matters for bootstrapping checkpoints whose downstream tests compare weights byte- for-byte against an independently-computed reference.

Parameters:

  • old_pos – Source position-embedding table, shape (N_old, dim) with N_old a perfect square.

  • target_num_patches – Desired output N_new, must be a perfect square.

Returns:

Resampled tensor of shape (target_num_patches, dim) and same dtype as old_pos. Returns old_pos unchanged when N_old == target_num_patches (no copy).

Raises:

AssertionError – if either patch count is not a perfect square.

opentau.utils.vision_utils.pad_to_patch_multiple(img: Tensor, patch_size: int, pad_value: float = -1.0) → Tensor[source]

Pad the trailing two (H, W) dims up to the next multiple of patch_size.

Padding goes on the bottom and right so the image origin stays pixel-aligned: every patch that contains real content keeps exactly the pixels it would have at a divisible resolution, and only the last patch row/column mixes in padding. The default pad_value=-1.0 is black in the [-1, 1] range SigLIP consumes; callers padding [0, 1]-range pixels should pass 0.0.

Parameters:

  • img – Image tensor of shape (..., H, W).

  • patch_size – ViT patch size to pad up to a multiple of.

  • pad_value – Fill value for the padded band.

Returns:

The input unchanged (no copy) when H and W are already multiples of patch_size, else a padded copy of shape (..., ceil(H/p)*p, ceil(W/p)*p).

opentau.utils.vision_utils.patch_grid_hw(height: int, width: int, patch_size: int) → tuple[int, int][source]

Patch-grid shape (grid_h, grid_w) that fully covers an image.

Uses ceiling division, so a resolution that does not divide patch_size still gets a grid covering every pixel — the caller is expected to pad the image up to grid * patch_size with pad_to_patch_multiple() before the conv patch embedding (which would otherwise floor-crop the remainder).

Parameters:

  • height – Image height in pixels.

  • width – Image width in pixels.

  • patch_size – ViT patch size (14 for SigLIP so400m).

Returns:

(grid_h, grid_w) patch counts per axis.