opentau.policies.pi05.configuration_pi05

Configuration module for the PI05 Policy.

This module defines the PI05Config class, which handles the configuration parameters for the PI05 Vision-Language-Action Flow Model. It includes settings for the model architecture, optimization, scheduling, and data processing.

Classes

PI05Config(n_obs_steps, ...)	Configuration class for the PI05 Policy.
PI05ContinuousStateConfig(n_obs_steps, ...)	Deprecated: use PI05Config(state_type="continuous") instead.

class opentau.policies.pi05.configuration_pi05.PI05Config(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 = True, 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, state_type: ~typing.Literal['discrete', 'continuous'] = 'discrete', resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, prompt_max_length: int = 256, response_indicator_max_length: int = 3, discrete_action_indicator_max_length: int = 3, response_max_length: int = 52, discrete_action_max_length: int = 32, discrete_action_tokenizer_path: str = 'physical-intelligence/fast', proj_width: int = 1024, dropout: float = 0.1, num_steps: int = 10, max_delay: int = 0, use_modality_embedding: bool = False, 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)

Bases: PreTrainedConfig

Configuration class for the PI05 Policy.

This class defines the configuration parameters for the PI05 model, including input/output structure, model architecture, training settings, and preprocessing.

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. Defaults to identity for visual features and mean-std for state and action.

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

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

• predict_response – Whether to predict the response. Defaults to False.

• resize_imgs_with_padding – Target size (height, width) for image resizing with padding. Defaults to (224, 224).

• empty_cameras – Number of empty camera inputs to add. Used for specific adaptations like Aloha simulation. Defaults to 0.

• prompt_max_length – Maximum length for tokenizer. Defaults to 256.

• discrete_action_max_length – Maximum length for discrete action tokens. Defaults to 32.

• proj_width – Width of the projection layer. Defaults to 1024.

• dropout – Dropout rate. Defaults to 0.1.

• num_steps – Number of flow matching steps for decoding. Defaults to 10.

• attention_implementation – Attention implementation to use (“eager”, “sdpa”, or “fa2”). Defaults to “eager”. “sdpa” dispatches to torch.nn.functional.scaled_dot_product_attention (frees ~5.6 GiB on forward at the bs ceiling tested; see PR #182). “fa2” is accepted for backward compatibility but logs a warning and falls back to “eager”.

• freeze_vision_encoder – Whether to freeze the vision encoder during fine-tuning. Defaults to True.

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

• optimizer_lr – Learning rate for the optimizer. Defaults to 2.5e-5.

• optimizer_betas – Beta parameters for AdamW optimizer. Defaults to (0.9, 0.95).

• optimizer_eps – Epsilon parameter for AdamW optimizer. Defaults to 1e-8.

• optimizer_weight_decay – Weight decay for AdamW optimizer. Defaults to 1e-10.

• scheduler_warmup_steps – Number of warmup steps for the scheduler. Defaults to 1_000.

• scheduler_decay_steps – Number of decay steps for the scheduler. Defaults to 30_000.

• scheduler_decay_lr – Target learning rate after decay. 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 = True, 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, state_type: ~typing.Literal['discrete', 'continuous'] = 'discrete', resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, prompt_max_length: int = 256, response_indicator_max_length: int = 3, discrete_action_indicator_max_length: int = 3, response_max_length: int = 52, discrete_action_max_length: int = 32, discrete_action_tokenizer_path: str = 'physical-intelligence/fast', proj_width: int = 1024, dropout: float = 0.1, num_steps: int = 10, max_delay: int = 0, use_modality_embedding: bool = False, 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]

Indices for action deltas.

Returns:

A list of indices corresponding to the chunk size.

Return type:

list[int]

attention_implementation: str = 'eager'

chunk_size: int = 50

discrete_action_indicator_max_length: int = 3

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

Returns the default optimizer configuration.

Returns:

The optimizer configuration with default parameters.

Return type:

AdamWConfig

get_scheduler_preset() → LRSchedulerConfig

Returns the default scheduler configuration.

Returns:

The scheduler configuration with default parameters.

Return type:

CosineDecayWithWarmupSchedulerConfig

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 = 10

property observation_delta_indices: None

Indices for observation deltas.

Returns:

As observation deltas are not used.

Return type:

None

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 = 1024

prompt_max_length: int = 256

resize_imgs_with_padding: tuple[int, int] = (224, 224)

response_indicator_max_length: int = 3

response_max_length: int = 52

property reward_delta_indices: None

Indices for reward deltas.

Returns:

As reward deltas are not used.

Return type:

None

scheduler_decay_lr: float = 2.5e-06

scheduler_decay_steps: int = 30000

scheduler_warmup_steps: int = 1000

state_type: Literal['discrete', 'continuous'] = 'discrete'

train_expert_only: bool = False

use_modality_embedding: bool = False

use_torch_compile: bool = True

validate_features() → None

Validates the features and adds empty cameras if configured.

This method checks feature configurations and dynamically adds empty camera inputs to self.input_features based on the empty_cameras parameter.

class opentau.policies.pi05.configuration_pi05.PI05ContinuousStateConfig(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 = True, 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, state_type: ~typing.Literal['discrete', 'continuous'] = 'continuous', resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, prompt_max_length: int = 256, response_indicator_max_length: int = 3, discrete_action_indicator_max_length: int = 3, response_max_length: int = 52, discrete_action_max_length: int = 32, discrete_action_tokenizer_path: str = 'physical-intelligence/fast', proj_width: int = 1024, dropout: float = 0.1, num_steps: int = 10, max_delay: int = 0, use_modality_embedding: bool = False, 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)

Bases: PI05Config

Deprecated: use PI05Config(state_type="continuous") instead.

__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 = True, 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, state_type: ~typing.Literal['discrete', 'continuous'] = 'continuous', resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, prompt_max_length: int = 256, response_indicator_max_length: int = 3, discrete_action_indicator_max_length: int = 3, response_max_length: int = 52, discrete_action_max_length: int = 32, discrete_action_tokenizer_path: str = 'physical-intelligence/fast', proj_width: int = 1024, dropout: float = 0.1, num_steps: int = 10, max_delay: int = 0, use_modality_embedding: bool = False, 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

state_type: Literal['discrete', 'continuous'] = 'continuous'