opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level

π05: A Vision-Language-Action Flow Model for General Robot Control

[Paper](https://www.physicalintelligence.company/download/pi05.pdf)

Functions

`build_attention_inputs`(use_flash, pad_masks, ...)	Return `(attention_mask, attention_block_ids)` for the active backend.
`create_sinusoidal_pos_embedding`(time, ...[, ...])	Computes sine-cosine positional embedding vectors for scalar positions.
`flash_cuda_active`(attention_implementation)	True if the `flash_cuda` backend is selected.
`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 with the specified value.

Classes

`ContextItem`(data, item_type, pad_mask[, ...])	One token-block contributed to the prefix or suffix sequence.
`PI07PaligemmaLowLevelFlowMatching`(config[, ...])	π07 Low-Level — A Vision-Language-Action Flow Matching model with SpaceTime SigLIP video, subtask/subgoal conditioning, and episode metadata.
`PI07PaligemmaLowLevelPolicy`(config[, ...])	Wrapper class around PI07PaligemmaLowLevelFlowMatching model to train and run inference within OpenTau.

class opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level.ContextItem(data: Tensor, item_type: Literal['text', 'image', 'video', 'state', 'discrete_action', 'action'], pad_mask: Tensor, attention: Literal['continue', 'bidirectional', 'causal'] = 'continue', exclude_from_cross_attention: bool = False, obs_history_is_pad: Tensor | None = None)[source]

Bases: object

One token-block contributed to the prefix or suffix sequence.

The PI07PaligemmaLowLevelPolicy builds a list of ContextItem``s and hands it to the ``PI07PaligemmaLowLevelFlowMatching model. The model embeds each item by item_type (text, video, state, discrete_action, action), concatenates the embeddings, and derives the 1-D attention mask from each item’s attention setting. To add or rearrange blocks in the prefix/suffix, edit the policy’s list construction — the model is layout-agnostic.

Fields:

data: Per-type input. text/discrete_action: token IDs: (B, L). video: pixel video (B, T, C, H, W) in [0, 1]. image: single pixel image (B, C, H, W) in [-1, 1] (subgoal images, embedded via the VLM’s image tower). state: continuous state (B, T, max_state_dim). action: noisy actions (B, chunk_size, max_action_dim).

item_type: Dispatch key for which embedding path to take. pad_mask: Padding mask. (B,) per-sample is allowed for

video/image (the model expands it to (B, n_tokens) with n_tokens the grid-derived video_encoder.num_video_tokens). All other types must pass (B, L) matching the embedded sequence length. True = real, False = padded.

attention: 1-D attention pattern for this block:

"continue": [0]*L — token continues the previous block’s attention scope (bidirectional with everything before).
"bidirectional": [1] + [0]*(L-1) — opens a new block boundary, then bidirectional within the block.
"causal": [1]*L — every token starts its own scope (fully causal within this block).

exclude_from_cross_attention: When True, this item’s tokens are

not cross-attended to by the suffix (action expert). Used for the trailing "Action: " indicator + discrete-action block, which the action expert must not condition on.

obs_history_is_pad: Optional (B, T) mask threaded into the

video encoder’s temporal attention (only used by video items). True marks padded history frames.

__init__(data: Tensor, item_type: Literal['text', 'image', 'video', 'state', 'discrete_action', 'action'], pad_mask: Tensor, attention: Literal['continue', 'bidirectional', 'causal'] = 'continue', exclude_from_cross_attention: bool = False, obs_history_is_pad: Tensor | None = None) → None

attention: Literal['continue', 'bidirectional', 'causal'] = 'continue'

data: Tensor

exclude_from_cross_attention: bool = False

item_type: Literal['text', 'image', 'video', 'state', 'discrete_action', 'action']

obs_history_is_pad: Tensor | None = None

pad_mask: Tensor

class opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level.PI07PaligemmaLowLevelFlowMatching(config: PI07PaligemmaLowLevelConfig, discrete_action_vocab_size: int | None = None, num_datasets: int = 1)[source]

Bases: Module

π07 Low-Level — A Vision-Language-Action Flow Matching model with SpaceTime SigLIP video, subtask/subgoal conditioning, and episode metadata.

Architecturally inherits the π0.5 PaliGemma + Gemma-Expert flow-matching backbone ([π0.5 paper](https://www.physicalintelligence.company/download/pi05.pdf)) and extends the prefix to ingest, on top of the original (image, language, state, discrete actions) tokens:

Image history — encoded by SpaceTimeSiglipVideoEncoder. n_obs_steps frames per camera; T=1 is byte-identical to plain SigLIP, T>1 inserts space-time separable temporal attention every spacetime_layer_stride SigLIP layers (MEM-paper recipe).
Per-timestep robot state stream (B, T, max_state_dim) — one VLM token per history step, masked according to obs_history_is_pad (current step is always real).
Subtask response tokens — the natural-language subtask emitted by PI07HighLevelPlannerModel; fully masked out per-sample when the high-level planner is dropped, in which case only the leading attention “boundary” token is kept.
Subgoal image(s) — optional future-state visual goal frame(s), routed through the same SpaceTime SigLIP encoder; per-sample masked, and the subgoal SigLIP forward is skipped entirely when no sample in the batch has a subgoal (saves ~2× SigLIP compute).
Episode metadata tokens — optional "Speed: …, Quality: …, Mistake: …" string assembled from speed / quality / mistake features and tokenized to metadata_max_length.

The flow-matching action expert (Gemma expert with adaRMS time conditioning, num_steps-step Euler integration, FAST-token discrete-action prefix that is excluded from continuous-action cross-attention via n_cross_att_tokens) is unchanged from π0.5.

┌────────────────────────────────────────────────────────────┐ │ actions │ │ ▲ │ │ ┌──┴───┐ │ │ kv cache │Gemma │ │ │ ┌─────────────────►│Expert│ ◄── adaRMS(time) │ │ │ │ │ │ │ ┌┴────────────┐ │ x 10 │ │ │ │ │ └──▲───┘ │ │ │ PaliGemma │ │ │ │ │ │ noise │ │ └▲─▲─▲─▲─▲─▲─▲ │ │ │ │ │ │ │ │ └── discrete actions │ │ │ │ │ │ │ └──── episode metadata │ │ │ │ │ │ └────── subgoal image(s) │ │ │ │ │ └──────── subtask response │ │ │ │ └────────── robot state │ │ │ └──────────── language tokens │ │ └────────────── video (image history) │ └────────────────────────────────────────────────────────────┘

__init__(config: PI07PaligemmaLowLevelConfig, discrete_action_vocab_size: int | None = None, num_datasets: int = 1)[source]

Initializes the PI07PaligemmaLowLevelFlowMatching 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[source]

Apply one denoising step of the noise x_t at a given timestep.

Parameters:

prefix_pad_masks – Prefix padding masks.
past_key_values – Past key values from the VLM.
x_t – Current noise tensor.
time – Time tensor of shape (batch_size, action_chunk_length).

Returns:

The predicted velocity tensor (v_t).

embed_prefix(items: list[ContextItem], group_index: Tensor | None = None, sync_across_ranks: bool = True) → tuple[Tensor, Tensor, Tensor, int][source]

Embed a layout-agnostic list of ``ContextItem``s into the prefix.

The policy decides what blocks go into the prefix and in what order; this method only embeds them by type, concatenates, and derives the 1-D attention mask. Items flagged exclude_from_cross_attention are dropped from the num_cross_att_tokens count returned alongside the embeddings, so the action expert never cross-attends to them (e.g. trailing "Action: " indicator + discrete-action tokens).

Parameters:: items – Ordered list of ContextItem blocks. Per the current architecture, items flagged exclude_from_cross_attention must be a contiguous trailing run.
Returns:: (embs, pad_masks, att_masks, num_cross_att_tokens).

embed_suffix(items: list[ContextItem], timestep: Tensor, group_index: Tensor | None = None) → tuple[Tensor, Tensor, Tensor, Tensor][source]

Embed the action expert’s suffix from a list of ``ContextItem``s.

Mirrors embed_prefix(). The timestep stays a separate argument (not a ContextItem) because it conditions the action expert via adaRMS rather than appearing as a sequence token.

Parameters:

items – Ordered ContextItem blocks (today: just a single "action" item containing the noisy actions; the policy may add more in the future).
timestep – (B, chunk_size) flow-matching time, fed through a sine-cosine positional encoding and a 2-layer MLP to produce the adaRMS conditioning vector.

Returns:

(embs, pad_masks, att_masks, adarms_cond).

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

Encode a video through the SpaceTime SigLIP video encoder.

At T=1 this is byte-identical to paligemma_with_expert.embed_image.

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. Passed to the video encoder so that temporal attention blocks padded frames.

Returns:

(B, num_video_tokens, vlm_hidden_size)

forward(prefix_items: list[ContextItem], actions: Tensor, discrete_actions: Tensor, discrete_action_masks: Tensor, actions_is_pad: Tensor | None = None, noise: Tensor | None = None, time: Tensor | None = None, real_action_dim: Tensor | None = None, group_index: Tensor | None = None, return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss][source]

Do a full training forward pass and compute the loss.

Parameters:

prefix_items – Ordered ContextItem list defining the entire VLM prefix layout. The "Action: " indicator and the discrete-action item must be flagged exclude_from_cross_attention=True so the continuous action expert does not condition on them.
actions – Continuous action targets (B, chunk_size, max_action_dim).
discrete_actions – FAST discrete-action token IDs (B, discrete_action_max_length), used for the CE loss.
discrete_action_masks – Mask for discrete_actions.
actions_is_pad – Optional (B, chunk_size) mask for padded actions.
noise – Optional pre-sampled noise tensor.
time – Optional pre-sampled flow-matching time tensor.

Returns:

A dictionary containing the loss components (“MSE” and “CE”).

generate_discrete_actions(prefix_items: list[ContextItem], max_new_tokens: int, stop_fn: Callable[[list[list[int]]], list[bool]], group_index: Tensor | None = None) → list[list[int]][source]

Greedy autoregressive generation of FAST discrete-action tokens.

Replays the training-time token layout: the prefix (which must end with the "Action: " indicator block) runs through the VLM backbone once with fill_kv_cache=True over its FULL length — unlike the flow-matching path, the cache here exists for token generation, not for action-expert cross-attention, so the indicator is not excluded. The first token is predicted from the last indicator hidden state (training computes the CE logits from prefix_out[:, -L-1:-1]; position -L-1 is exactly this token). Each subsequent step re-feeds the generated-so-far block with use_cache=True, fill_kv_cache=False — the cached prefix KV is prepended inside _run_layer, the same mechanism denoise_step() uses — with causal attention within the block and position ids continuing past the real prefix tokens, matching training’s cumsum(pad_masks) - 1 because generated tokens are never padded.

Parameters:

prefix_items – Prefix ContextItem list ending with the "Action: " indicator.
max_new_tokens – Hard cap on generated tokens (training’s discrete_action_max_length truncation).
stop_fn – list[list[int]] -> list[bool] per-sample completion predicate, re-evaluated after every step (the FAST coefficient-count rule lives in the policy).
group_index – Per-sample norm/projection group index.

Returns:

Per-sample token id lists, each frozen at its stop point (or at max_new_tokens if the predicate never fired).

sample_actions(prefix_items: list[ContextItem], action_prefix: Tensor, delay: Tensor, noise: Tensor | None = None, group_index: Tensor | None = None) → Tensor[source]

Do a full inference forward and compute the action.

Parameters:

prefix_items – Ordered ContextItem list defining the entire VLM prefix layout. Inference omits the discrete-action block, so no items need to be excluded from cross-attention.
action_prefix – Frozen action prefix (B, chunk_size, max_action_dim).
delay – Number of frozen delay actions at the start of the chunk.
noise – Optional pre-sampled noise.

Returns:

The sampled action tensor.

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

Samples Gaussian noise.

Parameters:

shape – The shape of the noise tensor.
device – The device to create the tensor on.

Returns:

A tensor containing the sampled noise.

sample_time(bsize: int, device: device | str) → Tensor[source]

Samples time steps from a Beta distribution.

Parameters:

bsize – Batch size.
device – The device to create the tensor on.

Returns:

A tensor containing the sampled time steps.

class opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level.PI07PaligemmaLowLevelPolicy(config: PI07PaligemmaLowLevelConfig, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Bases: PreTrainedPolicy

Wrapper class around PI07PaligemmaLowLevelFlowMatching model to train and run inference within OpenTau.

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

Initializes the PI07PaligemmaLowLevelPolicy.

Parameters:

config – Policy configuration class instance.
per_dataset_stats – Ordered list of per-dataset stat dicts used to fill the stacked Normalize/Unnormalize buffers.
dataset_names – Ordered list parallel to per_dataset_stats.

config_class: alias of PI07PaligemmaLowLevelConfig

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

Do a full training forward pass to compute the loss.

Parameters:

batch – Batch of data containing environment observations, actions, and targets.
noise – Optional noise tensor.
time – Optional time tensor.
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:

A dictionary containing the loss components (“MSE” and “CE”). When the warn_outlier_threshold check finds offending dims, a non-empty "outlier_records" list is also included for the training loop to log.

Override the from_pretrained method to handle key remapping.

Parameters:

pretrained_name_or_path – Path to the pretrained model or its name on the Hub.
config – Configuration object.
force_download – Whether to force download the model weights.
resume_download – Whether to resume download.
proxies – Proxy configuration.
token – Authentication token.
cache_dir – Directory to cache downloaded files.
local_files_only – Whether to only look for files locally.
revision – Specific model revision.
strict – Whether to strictly enforce state dict matching.
**kwargs – Additional keyword arguments.

Returns:

The loaded model instance.

Raises:

ValueError – If pretrained_name_or_path is None.

get_optim_params() → dict[source]

Returns the parameters to be optimized.

Returns:: A generator over the model parameters.

name: None = 'pi07_paligemma_low_level': The name of the policy. Must be defined in subclasses.

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

Predict a chunk of actions given environment observations.

Parameters:: batch – Batch of data containing environment observations.
Returns:: The predicted action chunk.
Raises:: NotImplementedError – Always, as this method is not implemented for PI05.

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

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

Parameters:

batch – Batch of data containing the key “discrete_actions”.

Returns:

discrete_action_tokens: A tensor of shape (batch_size, max_length) containing the tokenized actions.
discrete_action_masks: A tensor of shape (batch_size, max_length) indicating valid tokens.

Return type:

A tuple containing

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

Tokenize the text input.

Parameters:

batch – Batch of data containing the key “prompt” and “state”.

Returns:

lang_tokens: Tensor of language tokens.
lang_masks: Tensor of language attention masks.

Return type:

A tuple containing

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

Tokenize episode metadata into PaliGemma token IDs.

Builds strings Metadata: Speed: … Quality: … Mistake: … Robot: … FPS: … Control: … only from fields whose *_is_pad flag is False (robot_type / control_mode are strings, omitted when empty). For each sample, a padded or empty field is omitted entirely (not concatenated to segments). When every field is pad, the row is "" before tokenization.

fps is the effective per-sample frame rate (torch.long); the segment is omitted when fps_is_pad is True. Segment ordering is Speed → Quality → Mistake → Robot → FPS → Control to match the other pi07 / pi07_paligemma prepare_metadata implementations.

Always runs _hydrate_metadata_batch() first so callers see the same outcomes whether they hit sample_actions() (often no metadata keys), forward() with a partial-key batch, or a dataloader batch with every key set.

Values are normalized to shape (B,) on state’s device, with scalar tensors broadcast like subgoal_is_pad.

Note

speed_is_pad / quality_is_pad / mistake_is_pad / fps_is_pad each default to torch.ones(B, dtype=bool) (treat-as-pad) when the key is missing from batch. This is a behavior change from the previous treat-as-real (torch.zeros) default. Hand-built inference batches that supply "speed" / "quality" / "mistake" / "fps" but omit the corresponding _is_pad flag will now have those metadata fields silently dropped from the prefix string — pass ..._is_pad=torch.zeros(...) explicitly when the metadata is real.

Parameters:: batch – Batch dict; metadata keys may be missing or partial before hydration (see _hydrate_metadata_batch()).
Returns:: (metadata_tokens, metadata_masks) with shapes (B, metadata_max_length). Never None.

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

Tokenize the high-level planner subtask response.

Wraps each response string as "Subtask: {response}, " and pads/truncates to response_max_length. Uses add_special_tokens=False so no BOS token is inserted.

Parameters:: batch – Batch dict containing "response" (list of strings).
Returns:: (response_tokens, response_masks) each of shape (B, response_max_length).

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

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

Accepts either (B, D) (single timestep) or (B, T, D) (multi-timestep history). Single-timestep tensors are unsqueezed to (B, 1, D) so downstream code always receives 3-D state.

Parameters:: batch – Batch of data containing the "state" tensor.
Returns:: A tensor of shape (B, T, max_state_dim).
Raises:: ValueError – If the state dimension exceeds max_state_dim or if n_obs_steps > 1 but a 2-D state is provided.

prepare_subgoal_images(batch: dict[str, Tensor]) → tuple[list[Tensor], list[Tensor]][source]

Preprocess subgoal images for the VLM image tower.

Derives subgoal keys from config.image_features: for each camera{k} the corresponding batch key is subgoal{k}. If no subgoal{k} keys are present in the batch, zero-filled (B, C, H, W) tensors with masks all False are returned so that the prefix sequence length stays fixed.

Each image is resized with aspect-ratio padding and rescaled from [0, 1] to [-1, 1] (the range SigLIP expects), then embedded via PaliGemmaWithExpertModel.embed_image() as a 4-D (B, C, H, W) tensor (byte-identical to embed_video at T=1 on the shared vision tower).

When batch["subgoal_is_pad"] is True for a sample, the subgoal slots for that sample are zeroed out and masks cleared.

If no subgoal{k} keys are present (common during sample_actions when the environment does not provide subgoals), this returns the same tensor shapes as a padded training batch but with all masks False, which matches an all-padded batch from the dataloader in embed_prefix (no subgoal conditioning for any sample).

subgoal_is_pad may be omitted (defaults to all-True), scalar, or shape (batch,).

Parameters:: batch – Batch dict containing subgoal image tensors keyed as subgoal{k} for each camera{k} in config.image_features.
Returns:: A tuple (subgoal_images, subgoal_img_masks) of lists, where each image has shape (B, C, H, W) in [-1, 1].

prepare_videos(batch: dict[str, Tensor]) → tuple[list[Tensor], list[Tensor]][source]

Preprocess camera inputs into (B, T, C, H, W) video tensors.

Pixel values stay in [0, 1]; the SpaceTime SigLIP encoder rescales to [-1, 1] internally. Padded history frames are zeroed at the pixel level here (read from batch["obs_history_is_pad"]) AND masked inside the video encoder via temporal attention, mirroring pi07.

Parameters:: batch – Batch of data containing image/video tensors. May contain obs_history_is_pad (B, T) marking padded history frames. 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 (videos, vid_masks) where each element is a list with one entry per camera.

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.

Note: The provided action_prefix should NOT be normalized, as this method will handle normalization internally. The action_prefix should have shape (batch_size, action_chunk_length, action_dim) where action_chunk_length is less than or equal to config.chunk_size.

Parameters:

batch –
Batch of data containing environment observations. response may be omitted (defaults to "" per sample). Episode metadata is normalized before tokenization (same rules as training — see _hydrate_metadata_batch()):
- None of speed / quality / mistake / *_is_pad present (common at inference): behaves like an all-padded metadata batch (no conditioning).
- Mixed keys: any subset may be present; missing fields use those same defaults so tokenization matches a fully-specified batch.
speed / quality are floats; mistake is torch.bool. subgoal{k} / subgoal_is_pad may be omitted; omitted subgoals are treated like an all-padded dataloader batch (no subgoal conditioning), consistent with training when every sample has subgoal_is_pad=True.
action_prefix – Optional action prefix tensor of shape (batch_size, action_chunk_length, action_dim).
delay – Optional number of frozen delay actions from action_prefix.
noise – Optional noise tensor.

Returns:

The sampled actions tensor of shape (batch_size, action_chunk_length, action_dim).

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

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

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

Parameters:

batch – Environment observation dict.
noise – Optional pre-sampled noise for deterministic evaluation.

Returns:

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

opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level.build_attention_inputs(use_flash: bool, pad_masks: Tensor, att_masks: Tensor, n_cross_att_tokens: int | None = None, cross_att_pad_masks: Tensor | None = None) → tuple[Tensor | None, tuple[Tensor, ...] | None][source]

Return (attention_mask, attention_block_ids) for the active backend.

For the flash_cuda backend (use_flash=True) this builds the compact block-id representation and returns attention_mask=None so the dense (B, Sq, Sk) mask is never materialized. Otherwise it builds the dense mask via make_att_2d_masks() and returns attention_block_ids=None.

opentau.policies.pi07_paligemma.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[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) containing the positional embeddings.

Raises:

ValueError – If dimension is not divisible by 2 or if time tensor is not 2-D with shape (batch_size, action_chunk_length).

opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level.flash_cuda_active(attention_implementation: str) → bool[source]

True if the flash_cuda backend is selected.

No fallback: when flash_cuda is selected it is always used. If the custom CUDA kernel cannot be JIT-compiled (no CUDA / no compiler), the kernel call raises loudly rather than silently degrading to eager/sdpa.

opentau.policies.pi07_paligemma.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[source]

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

Tokens can attend to valid inputs tokens which have a cumulative att_masks smaller or equal to theirs. This way att_masks int[B, N] can be used to setup several types of attention, for example:

[[1 1 1 1 1 1]]: pure causal attention.

[[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
themselves and the last 3 tokens have a causal attention. The first entry could also be a 1 without changing behaviour.

[[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
block can attend all previous blocks and all tokens on the same block.

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 n_cross_att_tokens is provided.
cross_att_pad_masks – Padding masks for cross attention tokens. Required if n_cross_att_tokens is provided.

Returns:

A 2D attention mask tensor of shape (B, N + n_cross_att_tokens, N + n_cross_att_tokens) if n_cross_att_tokens is provided, else (B, N, N).

Raises:

ValueError – If att_masks or pad_masks are not 2D (including batch dimension).
AssertionError – If cross_att_pad_masks is missing when n_cross_att_tokens is set, or if its shape is incorrect.

opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level.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 (lists of integers).
max_length – The target length for the discrete action token sequences.

Returns:

discrete_action_tokens: A numpy array of shape (len(tokens), max_length) containing the padded discrete action tokens.
discrete_action_masks: A boolean numpy array of shape (len(tokens), max_length) indicating valid discrete action tokens (True) and padding (False).

Return type:

A tuple containing

opentau.policies.pi07_paligemma.low_level.modeling_pi07_low_level.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 with the specified value.

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

Raises:

ValueError – If the input image tensor does not have 4 dimensions.