opentau.policies.pi05_mem.rldx_video_encoder

RLDX SigLIP video encoder: plain SigLIP + STSS motion module.

RLDXVideoEncoder is a standalone video encoder that runs a stock SigLIP ViT on each frame and injects the RLDX-1 space-time self-similarity (STSS) motion module (https://github.com/RLWRLD/RLDX-1) as a residual update at a single encoder layer. The motion module is the only cross-frame mechanism here: unlike SpaceTimeSiglipVideoEncoder, this encoder has no space-time / temporal attention layers — the SigLIP tower is left exactly as pretrained, and temporal dynamics are captured solely by the STSS motion module (see opentau.policies.pi05_mem.motion_module).

It deliberately does not subclass or import the space-time encoder, so the two implementations stay fully independent.

Contract (identical I/O to the space-time encoder, so it is a drop-in):
forward(video, obs_history_is_pad=None)

video: (B, T, C, H, W) pixel values in [0, 1], 1 <= T. returns: (B, num_video_tokens, vlm_hidden_size) current-frame tokens.

Notes

  • Adds new learnable parameters (the motion module). A plain pi05 checkpoint loads with strict=False; with motion_zero_init=True (default) the motion residual is zero-gated at init so the policy is byte-identical at step 0 and the motion contribution warms up.

  • The motion module is registered under motion_module.* and stays trainable even when the SigLIP tower is frozen by the caller.

  • Motion runs only when there is a real time axis (T > 1); at T = 1 the encoder is a plain single-frame SigLIP forward.

  • obs_history_is_pad is honored: padded (start-of-episode zero) frames are replaced by the nearest real frame before the STSS correlation, so they contribute “no motion” instead of spurious motion against a blank frame. This mirrors the space-time encoder, which masks exactly these frames out of temporal attention. (The RLDX-1 reference does not mask them; we do, to avoid a train/inference distribution mismatch at episode start.)

Classes

RLDXVideoEncoder(vision_tower, ...[, ...])

Plain SigLIP video encoder with an RLDX-1 STSS motion module.
class opentau.policies.pi05_mem.rldx_video_encoder.RLDXVideoEncoder(vision_tower: SiglipVisionModel, multi_modal_projector: Module, max_num_frames: int, gradient_checkpointing: bool = False, *, motion_insert_layer: int | None = None, motion_hidden_dim: int = 256, motion_window: tuple[int, int, int] = (5, 9, 9), motion_corr_func: str = 'cosine', motion_n_encoders: int = 1, motion_norm: str = 'groupnorm', motion_int_mode: str = 'lite', motion_zero_init: bool = True)[source]

Bases: Module

Plain SigLIP video encoder with an RLDX-1 STSS motion module.

The caller owns vision_tower and multi_modal_projector (constructed by a PaliGemmaWithExpertModel / Gemma3WithExpertModel); this module holds them by reference (via a list, so their parameters are not re-registered under this module’s path) and adds only the motion module’s parameters.

Parameters:

  • vision_tower – A SiglipVisionModel (left unmodified — no layer wrapping).

  • multi_modal_projector – SigLIP-hidden -> VLA-hidden projector.

  • max_num_frames – Upper bound on T accepted by forward (validation cap).

  • gradient_checkpointing – Wrap each SigLIP layer forward in torch.utils.checkpoint during training.

  • motion_insert_layer – 0-indexed encoder layer after which the motion residual is injected. None -> mid-stack (n_layers // 3), which is layer 9 for the 27-layer so400m tower (RLDX placement).

  • motion_hidden_dim – Internal correlation/feature width (in_proj target).

  • motion_window – (L, kh, kw) space-time correlation window; spatial size must fit the patch grid.

  • motion_corr_func – “cosine” / “dotproduct” / “dotproduct_softmax”.

  • motion_n_encoders – Number of stacked STSS encoders (outputs summed).

  • motion_norm – “groupnorm” (default; per-sample, no cross-rank sync — safe under FSDP/DeepSpeed/multi-rank) / “batchnorm” (RLDX-faithful) / “syncbn”.

  • motion_int_mode – “lite” (single fuse conv) or “full” (3x3 conv stack).

  • motion_zero_init – Zero-init the residual so the module starts as a no-op.

__init__(vision_tower: SiglipVisionModel, multi_modal_projector: Module, max_num_frames: int, gradient_checkpointing: bool = False, *, motion_insert_layer: int | None = None, motion_hidden_dim: int = 256, motion_window: tuple[int, int, int] = (5, 9, 9), motion_corr_func: str = 'cosine', motion_n_encoders: int = 1, motion_norm: str = 'groupnorm', motion_int_mode: str = 'lite', motion_zero_init: bool = True)[source]

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

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].

  • obs_history_is_pad – Optional (B, T) bool mask, True for padded history frames. Padded frames are replaced by the nearest real frame before the STSS motion correlation so they don’t inject spurious motion into the current-frame residual.

Returns:

(B, num_video_tokens, vlm_hidden_size) current-frame tokens.

property multi_modal_projector: Module
property vision_tower: SiglipVisionModel