opentau.policies.pi06.configuration_pi06

Configuration module for the PI06 Policy.

This module defines the PI06Config class, which handles the configuration parameters for the PI06 Vision-Language-Action model. π06 inherits the π05 training recipe (FAST discrete action co-training, flow-matching continuous actions, Knowledge Insulation gradient-stop) but upgrades the backbone to Gemma 3 4B with 448×448 vision and a larger action expert.

Classes

PI06Config(n_obs_steps, ...)

Configuration class for the PI06 Policy.
class opentau.policies.pi06.configuration_pi06.PI06Config(n_obs_steps: int = 1, normalization_mapping: dict[str, ~opentau.configs.types.NormalizationMode] = <factory>, input_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, output_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, device: str | None = None, use_amp: bool = False, use_torch_compile: bool = False, torch_compile_mode: str = 'default', pretrained_path: str | None = None, skip_normalization_weights: bool = False, skip_input_resolution_check: bool = False, save_normalization_stats: bool = True, dataset_names: list[str] | None = None, dataset_to_norm_index: dict[str, int] | None = None, cloud_vlm_latency_mean: float = 0.0, cloud_vlm_latency_std: float = 0.0, cloud_vlm_latency_lower: float = 0.0, cloud_vlm_latency_upper: float = 0.0, action_decoder_latency_mean: float = 0.0, action_decoder_latency_std: float = 0.0, action_decoder_latency_lower: float = 0.0, action_decoder_latency_upper: float = 0.0, chunk_size: int = 50, n_action_steps: int = 50, max_state_dim: int = 32, max_action_dim: int = 32, predict_response: bool = False, resize_imgs_with_padding: tuple[int, int] = (448, 448), empty_cameras: int = 0, prompt_max_length: int = 256, response_max_length: int = 52, discrete_action_max_length: int = 32, discrete_action_tokenizer_path: str = 'physical-intelligence/fast', proj_width: int = 1280, dropout: float = 0.1, num_steps: int = 5, max_delay: int = 0, attention_implementation: str = 'eager', freeze_vision_encoder: bool = True, train_expert_only: bool = False, knowledge_insulation: bool = True, gradient_checkpointing: bool = False, optimizer_lr: float = 2.5e-05, optimizer_betas: tuple[float, float] = (0.9, 0.95), optimizer_eps: float = 1e-08, optimizer_weight_decay: float = 1e-10, scheduler_warmup_steps: int = 1000, scheduler_decay_steps: int = 30000, scheduler_decay_lr: float = 2.5e-06)[source]

Bases: PreTrainedConfig

Configuration class for the PI06 Policy.

Mirrors PI05Config but flips the defaults to match π0.6’s architecture: Gemma 3 4B backbone, 448×448 image input, ~860M action expert, and 5 flow-matching denoising steps.

Parameters:

  • n_obs_steps – Number of observation steps to use. Defaults to 1.

  • chunk_size – Trained action-chunk length, i.e. the prediction horizon the model always decodes at inference. Upper bound for n_action_steps. Defaults to 50.

  • n_action_steps – Inference execution horizon – how many actions from each predicted chunk are executed before the policy re-queries with fresh observations. Must be <= chunk_size. Defaults to 50.

  • normalization_mapping – Mapping of feature names to normalization modes.

  • max_state_dim – Maximum dimension for state vectors. Defaults to 32.

  • max_action_dim – Maximum dimension for action vectors. Defaults to 32.

  • predict_response – Whether to predict the response. Defaults to False. Enabling this is required to reproduce the paper’s hierarchical design (π0.6 model card §1: “preserves the hierarchical design of π0.5, providing high-level subtask prediction and low-level action generation”). When False, π0.6 reduces to a flat low-level-only model. The default is False because most LeRobot-style datasets do not carry subtask annotations. The same field is also used to supervise VQA / grounding-style textual targets during co-training (this is why the field is named “response” rather than “subtask” — it covers both uses, matching the π0.5 pretraining recipe).

  • resize_imgs_with_padding – Target image size. Defaults to (448, 448).

  • empty_cameras – Number of empty camera inputs to add. Defaults to 0. π0.6 pre-training uses up to 4 cameras (base + 2 wrist + optional backward for mobile manipulators); set this to match your robot.

  • prompt_max_length – Maximum tokenizer length. Defaults to 256.

  • response_max_length – Maximum response length. Defaults to 52.

  • discrete_action_max_length – Maximum discrete action token length. Defaults to 32.

  • proj_width – Width of the action projection layer. Defaults to 1280 to match the Gemma-v1 action expert hidden size.

  • dropout – Dropout rate. Defaults to 0.1.

  • num_steps – Number of flow matching denoising steps. Defaults to 5 (halved from π0.5’s 10, giving ~63 ms per chunk on an H100).

  • attention_implementation – “eager”, “sdpa”, or “fa2”. Defaults to “eager”. “sdpa” dispatches to torch.nn.functional.scaled_dot_product_attention (typically 2-3x faster on A100 + bf16). “fa2” is accepted for backward compatibility but logs a warning and falls back to eager.

  • freeze_vision_encoder – Whether to freeze the vision encoder. Defaults to True.

  • train_expert_only – Whether to train only the expert module. Defaults to False.

  • optimizer_lr – AdamW learning rate. Defaults to 2.5e-5.

  • optimizer_betas – AdamW betas. Defaults to (0.9, 0.95).

  • optimizer_eps – AdamW epsilon. Defaults to 1e-8.

  • optimizer_weight_decay – AdamW weight decay. Defaults to 1e-10.

  • scheduler_warmup_steps – Warmup steps. Defaults to 1_000.

  • scheduler_decay_steps – Decay steps. Defaults to 30_000.

  • scheduler_decay_lr – Target decay learning rate. Defaults to 2.5e-6.

__init__(n_obs_steps: int = 1, normalization_mapping: dict[str, ~opentau.configs.types.NormalizationMode] = <factory>, input_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, output_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, device: str | None = None, use_amp: bool = False, use_torch_compile: bool = False, torch_compile_mode: str = 'default', pretrained_path: str | None = None, skip_normalization_weights: bool = False, skip_input_resolution_check: bool = False, save_normalization_stats: bool = True, dataset_names: list[str] | None = None, dataset_to_norm_index: dict[str, int] | None = None, cloud_vlm_latency_mean: float = 0.0, cloud_vlm_latency_std: float = 0.0, cloud_vlm_latency_lower: float = 0.0, cloud_vlm_latency_upper: float = 0.0, action_decoder_latency_mean: float = 0.0, action_decoder_latency_std: float = 0.0, action_decoder_latency_lower: float = 0.0, action_decoder_latency_upper: float = 0.0, chunk_size: int = 50, n_action_steps: int = 50, max_state_dim: int = 32, max_action_dim: int = 32, predict_response: bool = False, resize_imgs_with_padding: tuple[int, int] = (448, 448), empty_cameras: int = 0, prompt_max_length: int = 256, response_max_length: int = 52, discrete_action_max_length: int = 32, discrete_action_tokenizer_path: str = 'physical-intelligence/fast', proj_width: int = 1280, dropout: float = 0.1, num_steps: int = 5, max_delay: int = 0, attention_implementation: str = 'eager', freeze_vision_encoder: bool = True, train_expert_only: bool = False, knowledge_insulation: bool = True, gradient_checkpointing: bool = False, optimizer_lr: float = 2.5e-05, optimizer_betas: tuple[float, float] = (0.9, 0.95), optimizer_eps: float = 1e-08, optimizer_weight_decay: float = 1e-10, scheduler_warmup_steps: int = 1000, scheduler_decay_steps: int = 30000, scheduler_decay_lr: float = 2.5e-06) → None
property action_delta_indices: list[int]

Get indices for action delta features.

Returns:

List of indices indicating which action features should be treated as deltas, or None if no delta features are used.

attention_implementation: str = 'eager'
chunk_size: int = 50
discrete_action_max_length: int = 32
discrete_action_tokenizer_path: str = 'physical-intelligence/fast'
dropout: float = 0.1
empty_cameras: int = 0
freeze_vision_encoder: bool = True
get_optimizer_preset() → AdamWConfig[source]

Returns the default optimizer configuration.

get_scheduler_preset() → LRSchedulerConfig[source]

Returns the default scheduler configuration.

gradient_checkpointing: bool = False
knowledge_insulation: bool = True
max_action_dim: int = 32
max_delay: int = 0
max_state_dim: int = 32
n_action_steps: int = 50
n_obs_steps: int = 1
normalization_mapping: dict[str, NormalizationMode]
num_steps: int = 5
property observation_delta_indices: None

Get indices for observation delta features.

Returns:

List of indices indicating which observation features should be treated as deltas, or None if no delta features are used.

optimizer_betas: tuple[float, float] = (0.9, 0.95)
optimizer_eps: float = 1e-08
optimizer_lr: float = 2.5e-05
optimizer_weight_decay: float = 1e-10
predict_response: bool = False
proj_width: int = 1280
prompt_max_length: int = 256
resize_imgs_with_padding: tuple[int, int] = (448, 448)
response_max_length: int = 52
property reward_delta_indices: None

Get indices for reward delta features.

Returns:

List of indices indicating which reward features should be treated as deltas, or None if no delta features are used.

scheduler_decay_lr: float = 2.5e-06
scheduler_decay_steps: int = 30000
scheduler_warmup_steps: int = 1000
train_expert_only: bool = False
validate_features() → None[source]

Validates the features and adds empty cameras if configured.