opentau.policies.pi05_mem.motion_module

Space-Time Self-Similarity (STSS) motion module.

Ported from RLDX-1 (rldx/model/modules/backbone/motion.py, https://github.com/RLWRLD/RLDX-1), whose motion module in turn builds on the SELFY space-time self-similarity formulation (Kwon et al., “Learning Self-Similarity in Space and Time as Generalized Motion for Video Action Recognition”, ICCV 2021).

Rather than optical flow or raw frame differences, the module captures temporal dynamics by computing local space-time self-similarity: for every spatio-temporal patch feature it correlates against its neighbors within a (L, kh, kw) window across frames, producing a similarity volume that a small 3D-conv encoder turns into a motion feature. The result is a residual that is added back onto the vision features:

v~(t) = v(t) + S_theta(STSS(v(t)))

In OpenTau this residual is injected at a single SigLIP encoder layer (see SpaceTimeSiglipVideoEncoder in video_encoder.py).

Interface (matching the RLDX-1 reference so the math is easy to compare):

forward(x, grid_sizes)

x: (sum_i T_i*H_i*W_i, C) flattened patch features, token order
(b t h w) (batch-major, then time, then row-major patches).

grid_sizes: (B, 3) int tensor; each row is [T, H, W] for one video. returns: (sum_i T_i*H_i*W_i, C) residual, same layout as x.

Differences from the RLDX-1 reference, all behind flags:

norm (default “groupnorm”): GroupNorm(1, C) — per-sample, no cross-rank stat sync, safe under FSDP/DeepSpeed/multi-rank (CLAUDE.md rule #5). Pass “batchnorm” for the RLDX-1-faithful BatchNorm3d, or “syncbn”.
zero_init_residual (default True): zero-initializes the output projection (or the LayerScale) so the module starts as an exact no-op. A policy fine-tuned from a pre-existing pi05 checkpoint is therefore byte-identical at step 0 and the motion contribution warms up during training. RLDX-1 dropped this zero-init so motion contributes from step 1 — pass zero_init_residual=False to match that.
The gradient-monitoring print hook from the reference is omitted.

Classes

`MotionModule`(d_in, d_hid, d_out[, window, ...])	STSS motion module producing a residual update to vision features.
`STSSEncoder`(d_in, d_hid, d_out[, window, ...])	One STSS block: LN -> in_proj -> transform -> extract -> integrate -> out_proj.
`STSSExtraction`([window, chnls, norm])	Project the `kh*kw` similarity channels of the STSS volume to features.
`STSSIntegration`(d_in[, window, chnls, norm, ...])	Fuse the `L` temporal-window slices into a single motion feature map.
`STSSTransformation`([window, corr_func])	Compute the local space-time self-similarity volume.

class opentau.policies.pi05_mem.motion_module.MotionModule(d_in: int, d_hid: int, d_out: int, window: tuple[int, int, int] = (5, 9, 9), ext_chnls: tuple[int, ...] = (256,), int_chnls: tuple[int, int, int] = (256, 256, 512), corr_func: str = 'cosine', n_encoders: int = 1, use_layerscale: bool = False, layerscale_init: float = 1e-05, norm: str = 'groupnorm', int_mode: str = 'lite', zero_init_residual: bool = True)[source]

Bases: Module

STSS motion module producing a residual update to vision features.

Parameters:

d_in – Input feature dimension (must equal the residual target dim).
d_hid – Internal correlation/feature dimension (in_proj target).
d_out – Output dimension; set equal to d_in so the residual adds back.
window – (L, kh, kw) space-time correlation window.
int_chnls (ext_chnls /) – channel widths for the extraction / integration convs.
corr_func – “cosine” (default), “dotproduct”, or “dotproduct_softmax”.
n_encoders – number of stacked STSS encoders (their outputs are summed).
use_layerscale – gate the output with a learnable per-channel LayerScale instead of an out_proj linear.
layerscale_init – initial LayerScale value (used only with use_layerscale).
norm – “groupnorm” (default, distributed-safe) / “batchnorm” (RLDX-faithful) / “syncbn”.
int_mode – “lite” (single fuse conv) or the full 3x3 conv stack.
zero_init_residual – when True, zero-initialize the residual output so the module starts as an exact no-op (warm start from a pretrained checkpoint); the contribution ramps up during training.

__init__(d_in: int, d_hid: int, d_out: int, window: tuple[int, int, int] = (5, 9, 9), ext_chnls: tuple[int, ...] = (256,), int_chnls: tuple[int, int, int] = (256, 256, 512), corr_func: str = 'cosine', n_encoders: int = 1, use_layerscale: bool = False, layerscale_init: float = 1e-05, norm: str = 'groupnorm', int_mode: str = 'lite', zero_init_residual: bool = True)[source]: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, grid_sizes: Tensor) → Tensor[source]

Parameters:

x – (sum_i T_i*H_i*W_i, C) flattened patch features, order (b t h w).
grid_sizes – (B, 3) int tensor; each row [T, H, W].

Returns:

(sum_i T_i*H_i*W_i, d_out) residual, same layout as x.

initialize_weights() → None[source]: Initialize all submodule weights.

class opentau.policies.pi05_mem.motion_module.STSSEncoder(d_in: int, d_hid: int, d_out: int, window: tuple[int, int, int] = (5, 9, 9), ext_chnls: tuple[int, ...] = (256,), int_chnls: tuple[int, int, int] = (256, 256, 512), corr_func: str = 'cosine', norm: str = 'groupnorm', int_mode: str = 'lite')[source]

Bases: Module

One STSS block: LN -> in_proj -> transform -> extract -> integrate -> out_proj.

__init__(d_in: int, d_hid: int, d_out: int, window: tuple[int, int, int] = (5, 9, 9), ext_chnls: tuple[int, ...] = (256,), int_chnls: tuple[int, int, int] = (256, 256, 512), corr_func: str = 'cosine', norm: str = 'groupnorm', int_mode: str = 'lite')[source]: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, grid_sizes: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.pi05_mem.motion_module.STSSExtraction(window: tuple[int, int, int] = (5, 9, 9), chnls: tuple[int, ...] = (256,), norm: str = 'groupnorm')[source]

Bases: Module

Project the kh*kw similarity channels of the STSS volume to features.

__init__(window: tuple[int, int, int] = (5, 9, 9), chnls: tuple[int, ...] = (256,), norm: str = 'groupnorm')[source]: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.pi05_mem.motion_module.STSSIntegration(d_in: int, window: tuple[int, int, int] = (5, 9, 9), chnls: tuple[int, int, int] = (64, 64, 64), norm: str = 'groupnorm', mode: str = 'lite')[source]

Bases: Module

Fuse the L temporal-window slices into a single motion feature map.

__init__(d_in: int, window: tuple[int, int, int] = (5, 9, 9), chnls: tuple[int, int, int] = (64, 64, 64), norm: str = 'groupnorm', mode: str = 'lite')[source]: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.pi05_mem.motion_module.STSSTransformation(window: tuple[int, int, int] = (5, 9, 9), corr_func: str = 'cosine')[source]

Bases: Module

Compute the local space-time self-similarity volume.

For each spatio-temporal patch feature, correlate it against the features of nearby frames within a window[0] temporal span, then keep only a local window[1] x window[2] spatial neighborhood of the correlation.

__init__(window: tuple[int, int, int] = (5, 9, 9), corr_func: str = 'cosine')[source]: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, grid_sizes: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.