opentau.policies.pi05_mem.configuration_pi05

Configuration module for the PI05 Mem Policy.

This module defines the PI05MemConfig class, which handles configuration parameters for the PI05 Mem variant. This variant extends the SigLIP image encoder from PaliGemma with space-time separable attention every spacetime_layer_stride-th layer (per the MEM paper’s low-level memory architecture), and processes temporal state sequences (one continuous token per timestep).

Classes

PI05MemConfig(n_obs_steps, ...)

Configuration class for the PI05 Mem Policy.

class opentau.policies.pi05_mem.configuration_pi05.PI05MemConfig(n_obs_steps: int = 8, 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, history_interval: int = 1, max_state_dim: int = 32, max_action_dim: int = 32, resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, prompt_max_length: int = 256, 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, attention_implementation: str = 'eager', rope_type: str = 'mrope_interleaved', freeze_vision_encoder: bool = True, train_expert_only: bool = False, knowledge_insulation: bool = True, gradient_checkpointing: bool = False, spacetime_layer_stride: int = 4, use_motion: bool = False, motion_insert_layer: int | None = None, motion_hidden_dim: int = 256, motion_window: tuple[int, int, int] = (5, 9, 9), motion_corr_func: str = 'cosine', motion_n_encoders: int = 1, motion_norm: str = 'groupnorm', motion_int_mode: str = 'lite', motion_zero_init: bool = True, 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 Mem Policy.

This variant uses PaliGemma’s SigLIP as a space-time video encoder: every spacetime_layer_stride-th ViT layer adds a causal temporal attention over frames (reusing the layer’s existing Q/K/V/O projections — no new learnable parameters). Past-timestep tokens are dropped after the encoder so the prefix matches a single-frame VLA’s 256 image tokens.

Parameters:

• n_obs_steps – Number of temporal video frames the video encoder sees per forward call. During training the dataloader must be configured with dataset_mixture.n_obs_history = n_obs_steps; during inference the observation-history buffer is stacked to produce exactly n_obs_steps frames (sampled at history_interval).

• history_interval – Temporal stride between stacked frames, in environment steps. 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. 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 video frame resizing with padding. Defaults to (224, 224).

• empty_cameras – Number of empty camera inputs to add. 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 (“eager”, “sdpa”, or “fa2”; “fa2” falls back to “eager” with a warning). Defaults to “eager”.

• freeze_vision_encoder – Whether to freeze the SigLIP vision tower. When True the multi_modal_projector remains trainable, matching the semantics in pi05_continuous_state. Defaults to True.

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

• spacetime_layer_stride – Every stride-th SigLIP encoder layer gets the temporal self-attention sublayer added. Defaults to 4, matching the MEM paper. The video encoder introduces no new learnable parameters and shares paligemma.vision_tower / multi_modal_projector with paligemma_with_expert, so any pi05 checkpoint loads directly with unchanged state_dict keys.

__init__(n_obs_steps: int = 8, 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, history_interval: int = 1, max_state_dim: int = 32, max_action_dim: int = 32, resize_imgs_with_padding: tuple[int, int] = (224, 224), empty_cameras: int = 0, prompt_max_length: int = 256, 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, attention_implementation: str = 'eager', rope_type: str = 'mrope_interleaved', freeze_vision_encoder: bool = True, train_expert_only: bool = False, knowledge_insulation: bool = True, gradient_checkpointing: bool = False, spacetime_layer_stride: int = 4, use_motion: bool = False, motion_insert_layer: int | None = None, motion_hidden_dim: int = 256, motion_window: tuple[int, int, int] = (5, 9, 9), motion_corr_func: str = 'cosine', motion_n_encoders: int = 1, motion_norm: str = 'groupnorm', motion_int_mode: str = 'lite', motion_zero_init: bool = True, 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

Returns the default optimizer configuration.

get_scheduler_preset() → LRSchedulerConfig

Returns the default scheduler configuration.

gradient_checkpointing: bool = False

history_interval: int = 1

knowledge_insulation: bool = True

max_action_dim: int = 32

max_delay: int = 0

max_state_dim: int = 32

motion_corr_func: str = 'cosine'

motion_hidden_dim: int = 256

motion_insert_layer: int | None = None

motion_int_mode: str = 'lite'

motion_n_encoders: int = 1

motion_norm: str = 'groupnorm'

motion_window: tuple[int, int, int] = (5, 9, 9)

motion_zero_init: bool = True

n_action_steps: int = 50

n_obs_steps: int = 8

normalization_mapping: dict[str, NormalizationMode]

num_steps: int = 10

property obs_buffer_size: int

Total raw frames the observation buffer must keep.

With n_obs_steps=T and history_interval=k, the buffer stores the most recent (T-1)*k + 1 frames so that T evenly-spaced frames can be selected.

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

proj_width: int = 1024

prompt_max_length: int = 256

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

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.

rope_type: str = 'mrope_interleaved'

scheduler_decay_lr: float = 2.5e-06

scheduler_decay_steps: int = 30000

scheduler_warmup_steps: int = 1000

spacetime_layer_stride: int = 4

train_expert_only: bool = False

use_motion: bool = False

validate_features() → None

Validates the features and adds empty cameras if configured.