opentau.policies.pi0.modeling_pi0

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

[Paper](https://www.physicalintelligence.company/download/pi0.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.
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

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

class opentau.policies.pi0.modeling_pi0.PI0FlowMatching(config: PI0Config)

Bases: Module

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

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

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

__init__(config: PI0Config)

Initializes the PI0FlowMatching model.

Parameters:

config – Model configuration.

denoise_step(state: Tensor, prefix_pad_masks: Tensor, past_key_values: list[dict[str, Tensor]], x_t: Tensor, timestep: Tensor) → Tensor

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

Parameters:

• state – State tensor.

• prefix_pad_masks – Prefix padding masks.

• past_key_values – Past key values from the VLM.

• x_t – Current noise tensor.

• timestep – Current timestep.

Returns:

The predicted velocity tensor (v_t).

embed_prefix(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor) → tuple[Tensor, Tensor, Tensor]

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

Parameters:

• images – List of image tensors.

• img_masks – List of image mask tensors.

• lang_tokens – Language token tensor.

• lang_masks – Language mask tensor.

Returns:

• embs: Concatenated embeddings tensor.

• pad_masks: Concatenated padding masks tensor.

• att_masks: Attention masks tensor.

Return type:

A tuple containing

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

Embed state, noisy_actions, timestep to prepare for Expert Gemma processing.

Parameters:

• state – State tensor.

• noisy_actions – Tensor containing noisy actions.

• timestep – Tensor containing timesteps.

Returns:

• embs: Concatenated embeddings tensor.

• pad_masks: Concatenated padding masks tensor.

• att_masks: Attention masks tensor.

Return type:

A tuple containing

forward(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, state: Tensor, actions: Tensor, noise: Tensor | None = None, time: Tensor | None = None) → Tensor

Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors).

Parameters:

• images – List of image tensors.

• img_masks – List of image mask tensors.

• lang_tokens – Language token tensor.

• lang_masks – Language mask tensor.

• state – State tensor.

• actions – Action tensor.

• noise – Optional noise tensor.

• time – Optional time tensor.

Returns:

The computed loss tensor.

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

Do a full inference forward and compute the action (batch_size x num_steps x num_motors).

Parameters:

• images – List of image tensors.

• img_masks – List of image mask tensors.

• lang_tokens – Language token tensor.

• lang_masks – Language mask tensor.

• state – State tensor.

• noise – Optional noise tensor.

Returns:

The sampled action tensor.

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

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

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.

set_requires_grad() → None

Sets the requires_grad attribute for state projection parameters.

class opentau.policies.pi0.modeling_pi0.PI0Policy(config: PI0Config, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)

Bases: PreTrainedPolicy

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

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

Initializes the PI0Policy.

Parameters:

• config – Policy configuration class instance.

• per_dataset_stats – Ordered list of per-dataset stat dicts used to fill the stacked Normalize/Unnormalize buffers. May be None when constructing for a checkpoint load — in that case config.dataset_names is consulted for the leading dim.

• dataset_names – Ordered list parallel to per_dataset_stats.

config_class

alias of PI0Config

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

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. CE is a zero stub for pi0, so CE_per_sample carries zero sum and count.

Returns:

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

classmethod from_pretrained(*args, **kwargs)

Override the from_pretrained method to display important disclaimer.

Parameters:

• *args – Positional arguments passed to super().from_pretrained.

• **kwargs – Keyword arguments passed to super().from_pretrained.

Returns:

The loaded model instance.

get_optim_params() → dict

Returns the parameters to be optimized.

Returns:

A generator over the model parameters.

name: None = 'pi0'

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

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

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

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

Apply Pi0 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]

Tokenize the text input.

Parameters:

batch – Batch of data containing “prompt” and potentially “advantage”.

Returns:

• lang_tokens: Tensor of language tokens.

• lang_masks: Tensor of language attention masks.

Return type:

A tuple containing

reset() → None

This should be called whenever the environment is reset.

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

Sample actions from the policy given environment observations.

Parameters:

• batch – Batch of data containing environment observations.

• 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

Select a single action given environment observations.

This method wraps select_actions in order to return one action at a time for execution in the environment. It works by managing the actions in a queue and only calling select_actions when the queue is empty.

Parameters:

• batch – Batch of data containing environment observations.

• noise – Optional noise tensor to be used during sampling.

Returns:

The selected action tensor.

opentau.policies.pi0.modeling_pi0.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 1-D tensor of shape (batch_size,).

• 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, dimension) containing the positional embeddings.

Raises:

ValueError – If dimension is not divisible by 2 or if time tensor is not 1-D.

opentau.policies.pi0.modeling_pi0.make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) → Tensor

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

Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way mask_ar 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.

Returns:

A 2D attention mask tensor of shape (B, N, N).

Raises:

ValueError – If att_masks or pad_masks are not 2D.

opentau.policies.pi0.modeling_pi0.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 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.