opentau.policies.pi05_mem.modeling_pi05

π05 Mem: A Vision-Language-Action Flow Model with space-time SigLIP video encoding and temporal state sequences.

Based on π05, this variant implements the low-level memory architecture from Torne, Pertsch, Walke et al. “MEM: Multi-Scale Embodied Memory for Vision Language Action Models” (Section III-C + Appendix C):

  1. Extends the PaliGemma SigLIP image encoder with space-time separable attention every spacetime_layer_stride-th ViT layer. The temporal sublayer re-uses each layer’s existing Q/K/V/O projections — no new learnable parameters are introduced. Past-timestep tokens are dropped after the encoder so the prefix matches a single-frame VLA’s 256 image tokens exactly.

  2. Accepts temporal state sequences (B, T, D) and projects each timestep into a separate continuous token for the Gemma backbone.

Functions

build_mrope_prefix_positions(seg_masks, ...)

Build interleaved-MRoPE (t, h, w) position ids for the prefix.

create_sinusoidal_pos_embedding(time, ...[, ...])

Computes sine-cosine positional embedding vectors for scalar positions.

make_att_2d_masks(pad_masks, att_masks[, ...])

Creates a 2-D attention mask given padding and 1-D attention masks.

mrope_suffix_position_ids(...)

Text-style interleaved-MRoPE positions for the action-expert suffix.

pad_discrete_tokens(tokens, max_length)

Pads or truncates a list of discrete action token sequences to a fixed length.

resize_with_pad(img, width, height[, pad_value])

Resizes an image to fit within the specified dimensions while maintaining aspect ratio, and pads the remaining area.

Classes

PI05MemFlowMatching(config[, ...])

π05 Mem: A Vision-Language-Action Flow Model with space-time SigLIP video encoding and temporal state sequences.

PI05MemPolicy(config[, per_dataset_stats, ...])

Wrapper class around PI05MemFlowMatching model.
class opentau.policies.pi05_mem.modeling_pi05.PI05MemFlowMatching(config: PI05MemConfig, discrete_action_vocab_size: int | None = None)[source]

Bases: Module

π05 Mem: A Vision-Language-Action Flow Model with space-time SigLIP video encoding and temporal state sequences.

┌──────────────────────────────────────────┐ │ actions │ │ ▲ │ │ ┌┴─────┐ │ │ kv cache │Gemma │ │ │ ┌──────────►│Expert│ │ │ │ │ │ │ │ ┌┴─────────┐ │x 10 │ │ │ │ │ └▲─────┘ │ │ │PaliGemma │ │ │ │ │ │ noise │ │ └▲──▲──▲──▲ │ │ │ │ │ └── discrete actions │ │ │ │ └───── state (T tokens) │ │ │ └──────── language tokens │ │ └─────────── video (SigLIP+ST) │ └──────────────────────────────────────────┘

__init__(config: PI05MemConfig, discrete_action_vocab_size: int | None = None)[source]

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

denoise_step(prefix_pad_masks: Tensor, prefix_position_ids: Tensor, past_key_values: list[dict[str, Tensor]], x_t: Tensor, time: Tensor) → Tensor[source]

Apply one denoising step.

embed_prefix(videos: list[Tensor], vid_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, state: Tensor, discrete_actions: Tensor | None = None, discrete_action_masks: Tensor | None = None, obs_history_is_pad: Tensor | None = None) → tuple[Tensor, Tensor, Tensor, Tensor][source]

Embed videos with the space-time SigLIP video encoder, language tokens with the embedding layer, and temporal state via per-timestep learned projection.

Parameters:

  • videos – List of video tensors, each (B, T, C, H, W).

  • vid_masks – List of video mask tensors, each (B,).

  • lang_tokens – Language token tensor.

  • lang_masks – Language mask tensor.

  • state – Temporal state tensor of shape (B, T, max_state_dim).

  • discrete_actions – Optional discrete action tensor.

  • discrete_action_masks – Optional discrete action mask tensor.

  • obs_history_is_pad – Optional bool tensor (B, T) from the dataloader. True for padded (clamped) timesteps, False for real ones. Used to mask state tokens during training; None during inference.

Returns:

(embs, pad_masks, att_masks, position_ids) tuple. position_ids is [B, L] for rope_type="rope" and [3, B, L] (interleaved MRoPE; only video tokens get 2-D spatial positions) otherwise.

embed_suffix(noisy_actions: Tensor, timestep: Tensor) → tuple[Tensor, Tensor, Tensor, Tensor][source]

Embed noisy_actions, timestep to prepare for Expert Gemma processing.

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

Encode a video through the space-time SigLIP video encoder.

The encoder applies standard SigLIP spatial attention on every layer plus a causal temporal attention sublayer every spacetime_layer_stride-th layer. Past-timestep tokens are dropped; only the current frame’s 256 tokens are returned.

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. Threaded into the SpaceTime SigLIP encoder so temporal attention blocks padded frames (pixel- zeroing alone is insufficient — the patch embedding bias and temporal PE for t < T-1 are non-zero, so zero pixels still produce non-zero hidden states the current frame would otherwise attend to).

Returns:

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

forward(videos: list[Tensor], vid_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, state: Tensor, actions: Tensor, actions_is_pad: Tensor | None = None, noise: Tensor | None = None, time: Tensor | None = None, discrete_actions: Tensor | None = None, discrete_action_masks: Tensor | None = None, obs_history_is_pad: Tensor | None = None, real_action_dim: Tensor | None = None, return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss][source]

Do a full training forward pass and compute the loss.

sample_actions(videos: list[Tensor], vid_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, state: Tensor, action_prefix: Tensor, delay: Tensor, noise: Tensor | None = None, obs_history_is_pad: Tensor | None = None) → Tensor[source]

Do a full inference forward and compute the action.

Parameters:

obs_history_is_pad – Optional (B, T) bool mask flagging padded history slots (True = padded). Emitted by PI05MemPolicy._build_history_batch so the encoder can use real mid-episode history while still masking out the start-of-episode zero-fill. None falls back to “all history padded except current” via embed_prefix and the encoder’s None-fallback.

sample_noise(shape: tuple[int, ...], device: device | str) → Tensor[source]
sample_time(bsize: int, device: device | str) → Tensor[source]
class opentau.policies.pi05_mem.modeling_pi05.PI05MemPolicy(config: PI05MemConfig, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Bases: PreTrainedPolicy

Wrapper class around PI05MemFlowMatching model.

Uses a space-time SigLIP video encoder (MEM paper low-level memory) and temporal state sequences projected into per-timestep continuous tokens in the VLM embedding space.

__init__(config: PI05MemConfig, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Initializes the PreTrainedPolicy.

Parameters:

  • config – The configuration object for the policy.

  • *inputs – Variable length argument list.

  • **kwargs – Arbitrary keyword arguments.

Raises:

ValueError – If config is not an instance of PreTrainedConfig.

config_class

alias of PI05MemConfig

forward(batch: dict[str, Tensor], noise: Tensor | None = None, time: Tensor | None = None, return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss][source]

Do a full training forward pass to compute the loss.

When return_per_sample is True, also returns per-sample MSE_per_sample/CE_per_sample (PerSampleLoss) for the validation per-(dataset, control_mode) breakdown; the scalar losses are unchanged.

classmethod from_pretrained(pretrained_name_or_path: str | Path, *, config: PreTrainedConfig | None = None, force_download: bool = False, resume_download: bool | None = None, proxies: dict | None = None, token: str | bool | None = None, cache_dir: str | Path | None = None, local_files_only: bool = False, revision: str | None = None, strict: bool = True, **kwargs) → T[source]

Override the from_pretrained method to handle key remapping.

get_optim_params() → dict[source]

Returns the policy-specific parameters dict to be passed on to the optimizer.

Returns:

A dictionary of parameters to optimize.

Return type:

dict

name: None = 'pi05_mem'

The name of the policy. Must be defined in subclasses.

predict_action_chunk(batch: dict[str, Tensor]) → Tensor[source]
prepare_discrete_actions(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Prepares discrete actions for the model by tokenizing and padding them.

prepare_language(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Tokenize the text input.

prepare_state(batch: dict[str, Tensor]) → Tensor[source]

Prepares the temporal state tensor, padding or truncating to max_state_dim.

Parameters:

batch – Batch of data containing “state” tensor of shape (B, T, D).

Returns:

A tensor of shape (B, T, max_state_dim).

prepare_videos(batch: dict[str, Tensor], obs_history_is_pad: Tensor | None = None) → tuple[list[Tensor], list[Tensor]][source]

Apply preprocessing to the video inputs.

Each camera key now contains a video tensor of shape (B, T, C, H, W). Frames are resized to 224x224 with padding. Pixel values remain in the [0, 1] range as produced by the dataset loader; the video encoder rescales to [-1, 1] (SigLIP’s expected range) inside its own forward pass.

Parameters:

  • batch – Batch of data containing video tensors.

  • obs_history_is_pad – Optional bool tensor (B, T) indicating which temporal frames are padded. Padded frames are zeroed out before encoding so the video encoder does not see clamped/repeated content. The current frame (t = T-1) is never zeroed, even when flagged: the dataset’s history_state_drop_prob augmentation marks obs_history_is_pad all-True while keeping the current step.

Returns:

A tuple of (videos, vid_masks) lists.

reset() → None[source]

This should be called whenever the environment is reset.

sample_actions(batch: dict[str, Tensor], action_prefix: Tensor | None = None, delay: Tensor | None = None, noise: Tensor | None = None) → Tensor[source]

Sample actions from the policy given environment observations.

select_action(batch: dict[str, Tensor], noise: Tensor | None = None) → Tensor[source]

Select a single action given environment observations.

opentau.policies.pi05_mem.modeling_pi05.build_mrope_prefix_positions(seg_masks: list[Tensor], seg_is_video: list[bool], grid: int) → Tensor[source]

Build interleaved-MRoPE (t, h, w) position ids for the prefix.

The prefix is a concatenation of ordered segments (videos, language, state, optionally discrete actions). Each segment is either a square video patch grid or a run of text-style tokens:

  • Video segment (grid * grid row-major patches): all patches share a single temporal index t = base (the encoder collapses history into one current frame, so there is no per-token time axis), and receive 2-D spatial positions h = base + row / w = base + col. A present video advances the running cursor by grid (= max(row, col) + 1), matching Qwen-style “text resumes after the vision block” continuity. A padded (absent) video advances nothing.

  • Text segment: t == h == w, incrementing by 1 per real (non-pad) token, so MRoPE degenerates to ordinary 1-D RoPE here.

The per-sample cursor mirrors the cumsum(pad_masks) - 1 progression of the 1-D path, but a video block consumes grid position units instead of grid * grid.

Parameters:

  • seg_masks – Per-segment bool pad masks, each [B, seg_len], in prefix order (True = real token).

  • seg_is_video – Parallel list flagging which segments are video grids.

  • grid – Side length of the (square) video patch grid.

Returns:

[3, B, L] long tensor of (temporal, height, width) positions.

opentau.policies.pi05_mem.modeling_pi05.create_sinusoidal_pos_embedding(time: Tensor, dimension: int, min_period: float, max_period: float, device: device | str = 'cpu') → Tensor[source]

Computes sine-cosine positional embedding vectors for scalar positions.

Parameters:

  • time – A 2-D tensor of shape (batch_size, action_chunk_length).

  • dimension – The dimension of the embedding vectors. Must be divisible by 2.

  • min_period – The minimum period of the sinusoidal functions.

  • max_period – The maximum period of the sinusoidal functions.

  • device – The device to create the tensors on. Defaults to “cpu”.

Returns:

A tensor of shape (batch_size, action_chunk_length, dimension).

opentau.policies.pi05_mem.modeling_pi05.make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor, n_cross_att_tokens: int | None = None, cross_att_pad_masks: Tensor | None = None) → Tensor[source]

Creates a 2-D attention mask given padding and 1-D attention masks.

Parameters:

  • pad_masks – bool[B, N] true if its part of the input, false if padding.

  • att_masks – int32[B, N] mask that’s 1 where previous tokens cannot depend on it and 0 where it shares the same attention mask as the previous token.

  • n_cross_att_tokens – Add attention mask for cross-attention tokens if provided.

  • cross_att_pad_masks – Padding masks for cross attention tokens.

Returns:

A 2D attention mask tensor.

opentau.policies.pi05_mem.modeling_pi05.mrope_suffix_position_ids(prefix_position_ids: Tensor, prefix_pad_masks: Tensor, suffix_pad_masks: Tensor, num_cross_att_tokens: int) → Tensor[source]

Text-style interleaved-MRoPE positions for the action-expert suffix.

The suffix continues from the position right after the prefix region the expert cross-attends to: offset = max(real prefix positions over the cross region) + 1 per sample. Suffix tokens are text-style (t == h == w), so this is the 1-D offset + cumsum(pad) - 1 progression broadcast to all three axes.

Only real (non-pad) prefix tokens contribute to the offset. This matters because a padded prefix token’s position is not “right after the real content”: an absent video does not advance the cursor yet its patch block still carries h/w = base..base+grid-1 (see build_mrope_prefix_positions), and a padded text/state token repeats a prior position. Masking padded tokens out of the max makes offset equal the MRoPE cursor after the cross region (real-token max + 1), so the suffix continues with no position gap. Note this is NOT the 1-D path’s token-count sum(pad_masks): a video block compacts grid*grid patches into grid position units, so the two schemes’ offsets differ whenever the cross region contains video.

Parameters:

  • prefix_position_ids – [3, B, L] prefix MRoPE positions.

  • prefix_pad_masks – [B, L] bool prefix pad mask (True = real).

  • suffix_pad_masks – [B, Ls] bool suffix pad mask.

  • num_cross_att_tokens – Number of leading prefix tokens cached for cross attention (the region the suffix offset continues from).

Returns:

[3, B, Ls] long tensor of suffix positions.

opentau.policies.pi05_mem.modeling_pi05.pad_discrete_tokens(tokens: list[list[int]], max_length: int) → tuple[ndarray, ndarray][source]

Pads or truncates a list of discrete action token sequences to a fixed length.

Parameters:

  • tokens – A list of discrete action token sequences.

  • max_length – The target length.

Returns:

A tuple of (discrete_action_tokens, discrete_action_masks) numpy arrays.

opentau.policies.pi05_mem.modeling_pi05.resize_with_pad(img: Tensor, width: int, height: int, pad_value: int = -1) → Tensor[source]

Resizes an image to fit within the specified dimensions while maintaining aspect ratio, and pads the remaining area.

Parameters:

  • img – Input image tensor of shape (batch_size, channels, current_height, current_width).

  • width – Target width.

  • height – Target height.

  • pad_value – Value to use for padding. Defaults to -1.

Returns:

The resized and padded image tensor of shape (batch_size, channels, height, width).