opentau.policies.pi07.low_level.modeling_pi07_low_level

π07 Low-Level Component: a Vision-Language-Action Flow Model for continuous action generation.

The low-level component is one half of the π07 hierarchical architecture. Given video observations (encoded by SpaceTimeSiglip), a language prompt, an optional subtask response from the high-level planner, temporal proprioceptive state, optional subgoal images, and optional metadata, it produces continuous action chunks via flow matching while simultaneously predicting discrete (FAST) action tokens through the VLM backbone.

Key differences from the base π05 policy:

1. SpaceTimeSiglip video encoder replaces SigLIP, compressing each camera’s temporal video into 256 tokens via a Perceiver cross-attention reducer.

2. Temporal state sequences (B, T, D) are projected per-timestep into separate continuous tokens for the Gemma backbone.

3. Supports optional subtask response, subgoal image, and metadata conditioning for hierarchical planning.

4. Knowledge Insulation: action-expert gradients are detached from the VLM backbone to preserve language understanding capabilities.

Functions

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

PI07LowLevelFlowMatching(config[, ...])	π07 Low-Level Component: flow-matching action generation with Knowledge Insulation.
PI07LowLevelPolicy(config[, ...])	Policy wrapper for the π07 low-level component.

class opentau.policies.pi07.low_level.modeling_pi07_low_level.PI07LowLevelFlowMatching(config: PI07LowLevelConfig, discrete_action_vocab_size: int | None = None, num_datasets: int = 1)

Bases: Module

π07 Low-Level Component: flow-matching action generation with Knowledge Insulation.

Architecture overview:

┌───────────────────────────────────────────────────┐
│                   actions                         │
│                   ▲                               │
│                  ┌┴─────┐                         │
│      kv cache    │Gemma │  (detached)             │
│      ┌──────────►│Expert│                         │
│      │           │      │                         │
│     ┌┴─────────┐ │x 10  │                         │
│     │          │ └▲─────┘                         │
│     │PaliGemma │  │                               │
│     │  (VLM)   │  noise                           │
│     └▲──▲──▲──▲──▲──▲──▲──▲                       │
│      │  │  │  │  │  └── ``Action:`` + discrete (training) │
│      │  │  │  │  └───── ``";

“`` + metadata │

│ │ │ │ └──────── subgoal images, Subgoal: │ │ │ │ └────────── response, commas, state, State: │ │ │ └──────────── language │ │ └────────────── video (SpaceTimeSiglip) │ └───────────────────────────────────────────────────┘

The VLM processes the same prefix layout as embed_prefix() (videos, language, State:, state, commas, response, Subgoal:, subgoal images, ``”;

“``,

optional Action:/discrete, metadata). The action expert receives the prefix KV-cache (detached for Knowledge Insulation) together with noisy continuous actions and flow-matching timestep embeddings to predict the velocity field.

__init__(config: PI07LowLevelConfig, discrete_action_vocab_size: int | None = None, num_datasets: int = 1)

Initializes the PI07LowLevelFlowMatching model.

Parameters:

• config – Model configuration.

• discrete_action_vocab_size – Size of the discrete action vocabulary.

• num_datasets – Number of (robot_type, control_mode) groups — the leading dim of the stacked Normalize/Unnormalize buffers. Used as the number of per-group projection heads when config.per_group_projection is set, so the projection axis stays in lockstep with the norm-head axis. Defaults to 1.

denoise_step(prefix_pad_masks: Tensor, past_key_values: list[dict[str, Tensor]], x_t: Tensor, time: Tensor, group_index: Tensor | None = None) → Tensor

Run one action-expert forward pass to predict the velocity field.

Embeds the suffix (noisy actions + timestep), constructs the cross-attention mask to the cached prefix, and runs the Gemma expert to produce the predicted velocity v_t.

Parameters:

• prefix_pad_masks – (B, prefix_len) padding mask from the prefix pass (used for cross-attention masking).

• past_key_values – Cached KV states from the VLM prefix pass.

• x_t – Current noisy actions (B, chunk_size, max_action_dim).

• time – Per-step timestep (B, chunk_size).

Returns:

Predicted velocity (B, n_action_steps, max_action_dim).

embed_prefix(videos: list[Tensor], vid_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, state: Tensor, response_tokens: Tensor, response_masks: Tensor, metadata_tokens: Tensor, metadata_masks: Tensor, discrete_actions: Tensor | None = None, discrete_action_masks: Tensor | None = None, obs_history_is_pad: Tensor | None = None, subgoal_images: list[Tensor] | None = None, subgoal_img_masks: list[Tensor] | None = None, group_index: Tensor | None = None, sync_across_ranks: bool = True) → tuple[Tensor, Tensor, Tensor]

Embed all prefix modalities and build the 1-D attention pattern.

Concatenation order – matches π0.7 paper Fig. 19’s “image goals come after the text prompt” rule (subgoal block sits at the tail, just before the optional discrete-action block):

[videos | language | State: | state(T) | <state-end> | response? | metadata? | ";\n "? | Subgoal: subgoal_images... ", "? | ("Action:" + discrete_actions only when training)]

<state-end> is ", " when at least one optional middle block (response / subgoal / metadata) contributes real tokens, else ":\n". The trailing ";\n " prefix-end is only emitted in the former case; without optional content the state-end already serves as the separator before "Action: ", so appending another would dangle spurious tokens.

Note (batch-wide gate): has_any_optional is computed by OR-ing the masks across the whole batch, so a sample with no optional content sharing a batch with samples that do have optionals still receives the ", " state-end and ";\n " prefix-end (with all-False masks on the absent optional block). This keeps every sample in the batch aligned to the same prefix layout — all-or-nothing per batch — which is needed because the prefix length determines per-sample position IDs and cross-attention offsets that must be uniform within a batch. Acceptable in practice for the homogeneous batches used in training; per-sample gating would require variable-length prefixes per sample.

Attention pattern (via att_masks cumsums). Paper §VI.B says “observation tokens and subgoal image tokens use bidirectional attention within themselves … the following text tokens use causal attention”:

• Video patches: one bidirectional block ([0] * N).

• All text spans (language, State:, state-end ", " / ":\n", response, metadata, ";\n ", Subgoal:, trailing ", " after subgoals, Action:): one causal block per token ([1] * N).

• State tokens (per-history-step linear projections): one bidirectional block ([1] + [0] * (T-1)).

• Subgoal image patches: one bidirectional block per camera ([1] + [0] * (N-1)).

• Discrete-action FAST tokens (training only): causal [1] * N.

Parameters:

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

• vid_masks – List of boolean masks, each (B,).

• lang_tokens – Language token IDs (B, prompt_max_length).

• lang_masks – Boolean mask for language tokens.

• state – Temporal state (B, T, max_state_dim).

• response_tokens – Subtask response token IDs (B, response_max_length).

• response_masks – Boolean mask for response tokens.

• metadata_tokens – Metadata token IDs (B, metadata_max_length).

• metadata_masks – Boolean mask for metadata tokens.

• discrete_actions – Optional FAST token IDs (B, discrete_action_max_length). Provided during training.

• discrete_action_masks – Boolean mask for discrete actions.

• obs_history_is_pad – Optional (B, T) bool tensor; True for padded timesteps. Used to mask state tokens during training.

• subgoal_images – List of subgoal image tensors (B, C, H, W).

• subgoal_img_masks – List of boolean masks (B,).

Returns:

• embs: (B, total_seq_len, D)

• pad_masks: (B, total_seq_len)

• att_masks: (B, total_seq_len) for make_att_2d_masks().

Return type:

A tuple (embs, pad_masks, att_masks) where

embed_suffix(noisy_actions: Tensor, timestep: Tensor, group_index: Tensor | None = None) → tuple[Tensor, Tensor, Tensor, Tensor]

Embed noisy actions and flow-matching timestep for the action expert.

Projects actions through action_in_proj and computes an adaRMS-style conditioning vector from sinusoidal timestep embeddings via a two-layer MLP. The suffix forms a single bidirectional block (att_masks = [1, 0, …, 0]).

Parameters:

• noisy_actions – (B, chunk_size, max_action_dim) noisy action tensor at the current flow-matching timestep.

• timestep – (B, chunk_size) per-step timestep values.

Returns:

A tuple (embs, pad_masks, att_masks, adarms_cond) where adarms_cond is the conditioning vector for adaptive RMSNorm in the Gemma expert layers.

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

Encode a video through SpaceTimeSiglip + Perceiver reducer + projection.

Parameters:

• video – (B, T, C, H, W)

• 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)

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, subgoal_images: list[Tensor] | None = None, subgoal_img_masks: list[Tensor] | None = None, metadata_tokens: Tensor | None = None, metadata_masks: Tensor | None = None, response_tokens: Tensor | None = None, response_masks: Tensor | None = None, real_action_dim: Tensor | None = None, group_index: Tensor | None = None, return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss]

Training forward pass: embed all modalities and compute losses.

Runs the VLM on the prefix (video, language, response, state, subgoal images, metadata, discrete actions), then the action expert on the noisy action suffix. Returns both the flow-matching MSE loss and the discrete-action cross-entropy loss.

The VLM’s KV-cache is detached before being passed to the action expert (Knowledge Insulation).

Parameters:

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

• vid_masks – List of boolean masks (B,).

• lang_tokens – Language token IDs (B, prompt_max_length).

• lang_masks – Boolean mask for language tokens.

• state – Temporal state (B, T, max_state_dim).

• actions – Ground-truth continuous actions (B, chunk_size, max_action_dim).

• actions_is_pad – Optional (B, chunk_size) bool mask for padded actions.

• noise – Optional pre-sampled noise.

• time – Optional pre-sampled flow-matching timesteps.

• discrete_actions – Optional FAST token IDs.

• discrete_action_masks – Optional mask for discrete actions.

• obs_history_is_pad – Optional (B, T) mask for padded frames.

• subgoal_images – Optional list of subgoal image tensors.

• subgoal_img_masks – Optional list of masks.

• metadata_tokens – Optional metadata token IDs.

• metadata_masks – Optional mask for metadata tokens.

• response_tokens – Optional subtask response token IDs.

• response_masks – Optional mask for response tokens.

Returns:

Dict with "MSE" (flow-matching loss) and "CE" (discrete action loss) scalar tensors.

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, subgoal_images: list[Tensor] | None = None, subgoal_img_masks: list[Tensor] | None = None, metadata_tokens: Tensor | None = None, metadata_masks: Tensor | None = None, response_tokens: Tensor | None = None, response_masks: Tensor | None = None, obs_history_is_pad: Tensor | None = None, group_index: Tensor | None = None) → Tensor

Inference: iteratively denoise to produce a continuous action chunk.

Embeds the prefix (without discrete actions), caches the VLM KV states, then runs num_steps denoising iterations through the action expert. Prefix action steps (up to delay) are held fixed from action_prefix.

Parameters:

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

• vid_masks – List of boolean masks (B,).

• lang_tokens – Language token IDs (B, prompt_max_length).

• lang_masks – Boolean mask for language tokens.

• state – Temporal state (B, T, max_state_dim).

• action_prefix – Previously committed actions (B, chunk_size, max_action_dim) (zero-padded beyond delay).

• delay – Scalar tensor indicating how many prefix steps are fixed.

• noise – Optional pre-sampled noise.

• subgoal_images – Optional list of subgoal image tensors.

• subgoal_img_masks – Optional list of masks.

• metadata_tokens – Optional metadata token IDs.

• metadata_masks – Optional mask for metadata tokens.

• response_tokens – Optional subtask response token IDs.

• response_masks – Optional mask for response tokens.

• obs_history_is_pad – Optional (B, T) bool mask flagging padded history slots (True = padded). Emitted by PI07LowLevelPolicy._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.

Returns:

Denoised action chunk (B, chunk_size, max_action_dim).

sample_noise(shape: tuple[int, ...], device: device | str) → Tensor

sample_time(bsize: int, device: device | str) → Tensor

class opentau.policies.pi07.low_level.modeling_pi07_low_level.PI07LowLevelPolicy(config: PI07LowLevelConfig, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)

Bases: PreTrainedPolicy

Policy wrapper for the π07 low-level component.

Handles tokenization, normalization, observation-history buffering, and action queue management around the inner PI07LowLevelFlowMatching model.

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

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 PI07LowLevelConfig

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

Training forward pass: normalize, prepare modalities, and compute losses.

Returns a dict with "MSE" (flow-matching velocity loss) and "CE" (discrete action cross-entropy loss). When the warn_outlier_threshold check finds offending dims, a non-empty "outlier_records" list is also included for the training loop to log.

Parameters:

• batch – Training batch dict with observations, actions, and prompts.

• noise – Optional pre-sampled noise tensor.

• time – Optional pre-sampled flow-matching timesteps.

• return_per_sample – When True, also returns per-sample MSE_per_sample/CE_per_sample (PerSampleLoss) for the validation per-(dataset, control_mode) breakdown. Scalars unchanged.

Returns:

Dict with "MSE" and "CE" scalar loss tensors (plus per-sample entries when return_per_sample is True).

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

Override the from_pretrained method to handle key remapping.

get_optim_params() → dict

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 = 'pi07_low_level'

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

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

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

Tokenize continuous actions into discrete FAST tokens and pad to fixed length.

Parameters:

batch – Batch dict containing "discrete_actions" (min-max normalized).

Returns:

A tuple (token_ids, token_masks) with shapes (B, discrete_action_max_length).

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

Tokenize the language prompt into PaliGemma token IDs.

Wraps each prompt string as "Task: {task}<eos>" and pads/truncates to prompt_max_length.

Parameters:

batch – Batch dict containing "prompt" (list of strings).

Returns:

A tuple (lang_tokens, lang_masks) with shapes (B, prompt_max_length).

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

Tokenize episode metadata into Gemma 3 token IDs.

Wraps non-empty per-sample metadata segments into a single "Metadata: {seg1} {seg2} ..." string, then pads/truncates to metadata_max_length. Samples with no active segments emit an empty string.

Parameters:

batch – Batch dict that may contain any of: "speed", "quality", "mistake", "fps" (numeric tensors with a corresponding _is_pad bool tensor — entries marked as pad are dropped), and "robot_type", "control_mode" (lists of strings — empty string is the pad signal, no separate _is_pad flag). Missing keys are treated as fully padded.

Returns:

A tuple (metadata_tokens, metadata_masks) with shapes (B, metadata_max_length).

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

Tokenize the high-level planner subtask response into PaliGemma token IDs.

Wraps each response string as "Subtask: {response}, " and pads/truncates to response_max_length. Uses add_special_tokens=False so no BOS (or other special tokens) are inserted; the prefix already encodes a "Subtask: " span in PI07LowLevelFlowMatching.embed_prefix().

Parameters:

batch – Batch dict containing "response" (list of strings).

Returns:

A tuple (response_tokens, response_masks) with shapes (B, response_max_length).

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

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_subgoal_images(batch: dict[str, Tensor]) → tuple[list[Tensor], list[Tensor]]

Preprocess subgoal images for SigLIP embedding.

Derives subgoal keys from config.image_features: for each camera{k} the corresponding batch key is subgoal{k} (the naming convention used by LeRobotDataset._emit_optional_keys). If no subgoal{k} keys are present, a warning is logged and ([], []) is returned; downstream gating in embed_prefix short-circuits the subgoal branch when the lists are empty.

Resizes each subgoal image to 224×224 with aspect-ratio padding and converts the pixel range from [0, 1] to [-1, 1] as expected by SigLIP. Missing cameras are filled with -1 padding tensors up to empty_cameras.

When batch["subgoal_is_pad"] is True for a sample, all subgoal slots for that sample are zeroed out and their masks set to False so that downstream attention ignores them.

Parameters:

batch – Batch dict containing subgoal image tensors keyed as subgoal{k} for each camera{k} in config.image_features. If all are absent, see warning + fallback above.

Returns:

A tuple (subgoal_images, subgoal_img_masks) of lists.

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

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. The dataset loader is expected to yield raw [0, 1] floats — the SpaceTime SigLIP encoder rescales to [-1, 1] internally (see video_encoder.SpaceTimeSiglipVideoEncoder.forward), so callers should NOT pre-apply ImageNet normalization.

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 SpaceTimeSiglip does not process 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

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

Sample a full action chunk via flow-matching inference.

Normalizes inputs, prepares all modalities (video, language, response, state, subgoal images, metadata), and delegates to the inner model’s sample_actions for iterative denoising.

Parameters:

• batch – Environment observation dict.

• action_prefix – Optional previously-committed actions for real-time inference with delay.

• delay – Number of prefix action steps already committed.

• noise – Optional pre-sampled noise for deterministic evaluation.

Returns:

Action chunk tensor of shape (B, n_action_steps, action_dim).

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

Select a single action from the queue, regenerating if needed.

Builds temporal observation history when configured, runs flow-matching inference to fill the action queue, and pops the next action.

Parameters:

• batch – Environment observation dict with "state", image keys, "prompt", "response", and optional "metadata".

• noise – Optional pre-sampled noise for deterministic evaluation.

Returns:

A single action tensor of shape (B, action_dim).

supports_torch_compile: bool = True

Whether maybe_compile_for_training() may compile this policy’s self.model. Default False: opt-in per policy, because the in-place nn.Module.compile only takes effect if the policy’s training forward invokes the submodule via self.model(...) (__call__) rather than self.model.forward(...). Subclasses whose forward has been switched to self.model(...) (currently PI05Policy and PI07LowLevelPolicy) set this True. Leaving it False makes use_torch_compile=True a loud no-op on unwired policies instead of a silent compile-but-never-dispatch.

opentau.policies.pi07.low_level.modeling_pi07_low_level.create_sinusoidal_pos_embedding(time: Tensor, dimension: int, min_period: float, max_period: float, device: device | str = 'cpu') → Tensor

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.pi07.low_level.modeling_pi07_low_level.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

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.pi07.low_level.modeling_pi07_low_level.pad_discrete_tokens(tokens: list[list[int]], max_length: int) → tuple[ndarray, ndarray]

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.pi07.low_level.modeling_pi07_low_level.resize_with_pad(img: Tensor, width: int, height: int, pad_value: int = -1) → Tensor

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