opentau.policies.pi06.modeling_pi06

π06: a Vision-Language-Action model built on Gemma 3 4B.

Relative to π05 this policy swaps PaliGemma-3B for Gemma 3 4B (34 interleaved sliding-window/global layers, SigLIP at 448×448), enlarges the action expert to ~860M parameters so it matches the backbone depth, and halves the default flow-matching denoising schedule to 5 steps.

References

• π0.6 Model Card, Physical Intelligence, 2025-11-17.

• π*0.6: a VLA That Learns From Experience, arXiv:2511.14759.

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 / truncates a ragged list of FAST-tokenized action chunks to a fixed length, returning (B, max_length) arrays.
resize_with_pad(img, width, height[, pad_value])	Resizes an image to fit within target dimensions while maintaining aspect ratio, padding the remainder with pad_value.

Classes

PI06FlowMatching(config[, ...])	π06: Gemma 3 4B backbone + Gemma-v1 action expert + flow matching.
PI06Policy(config[, per_dataset_stats, ...])	Wrapper around PI06FlowMatching for training and inference in OpenTau.

class opentau.policies.pi06.modeling_pi06.PI06FlowMatching(config: PI06Config, discrete_action_vocab_size: int | None = None, language_tokenizer: AutoTokenizer | None = None)

Bases: Module

π06: Gemma 3 4B backbone + Gemma-v1 action expert + flow matching.

┌──────────────────────────────────────────┐ │ actions │ │ ▲ │ │ ┌┴─────┐ │ │ kv cache │Gemma │ │ │ ┌──────────►│Expert│ │ │ │ │ 34L │ │ │ ┌┴─────────┐ │AdaRMS│ │ │ │ │ └▲─────┘ │ │ │Gemma 3 4B│ │ │ │ │(SigLIP │ noise │ │ │448×448) │ │ │ └▲──▲──▲──▲ │ │ │ │ │ └── discrete actions │ │ │ │ └───── robot state │ │ │ └──────── language tokens │ │ └─────────── image(s) │ └──────────────────────────────────────────┘

__init__(config: PI06Config, discrete_action_vocab_size: int | None = None, language_tokenizer: AutoTokenizer | None = None)

Initializes the PI06FlowMatching model.

Parameters:

• config – PI06Config instance.

• discrete_action_vocab_size – FAST tokenizer vocabulary size.

• language_tokenizer – Optional pre-loaded Gemma 3 tokenizer to share with the enclosing PI06Policy. 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 — there is no second copy to fall out of sync.

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

Apply one Euler step of the flow-matching denoiser.

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) → tuple[Tensor, Tensor, Tensor]

Embed the prefix (image + language + optional response + discrete action tokens) and emit the block-pattern attention masks π0.x uses.

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

Embed the noisy action chunk + timestep for the expert’s bidirectional block.

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, real_action_dim: Tensor | None = None, return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss]

Full training forward pass. Returns {“MSE”: …, “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]]]

Autoregressive response-token generation for one 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) → Tensor

Inference: encode prefix once, run num_steps Euler steps.

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

Standard Gaussian noise (float32).

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

π0-style flow-matching time sampler: Beta(1.5, 1.0) on (0.001, 1.000).

class opentau.policies.pi06.modeling_pi06.PI06Policy(config: PI06Config, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)

Bases: PreTrainedPolicy

Wrapper around PI06FlowMatching for training and inference in OpenTau.

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

Initializes the PI06Policy.

Parameters:

• config – PI06Config 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 PI06Config

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

Full training forward pass. Returns {“MSE”: …, “CE”: …}.

When return_per_sample is True, also returns 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

Load a pretrained π0.6 checkpoint.

Mirrors PI05Policy.from_pretrained — the only π0.6 specific logic is in _fix_pytorch_state_dict_keys, which tolerates both a native π0.6 checkpoint layout and weights migrated from π0.5.

get_optim_params() → dict

Return all parameters for the optimizer.

name: None = 'pi06'

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

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

Action-chunk prediction API (not used for π0.6).

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

Tokenize continuous actions with the FAST processor and pad to discrete_action_max_length.

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

Discretize each state dim into 256 bins and format as a space-joined string, matching the π0.5 / π0.6 “State:” prompt template.

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

Resize (with padding) each camera view to config.resize_imgs_with_padding (π0.6 default: 448×448) and normalize from [0, 1] to [-1, 1] as SigLIP expects. Missing views are replaced by all--1 tensors up to empty_cameras.

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

Tokenize the task prompt together with the discretized state string.

Format matches π0.5:

“Task: {task}<eos>State: {state}<eos>Response:” if predict_response “Task: {task}<eos>State: {state}<eos>Actions:” otherwise

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

Tokenize the response field for supervised co-training.

Returns (None, None) when response prediction is disabled.

reset() → None

Clears the rolling action queue; call on every environment reset.

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

Sample an action chunk.

The provided action_prefix must be unnormalized — this method normalizes internally before feeding it to the flow-matching module.

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

Return a single action from the rolling queue, refilling when needed.

Matches pi05 semantics — only safe for simulation loops. Real-robot inference should pipeline sample_actions in the ROS node directly.

opentau.policies.pi06.modeling_pi06.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 (B, action_chunk_length).

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

• min_period – Minimum period of the sinusoidal functions.

• max_period – Maximum period of the sinusoidal functions.

• device – The device to create the tensors on.

Returns:

Tensor of shape (B, action_chunk_length, dimension) with the embedding.

opentau.policies.pi06.modeling_pi06.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.

Tokens can attend to valid tokens whose cumulative att_masks is smaller or equal to theirs. Block semantics match π0.5 exactly — see the pi05 docstring for the full table of examples, condensed:

[[1 1 1 1 1 1]]: pure causal [[0 0 0 1 1 1]]: prefix-LM (first 3 bidirectional, last 3 causal) [[1 0 1 0 1 0 0 1 0 0]]: multi-block causal

Parameters:

• pad_masks – bool (B, N) — True for real tokens.

• att_masks – int32 (B, N) — 1 starts a new block, 0 continues.

• n_cross_att_tokens – If set, prepend a cross-attention mask of that width so suffix tokens can attend back into the cached prefix.

• cross_att_pad_masks – Prefix pad mask used when building the cross attention block.

Returns:

Boolean 2-D attention mask of shape (B, N, N) or (B, N, N + n_cross_att_tokens) when cross attention is requested.

opentau.policies.pi06.modeling_pi06.pad_discrete_tokens(tokens: list[list[int]], max_length: int) → tuple[ndarray, ndarray]

Pads / truncates a ragged list of FAST-tokenized action chunks to a fixed length, returning (B, max_length) arrays.

opentau.policies.pi06.modeling_pi06.resize_with_pad(img: Tensor, width: int, height: int, pad_value: int = -1) → Tensor

Resizes an image to fit within target dimensions while maintaining aspect ratio, padding the remainder with pad_value.

Parameters:

• img – (B, C, H_src, W_src) tensor.

• width – Target width.

• height – Target height.

• pad_value – Padding value (defaults to -1 to match SigLIP’s [-1, 1] range).

Returns:

Resized/padded tensor of shape (B, C, height, width).