opentau.policies.pi07.video_encoder

SigLIP video encoder with space-time separable attention (MEM paper).

Implements the low-level memory video encoder from Torne, Pertsch, Walke et al. “MEM: Multi-Scale Embodied Memory for Vision Language Action Models” (Section III-C + Appendix C): a standard SigLIP ViT extended with space-time separable attention at every spacetime_layer_stride-th layer and a fixed sinusoidal temporal position encoding whose current-frame row is zero. Past-timestep tokens are dropped after the encoder so the output shape matches a single-image VLA.

This module is the single, canonical implementation of the SigLIP video encoder; all callers — pi05_mem, pi07/low_level (Gemma 3 backbone), and pi07_paligemma/low_level (legacy PaliGemma backbone) — import from here.

Key properties:

  • Introduces no new learnable parameters on top of the pretrained SigLIP weights (temporal attention re-uses each layer’s own Q/K/V/O projections). Any pi05/pi05_continuous_state checkpoint can be loaded directly — the space-time layers just wrap the existing SiglipEncoderLayer weights.

  • Single-frame invariance: with T=1 the output is byte-identical to PaliGemmaModel.get_image_features (see the single-frame invariance tests).

  • Convention: the current frame lives at the last time index (t = T-1). This matches src/opentau/datasets/factory.py:136 (delta_timestamps) and PI05MemPolicy._build_history_batch. obs_history_is_pad[:, -1] is always False by construction.

  • Variable number of frames per forward. The encoder is constructed with a max_num_frames cap (used to size the cached temporal PE buffer); each forward accepts any T in [1, max_num_frames]. Slicing the precomputed PE as pe[max_num_frames - T:] preserves the “row T-1 is zero” invariant — the PE values are functions of the relative offset from the current frame, so the last T rows of a PE built for M frames are byte-identical to a PE built for T frames directly.

The encoder does NOT own its own copy of the SigLIP weights. The caller constructs a PaliGemmaWithExpertModel (pi05_mem, pi07_paligemma) or Gemma3WithExpertModel (pi07/low_level) — which already owns vision_tower and multi_modal_projector — and passes them in by reference. This avoids duplicating ~400M parameters in memory.

Functions

suppress_spacetime_temporal(module)

Context manager that flips _temporal_active=False on every SpaceTimeEncoderLayerWrapper in module's subtree, and restores the previous value on exit.

Classes

SpaceTimeEncoderLayerWrapper(base_layer, ...)

Replaces a SiglipEncoderLayer in-place; adds factorized space-time attention via a single composed attention sublayer (Reading B of MEM Eq.

SpaceTimeSiglipVideoEncoder(vision_tower, ...)

SigLIP-based video encoder with space-time separable attention.
class opentau.policies.pi07.video_encoder.SpaceTimeEncoderLayerWrapper(base_layer: SiglipEncoderLayer, max_num_frames: int, num_tokens_per_frame: int)[source]

Bases: Module

Replaces a SiglipEncoderLayer in-place; adds factorized space-time attention via a single composed attention sublayer (Reading B of MEM Eq. 3).

The wrapper adopts the original layer’s submodules by reference — self_attn, layer_norm1, layer_norm2, mlp — so its state_dict keys are identical to a vanilla SiglipEncoderLayer. That means any pi05 / pi05_continuous_state checkpoint can load directly into the wrapped layer without any key remapping. The only new state is a non-persistent _temporal_pe buffer (excluded from state_dict).

The forward computes (single QKV projection per block; two SDPAs share those projections; one residual; one out_proj):

h_pe = h + e(t) # broadcast over (B, N) z = LN1(h_pe) # ONE LN Q,K,V = W_Q·z, W_K·z, W_V·z # ONE projection V’ = SDPA(Q, K, V; causal mask over T) # temporal pass per patch out = SDPA(Q, K, V’; no mask, over N) # spatial pass per timestep h = h + W_O(out) # ONE residual on h (not h_pe) h = h + MLP( LN2(h) ) # standard MLP block

Q and K are reused across both passes (same tensor, just permuted into each layout); only V is replaced by V' for the spatial pass. W_O fires once at the end. This matches the paper’s “no new learnable parameters” claim at the projection level: W_Q, W_K, W_V, W_O are each applied exactly once per block, not twice.

At T=1 the temporal SDPA collapses to the identity (a single key, so the attention weight is 1 and V passes through unchanged) and e(t=0)=0 makes h_pe = h. The block is therefore mathematically identical to a vanilla SiglipEncoderLayer forward at T=1 — no short-circuit is required for correctness. The wrapper still routes T=1 inputs through _spatial_block_forward for compute savings and bit-exact match in low-precision dtypes (where the no-op temporal SDPA can drift by ~1 ULP).

Variable T per forward: the cached PE is built once for max_num_frames rows and sliced as pe[max_num_frames - num_frames:] each forward. Because PE values depend only on the relative offset from the current frame (not on the absolute index), this slice is byte-identical to a PE freshly built for the actual num_frames.

__init__(base_layer: SiglipEncoderLayer, max_num_frames: int, num_tokens_per_frame: int)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states: Tensor, attention_mask: Tensor | None = None, output_attentions: bool = False, temporal_attn_mask: Tensor | None = None, num_frames: int | None = None) → Tensor[source]

hidden_states: (B*T, N, D) -> (B*T, N, D).

Signature extends SiglipEncoderLayer.forward with two extra kwargs: temporal_attn_mask and num_frames. Vanilla SiglipEncoderLayer instances never receive them: SpaceTimeSiglipVideoEncoder.forward dispatches them only to SpaceTimeEncoderLayerWrapper instances via an isinstance check, bypassing SiglipEncoder.forward entirely.

Parameters:

  • temporal_attn_mask – Optional (B*N, 1, T, T) additive float mask for temporal attention. 0.0 = attend, -inf = block. When None, the standard causal mask is used.

  • num_frames – Number of frames T for this forward. Required when _temporal_active is True and the input is multi-frame (the wrapper cannot infer T from the (B*T, N, D) shape alone). May be None when _temporal_active is False (suppress path) or when the caller knows the input is a single frame.

gradient_checkpointing: bool = False
class opentau.policies.pi07.video_encoder.SpaceTimeSiglipVideoEncoder(vision_tower: SiglipVisionModel, multi_modal_projector: Module, max_num_frames: int, spacetime_layer_stride: int = 4, gradient_checkpointing: bool = False, expected_image_size: tuple[int, int] | None = None)[source]

Bases: Module

SigLIP-based video encoder with space-time separable attention.

Takes video tensors of shape (B, T, 3, H, W) in the [0, 1] range and produces (B, num_video_tokens, vlm_hidden_size). Rescales pixels to [-1, 1] internally (SigLIP’s expected range).

T may vary per forward in [1, max_num_frames]; the encoder is constructed with a max_num_frames cap that sizes the cached temporal PE buffer.

Past-timestep tokens are dropped after the encoder; only the current frame’s num_video_tokens tokens are returned, so the output shape is identical to a single-frame VLA’s vision-token prefix.

The multi_modal_projector is applied to match the output space of PaliGemmaModel.get_image_features. We intentionally do not apply the / sqrt(text_hidden_size) scaling, matching opentau.utils.transformers_patch.patched_paligemma_model_get_image_features which removes it from stock HuggingFace.

The caller owns vision_tower and multi_modal_projector. This module holds them by reference (via a list, so nn.Module does not re-register their parameters under this module’s path) and mutates the vision_tower’s encoder in place to wrap every spacetime_layer_stride-th layer. Callers include PI07LowLevelFlowMatching (Gemma 3 backbone), PI05MemFlowMatching (PaliGemma backbone), and the legacy PI07LowLevelFlowMatching under pi07_paligemma (PaliGemma backbone).

__init__(vision_tower: SiglipVisionModel, multi_modal_projector: Module, max_num_frames: int, spacetime_layer_stride: int = 4, gradient_checkpointing: bool = False, expected_image_size: tuple[int, int] | None = None)[source]

See the class docstring; only the non-obvious arg is documented here.

Parameters:

expected_image_size – (H, W) of the frames each forward will receive. None (default) means the SigLIP config’s square image_size — the historical behaviour. A non-default value enables native-resolution encoding: frames are padded up to the next patch_size multiple (so the conv patch embedding never floor-crops pixels) and the pretrained position embeddings are bicubically interpolated to the resulting patch grid. num_video_tokens then reflects that grid, so every consumer of the token count (skip-path zero fills, temporal-mask expansion, per-layer shape checks) stays consistent by construction. Determinism caveat: with freeze_vision_encoder=False the interpolation backprops into the trainable position-embedding table every step, and CUDA’s backward for bicubic F.interpolate is non-deterministic (atomicAdd) — GPU native-resolution runs with an unfrozen vision tower are therefore not bit-reproducible. The default (frozen tower) is unaffected, as is the config-resolution path (no interpolation).

forward(video: Tensor, obs_history_is_pad: Tensor | None = None) → Tensor[source]

Encode a video clip and return the current-frame tokens.

Parameters:

  • video – (B, T, C, H, W) pixel values in [0, 1], with 1 <= T <= max_num_frames, C == 3, and spatial size matching expected_image_size (the SigLIP config’s square image_size unless overridden at construction).

  • obs_history_is_pad – Optional (B, T) bool mask where True marks padded history frames. Padded frames are blocked in the temporal attention so the current frame cannot read contaminated hidden states from them. When None and T > 1, falls back to “only the current frame is real” — a defensive default for callers that omit the mask (the built-in select_action -> _build_history_batch path does emit it).

Returns:

(B, num_video_tokens, vlm_hidden_size) current-frame tokens, ready to concatenate into the VLA prefix.

property multi_modal_projector: Module
property vision_tower: SiglipVisionModel
opentau.policies.pi07.video_encoder.suppress_spacetime_temporal(module: Module) → Iterator[None][source]

Context manager that flips _temporal_active=False on every SpaceTimeEncoderLayerWrapper in module’s subtree, and restores the previous value on exit.

Used by Gemma3WithExpertModel.embed_image so that single-image forwards through a vision_tower that has been wrapped with space-time attention skip the temporal sublayer (which has no time axis to attend over for non-video inputs). When module contains no wrappers (e.g. no video encoder has been constructed yet), this is a no-op.