# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""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.
"""
from dataclasses import dataclass, field
from opentau.configs.policies import PreTrainedConfig
from opentau.configs.types import FeatureType, NormalizationMode, PolicyFeature
from opentau.optim.optimizers import AdamWConfig
from opentau.optim.schedulers import (
CosineDecayWithWarmupSchedulerConfig,
LRSchedulerConfig,
)
[docs]
@PreTrainedConfig.register_subclass("pi07_paligemma_low_level")
@dataclass
class PI07PaligemmaLowLevelConfig(PreTrainedConfig):
"""Configuration class for the pi07_paligemma_low_level policy.
This class defines the configuration parameters for the PI07PaligemmaLowLevel model, including
input/output structure, model architecture, training settings, and preprocessing.
Args:
n_obs_steps: Number of observation steps to use. Defaults to 6.
history_interval: Temporal stride between the ``n_obs_steps`` stacked
frames in the inference buffer, 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.
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.
Must match the resolution of the input images (``TrainPipelineConfig.resolution`` /
the bound image features) — training fails fast on a mismatch, since the policy
would silently letterbox every frame a second time
(``skip_input_resolution_check=true`` downgrades that to a warning for legacy
checkpoints). ``null`` passes frames through at the input resolution. Non-224
values are supported natively: frames are padded up to the next patch multiple
and the SigLIP position embeddings are interpolated to the resulting grid, so no
pixels are cropped and no letterboxing to 224x224 occurs. 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.
per_group_projection: When True, use an independent state/action
projection per (robot_type, control_mode) group (shares the
normalization-head grouping). Defaults to False.
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", "fa2", or
"flash_cuda"). 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". "flash_cuda" uses the custom block-causal CUDA flash
kernel (``opentau.policies.flash_attn_cuda``, Tensor Core forward + backward): the SxS
attention mask is reconstructed in-kernel from compact block-ids so it is never
materialized. There is no fallback — if the kernel cannot be JIT-compiled (no
CUDA/compiler) the forward raises rather than silently using sdpa/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.
"""
# Input / output structure.
n_obs_steps: int = 6
chunk_size: int = 50
n_action_steps: int = 50
# Inference observation-history buffer: ``history_interval`` is the
# temporal stride between the ``n_obs_steps`` stacked frames. Together they
# determine ``obs_buffer_size = (n_obs_steps - 1) * history_interval + 1``.
history_interval: int = 1
normalization_mapping: dict[str, NormalizationMode] = field(
default_factory=lambda: {
"VISUAL": NormalizationMode.IDENTITY,
"STATE": NormalizationMode.MEAN_STD,
"ACTION": NormalizationMode.MEAN_STD,
}
)
# Shorter state and action vectors will be padded
max_state_dim: int = 32
max_action_dim: int = 32
# Image preprocessing
resize_imgs_with_padding: tuple[int, int] = (224, 224)
# Add empty images. Used by pi05_aloha_sim which adds the empty
# left and right wrist cameras in addition to the top camera.
empty_cameras: int = 0
# Every spacetime_layer_stride-th SigLIP encoder layer gets a temporal
# attention sub-layer when n_obs_steps > 1. When n_obs_steps == 1 the
# encoder is byte-identical to plain SigLIP (no temporal attention).
spacetime_layer_stride: int = 4
# Language Tokenizer
prompt_max_length: int = 256
# Maximum token length for subtask response from the high-level planner
response_max_length: int = 52
# Maximum token length for optional episode metadata (speed / quality / mistake)
metadata_max_length: int = 52
# Maximum length of the action tokens
discrete_action_max_length: int = 32
# HF repo id or local path for the FAST action tokenizer
# (``AutoProcessor.from_pretrained(..., trust_remote_code=True)``).
# Override to use a tokenizer specialized to your mixture (see
# ``opentau.scripts.fit_fast_tokenizer``).
discrete_action_tokenizer_path: str = "physical-intelligence/fast"
# EXPERIMENTAL eval-only path: when True, ``sample_actions`` ignores the
# flow-matching action expert and instead autoregressively generates the
# FAST discrete-action tokens with the VLM backbone (greedy argmax over
# ``da_head`` logits, same "Action: " prefix layout as training), decodes
# them to an action chunk, and inverts the MIN_MAX discrete-action
# normalization. Generation stops when the BPE expansion of the generated
# tokens reaches chunk_size * max_action_dim DCT coefficients (FAST
# sequences are self-delimiting given the (T, D) shape — there is no
# learned EOS because training masks the CE loss at pad positions), capped
# at ``discrete_action_max_length`` tokens (the same cap training applies).
eval_use_discrete_actions: bool = False
# Projector
proj_width: int = 1024
# When True, give the continuous STATE/ACTION projections (state_proj,
# action_in_proj, action_out_proj) an independent (weight, bias) per
# (robot_type, control_mode) group — the same grouping the stacked
# Normalize/Unnormalize heads use, selected per-sample by
# `_resolve_dataset_index`. The number of groups is len(dataset_names)
# (falls back to 1 when dataset_names is unset). time_mlp_in/out stay
# shared. Defaults to False: when off the three projections are plain
# nn.Linear and the model is byte-identical to the pre-flag version (same
# params, same state_dict keys, same numerics). Flipping it on changes the
# param count and the projection state_dict shapes (leading group axis);
# legacy single-group checkpoints are promoted + duplicated on load, and a
# group new to the checkpoint gets a fresh row copied from row 0. Loading a
# per-group checkpoint with this flag off raises rather than silently
# collapsing the group axis.
per_group_projection: bool = False
# Dropout
dropout: float = 0.1
# Decoding
num_steps: int = 10
# Real Time Inference
# maximum number of frozen actions
max_delay: int = 0
# Attention utils
attention_implementation: str = "eager"
# Finetuning settings
freeze_vision_encoder: bool = True
train_expert_only: bool = False
# Knowledge insulation (π0.5): when True (default), the prefix/VLM KV cache
# is detached before the action expert reads it, so the flow-matching action
# loss does NOT backpropagate into the VLM backbone. Set False to let the
# action gradient flow into the VLM (end-to-end action training). Default
# True preserves existing behavior and is what current checkpoints expect.
knowledge_insulation: bool = True
# Wrap each transformer-layer forward in torch.utils.checkpoint to trade
# ~25-33%% same-batch compute for ~30-40 GB of activation memory per rank,
# typically netting +10-25%% throughput once the freed memory is spent on
# a larger per-rank batch. Only supported with distributed_type=MULTI_GPU
# (DDP), NO (single process), or DeepSpeed ZeRO-1/2 — src/opentau/scripts/
# train.py raises if the accelerator's distributed_type is anything else
# (ZeRO-3, FSDP) because pi05's custom per-layer forward does not wire up
# the backend-specific activation-checkpointing hooks those strategies
# require. Defaults to False (no ckpt, lowest risk).
gradient_checkpointing: bool = False
# Training presets
optimizer_lr: float = 2.5e-5
optimizer_betas: tuple[float, float] = (0.9, 0.95)
optimizer_eps: float = 1e-8
optimizer_weight_decay: float = 1e-10
scheduler_warmup_steps: int = 1_000
scheduler_decay_steps: int = 30_000
scheduler_decay_lr: float = 2.5e-6
# Debug aid: during the training forward, emit a logging.warning when any
# individual NORMALIZED state/action feature dim (after taking abs) exceeds
# this value — a symptom of bad normalization stats (e.g. near-zero std on a
# constant dim) or corrupt data. The warning names the offending
# source/episode/frame when those batch fields are present, to point at the
# dataset/frame to inspect. Set to 0.0 (or negative) to disable entirely and
# skip the per-step device sync. The default 32 is deliberately permissive: the failure it
# targets drives normalized values to ~1e8, so 32 flags those by a wide margin while
# tolerating benign large-but-finite dims (the zero-variance guard and pad-aware masking
# already remove the common false positives).
warn_outlier_threshold: float = 32.0
@property
def obs_buffer_size(self) -> int:
"""Total raw frames the state history 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.
"""
if self.n_obs_steps <= 1:
return 1
return (self.n_obs_steps - 1) * self.history_interval + 1
def __post_init__(self):
"""Post-initialization validation."""
super().__post_init__()
# TODO(Steven): Validate device and amp? in all policy configs?
"""Input validation (not exhaustive)."""
if not isinstance(self.n_obs_steps, int) or self.n_obs_steps < 1:
raise ValueError(f"`n_obs_steps` must be a positive integer, got {self.n_obs_steps}.")
if not isinstance(self.history_interval, int) or self.history_interval < 1:
raise ValueError(f"`history_interval` must be a positive integer, got {self.history_interval}.")
if self.n_action_steps > self.chunk_size:
raise ValueError(
f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
f"{self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
)
if self.max_delay > self.chunk_size:
raise ValueError(
f"The max delay must be less than or equal to the chunk size. Got {self.max_delay} for `max_delay` and {self.chunk_size} for `chunk_size`."
)
if self.n_action_steps < self.chunk_size and self.max_delay != 0:
raise ValueError(
"A shortened execution horizon (n_action_steps < chunk_size) is not yet "
"supported together with real-time inference delay (max_delay > 0); they "
"would entangle the action-queue prefix logic. Got "
f"n_action_steps={self.n_action_steps}, chunk_size={self.chunk_size}, "
f"max_delay={self.max_delay}."
)
[docs]
def validate_features(self) -> 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.
"""
for i in range(self.empty_cameras):
key = f"observation.images.empty_camera_{i}"
empty_camera = PolicyFeature(
type=FeatureType.VISUAL,
shape=(3, 480, 640),
)
self.input_features[key] = empty_camera
[docs]
def get_optimizer_preset(self) -> AdamWConfig:
"""Returns the default optimizer configuration.
Returns:
AdamWConfig: The optimizer configuration with default parameters.
"""
return AdamWConfig(
lr=self.optimizer_lr,
betas=self.optimizer_betas,
eps=self.optimizer_eps,
weight_decay=self.optimizer_weight_decay,
)
[docs]
def get_scheduler_preset(self) -> LRSchedulerConfig:
"""Returns the default scheduler configuration.
Returns:
CosineDecayWithWarmupSchedulerConfig: The scheduler configuration with default parameters.
"""
return CosineDecayWithWarmupSchedulerConfig(
peak_lr=self.optimizer_lr,
decay_lr=self.scheduler_decay_lr,
num_warmup_steps=self.scheduler_warmup_steps,
num_decay_steps=self.scheduler_decay_steps,
)
@property
def observation_delta_indices(self) -> None:
"""Indices for observation deltas.
Returns:
None: As observation deltas are not used.
"""
return None
@property
def action_delta_indices(self) -> list[int]:
"""Indices for action deltas.
Returns:
list[int]: A list of indices corresponding to the chunk size.
"""
return list(range(self.chunk_size))
@property
def reward_delta_indices(self) -> None:
"""Indices for reward deltas.
Returns:
None: As reward deltas are not used.
"""
return None