opentau.policies.pi0.configuration_pi0

Configuration module for the PI0 Policy.

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

Classes

PI0Config(n_obs_steps, ...)

Configuration class for the PI0 Policy.

class opentau.policies.pi0.configuration_pi0.PI0Config(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, safety_buffer: int = 0, max_state_dim: int = 32, max_action_dim: int = 32, resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, tokenizer_max_length: int = 48, proj_width: int = 1024, dropout: float = 0.1, num_steps: int = 10, advantage_threshold: float = 0.0, advantage: ~typing.Literal['ignore', 'on', 'use'] = 'use', use_cache: bool = True, attention_implementation: str = 'eager', freeze_vision_encoder: bool = True, train_expert_only: bool = False, train_state_proj: 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, use_awr: bool = False, awr_max_weight: float = 10.0)

Bases: PreTrainedConfig

Configuration class for the PI0 Policy.

This class defines the configuration parameters for the PI0 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.

• safety_buffer – Safety buffer size. Defaults to 0.

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

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

• tokenizer_max_length – Maximum length for tokenizer. Defaults to 48.

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

• advantage_threshold – Advantage binarization threshold for AWR. Defaults to 0.0.

• advantage – Advantage conditioning mode. One of “ignore”, “on”, “use”. “use” uses values from dataset, “ignore” disables conditioning, “on” sets advantage to True (for expert demos). Defaults to “use”.

• use_cache – Whether to use KV cache during inference. Defaults to True.

• attention_implementation – Attention implementation to use (“eager”, “sdpa”, or “fa2”). Defaults to “eager”. “sdpa” dispatches to torch.nn.functional.scaled_dot_product_attention and is typically 2-3x faster on A100 + bf16 (FlashAttention-2 backend). “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.

• train_state_proj – Whether to train the state projection layer. Defaults to True.

• 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 = 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, safety_buffer: int = 0, max_state_dim: int = 32, max_action_dim: int = 32, resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, tokenizer_max_length: int = 48, proj_width: int = 1024, dropout: float = 0.1, num_steps: int = 10, advantage_threshold: float = 0.0, advantage: ~typing.Literal['ignore', 'on', 'use'] = 'use', use_cache: bool = True, attention_implementation: str = 'eager', freeze_vision_encoder: bool = True, train_expert_only: bool = False, train_state_proj: 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, use_awr: bool = False, awr_max_weight: float = 10.0) → 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]

advantage: Literal['ignore', 'on', 'use'] = 'use'

advantage_threshold: float = 0.0

attention_implementation: str = 'eager'

awr_max_weight: float = 10.0

chunk_size: int = 50

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

max_action_dim: int = 32

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

proj_width: int = 1024

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

property reward_delta_indices: None

Indices for reward deltas.

Returns:

As reward deltas are not used.

Return type:

None

safety_buffer: int = 0

scheduler_decay_lr: float = 2.5e-06

scheduler_decay_steps: int = 30000

scheduler_warmup_steps: int = 1000

tokenizer_max_length: int = 48

train_expert_only: bool = False

train_state_proj: bool = True

use_awr: bool = False

use_cache: 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.