opentau.policies.pi05.modeling_pi05

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

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

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 with the specified value.

Classes

`PI05FlowMatching`(config[, ...])	π05: A Vision-Language-Action Flow Model for General Robot Control
`PI05Policy`(config[, per_dataset_stats, ...])	Wrapper class around PI05FlowMatching model to train and run inference within OpenTau.

class opentau.policies.pi05.modeling_pi05.PI05FlowMatching(config: PI05Config, discrete_action_vocab_size: int | None = None, language_tokenizer: AutoTokenizer | None = None)[source]

Bases: Module

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

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

┌──────────────────────────────────────────┐ │ actions │ │ ▲ │ │ ┌┴─────┐ │ │ kv cache │Gemma │ │ │ ┌──────────►│Expert│ │ │ │ │ │ │ │ ┌┴─────────┐ │x 10 │ │ │ │ │ └▲─────┘ │ │ │PaliGemma │ │ │ │ │ │ noise │ │ └▲──▲──▲──▲ │ │ │ │ │ └── discrete actions │ │ │ │ └───── robot state │ │ │ └──────── language tokens │ │ └─────────── image(s) │ └──────────────────────────────────────────┘

__init__(config: PI05Config, discrete_action_vocab_size: int | None = None, language_tokenizer: AutoTokenizer | None = None)[source]

Initializes the PI05FlowMatching model.

Parameters:

config – Model configuration.
discrete_action_vocab_size – Size of the discrete action vocabulary.
language_tokenizer – Optional pre-loaded PaliGemma tokenizer to share with the enclosing PI05Policy. When None (e.g. unit tests that construct the inner module directly) the tokenizer is loaded here. Either way, the same instance is used by both layers, and ensure_loc_tokens runs once.

denoise_step(prefix_pad_masks: Tensor, past_key_values: list[dict[str, Tensor]], x_t: Tensor, time: Tensor) → 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(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, response_tokens: Tensor | None = None, response_masks: Tensor | None = None, discrete_actions: Tensor | None = None, discrete_action_masks: Tensor | None = None, state: Tensor | None = None) → tuple[Tensor, Tensor, Tensor, Tensor][source]

Embed images with SigLIP and language tokens with embedding layer to prepare for PaliGemma transformer processing.

When state_type == "continuous" and state is provided, the raw state vector is projected into VLM embedding space and appended as a single token (with full attention to images and language). Response tokens are ignored in this mode.

Parameters:

images – List of image tensors.
img_masks – List of image mask tensors.
lang_tokens – Language token tensor.
lang_masks – Language mask tensor.
response_tokens – Optional Response language token tensor.
response_masks – Optional Response language mask tensor.
discrete_actions – Optional discrete action tensor.
discrete_action_masks – Optional discrete action mask tensor.
state – Optional continuous state tensor of shape (batch_size, max_state_dim).
indicator_tokens – Optional indicator token tensor.
indicator_masks – Optional indicator mask tensor.

Returns:

embs: Concatenated embeddings tensor.
pad_masks: Concatenated padding masks tensor.
att_masks: Attention masks tensor.
segment_ids: Per-token modality slot tensor (shape matches pad_masks). Used to add the modality-specific embedding on top of embs when config.use_modality_embedding is set; also returned for downstream/diagnostic use.

Return type:

A tuple containing

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

Embed noisy_actions, timestep to prepare for Expert Gemma processing.

Parameters:

noisy_actions – Tensor containing noisy actions.
timestep – Tensor containing timesteps of shape (batch_size, action_chunk_length).

Returns:

embs: Concatenated embeddings tensor.
pad_masks: Concatenated padding masks tensor.
att_masks: Attention masks tensor.
adarms_cond: AdaRMS conditioning tensor.

Return type:

A tuple containing

forward(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, actions: Tensor, actions_is_pad: Tensor | None = None, response_tokens: Tensor | None = None, response_masks: Tensor | None = None, noise: Tensor | None = None, time: Tensor | None = None, discrete_actions: Tensor | None = None, discrete_action_masks: Tensor | None = None, state: 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.

Parameters:

images – List of image tensors.
img_masks – List of image mask tensors.
lang_tokens – Language token tensor.
lang_masks – Language mask tensor.
response_tokens – Response language token tensor.
response_masks – Response language mask tensor.
actions – Action tensor.
actions_is_pad – Optional action is padded mask tensor.
noise – Optional noise tensor.
time – Optional time tensor.
discrete_actions – Optional discrete action tensor.
discrete_action_masks – Optional discrete action mask tensor.
state – Optional continuous state tensor of shape (batch_size, max_state_dim).
indicator_tokens – Optional indicator token tensor.
indicator_masks – Optional indicator mask tensor.

Returns:

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

infer_response(prefix_out: Tensor, prefix_embs: Tensor, prefix_pad_masks: Tensor, prefix_att_masks: Tensor, past_key_values: list[dict[str, Tensor]], prefix_offsets: Tensor, response_tokens: Tensor, auto_step: int, bsize: int, device: device) → tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, list[dict[str, Tensor]], Tensor][source]

Perform autoregressive inference for response generation.

This method generates the next token in the response sequence, updating the various state tensors required for maintaining the generation context. It handles the initial BOS token generation as well as subsequent tokens, and manages padding masks to handle variable-length sequences properly.

Parameters:

prefix_out – Output tensor from the previous step.
prefix_embs – Embeddings for the current prefix context.
prefix_pad_masks – Boolean mask indicating valid (non-padding) tokens in the prefix.
prefix_att_masks – Attention mask for the prefix.
past_key_values – KV cache from previous transformer steps.
prefix_offsets – Position offsets for the current generation step.
response_tokens – Accumulated tokens generated so far.
auto_step – Current autoregressive step index (0 for first token).
bsize – Batch size.
device – Device to run the computation on.

Returns:

(prefix_out, prefix_embs, prefix_pad_masks, prefix_att_masks,: prefix_offsets, response_tokens, past_key_values, response_token)

Return type:

A tuple containing updated tensors for the next step

sample_actions(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, action_prefix: Tensor, delay: Tensor, noise: Tensor | None = None, state: Tensor | None = None) → Tensor[source]

Do a full inference forward and compute the action.

Parameters:

images – List of image tensors.
img_masks – List of image mask tensors.
lang_tokens – Language token tensor.
lang_masks – Language mask tensor.
action_prefix – Action prefix tensor.
delay – Number of delay actions, aka number of actions frozen from the action_prefix.
noise – Optional noise tensor.
state – Optional continuous state tensor of shape (batch_size, max_state_dim).
indicator_tokens – Optional indicator token tensor.
indicator_masks – Optional indicator mask tensor.

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.pi05.modeling_pi05.PI05Policy(config: PI05Config, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Bases: PreTrainedPolicy

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

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

Initializes the PI05Policy.

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 PI05Config

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.

Parameters:

batch – Batch of data containing environment observations, actions, and targets.
noise – Optional noise tensor.
time – Optional time tensor.
return_per_sample – When True, additionally return per-sample MSE_per_sample / CE_per_sample (PerSampleLoss) so the validation loop can bucket the loss by (dataset, control_mode) provenance. The scalar MSE/CE are unchanged, so the training path is unaffected.

Returns:

A dictionary with the loss components (“MSE” and “CE”), plus “MSE_per_sample”/”CE_per_sample” when return_per_sample is True.

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 = 'pi05': 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_discrete_state(batch: dict[str, Tensor]) → list[str][source]

Discretizes the state into bins and converts it to a string representation.

Each dimension of the state vector is discretized into 256 bins. The values of each dimension of the state are expected to be in the range [-1, 1]. The discretization bins are linearly spaced between -1 and 1. The index of the bin for each dimension is then concatenated into a space-separated string.

Parameters:: batch – Batch of data containing the “state” tensor.
Returns:: A list of strings, where each string is a space-separated list of discretized state values.
Raises:: ValueError – If the state values are not normalized between -1 and 1.

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

Apply preprocessing to the images.

Resizes to 224x224 and padding to keep aspect ratio, and converts pixel range from [0.0, 1.0] to [-1.0, 1.0] as requested by SigLIP.

Parameters:

batch – Batch of data containing image tensors.

Returns:

images: A list of processed image tensors.
img_masks: A list of image mask tensors.

Return type:

A tuple containing

Raises:

ValueError – If no image features are present in the batch.

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

Tokenize the text input.

When state_type == "discrete", the state is discretized into bins and embedded into the prompt string. When state_type == "continuous", the state is handled separately via prepare_state(), so the prompt only contains the task description.

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

Tokenize the response input.

Parameters:

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

Returns:

response_tokens: Tensor of response language tokens.
response_masks: Tensor of response language attention masks.

Return type:

A tuple containing

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

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

Only used when state_type == "continuous".

Parameters:: batch – Batch of data containing the “state” tensor of shape (batch_size, state_dim).
Returns:: A tensor of shape (batch_size, max_state_dim).
Raises:: ValueError – If the state dimension exceeds max_state_dim.

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.
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 given environment observations.

This method uses an action queue that is replenished when it has config.max_delay or fewer actions (or is empty). When replenishing, the current queue contents are used as action_prefix for sample_actions, then the queue is refilled with the new chunk.

Note: This method should only be called when running a policy in simulation. For real world inference, this method should be written in the ROS client node.

Parameters:

batch – Batch of data containing environment observations.
noise – Optional noise tensor to be used during sampling.

Returns:

The selected action tensor.

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

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