opentau 0.1.0__py3-none-any.whl

opentau/policies/normalize.py

@@ -0,0 +1,315 @@
+#!/usr/bin/env python
+# 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.
+
+"""Normalization and Unnormalization utilities for policies.
+
+This module provides classes and functions to normalize and unnormalize data
+(e.g., observations and actions) based on statistical properties (mean, std, min, max).
+It handles different normalization modes and supports creating buffers for statistics.
+"""
+
+import sys
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+from opentau.configs.types import FeatureType, NormalizationMode, PolicyFeature
+
+EPS = 1e-8  # Small epsilon value for numerical stability in normalization
+
+
+def warn_missing_keys(features: dict[str, PolicyFeature], batch: dict[str, Tensor], mode: str) -> None:
+    """Warns if expected features are missing from the batch.
+
+    Args:
+        features: Dictionary of expected policy features.
+        batch: Dictionary containing the data batch.
+        mode: The operation mode (e.g., "normalization" or "unnormalization") for the warning message.
+    """
+    for missing_key in set(features) - set(batch):
+        red_seq = "\033[91m"
+        reset_seq = "\033[0m"
+        print(
+            f"{red_seq}Warning: {missing_key} was missing from the batch during {mode}.{reset_seq}",
+            file=sys.stderr,
+        )
+
+
+def create_stats_buffers(
+    features: dict[str, PolicyFeature],
+    norm_map: dict[str, NormalizationMode],
+    stats: dict[str, dict[str, Tensor]] | None = None,
+) -> dict[str, dict[str, nn.ParameterDict]]:
+    """
+    Create buffers per modality (e.g. "observation.image", "action") containing their mean, std, min, max
+    statistics.
+
+    Args:
+        features: Dictionary mapping feature names to PolicyFeature objects.
+        norm_map: Dictionary mapping feature types to NormalizationMode.
+        stats: Optional dictionary containing pre-computed statistics (mean, std, min, max)
+            for each feature. If None, buffers are initialized with infinity.
+
+    Returns:
+        dict: A dictionary where keys are modalities and values are `nn.ParameterDict` containing
+            `nn.Parameters` set to `requires_grad=False`, suitable to not be updated during backpropagation.
+
+    Raises:
+        ValueError: If stats contain types other than np.ndarray or torch.Tensor.
+    """
+    stats_buffers = {}
+
+    for key, ft in features.items():
+        norm_mode = norm_map.get(ft.type, NormalizationMode.IDENTITY)
+        if norm_mode is NormalizationMode.IDENTITY:
+            continue
+
+        assert isinstance(norm_mode, NormalizationMode)
+
+        shape = tuple(ft.shape)
+
+        if ft.type is FeatureType.VISUAL:
+            # sanity checks
+            assert len(shape) == 3, f"number of dimensions of {key} != 3 ({shape=}"
+            c, h, w = shape
+            assert c < h and c < w, f"{key} is not channel first ({shape=})"
+            # override image shape to be invariant to height and width
+            shape = (c, 1, 1)
+
+        # Note: we initialize mean, std, min, max to infinity. They should be overwritten
+        # downstream by `stats` or `policy.load_state_dict`, as expected. During forward,
+        # we assert they are not infinity anymore.
+
+        buffer = {}
+        if norm_mode is NormalizationMode.MEAN_STD:
+            mean = torch.ones(shape, dtype=torch.float32) * torch.inf
+            std = torch.ones(shape, dtype=torch.float32) * torch.inf
+            buffer = nn.ParameterDict(
+                {
+                    "mean": nn.Parameter(mean, requires_grad=False),
+                    "std": nn.Parameter(std, requires_grad=False),
+                }
+            )
+        elif norm_mode is NormalizationMode.MIN_MAX:
+            min = torch.ones(shape, dtype=torch.float32) * torch.inf
+            max = torch.ones(shape, dtype=torch.float32) * torch.inf
+            buffer = nn.ParameterDict(
+                {
+                    "min": nn.Parameter(min, requires_grad=False),
+                    "max": nn.Parameter(max, requires_grad=False),
+                }
+            )
+
+        # TODO(aliberts, rcadene): harmonize this to only use one framework (np or torch)
+        if stats:
+            if isinstance(stats[key]["mean"], np.ndarray):
+                if norm_mode is NormalizationMode.MEAN_STD:
+                    buffer["mean"].data = torch.from_numpy(stats[key]["mean"]).to(dtype=torch.float32)
+                    buffer["std"].data = torch.from_numpy(stats[key]["std"]).to(dtype=torch.float32)
+                elif norm_mode is NormalizationMode.MIN_MAX:
+                    buffer["min"].data = torch.from_numpy(stats[key]["min"]).to(dtype=torch.float32)
+                    buffer["max"].data = torch.from_numpy(stats[key]["max"]).to(dtype=torch.float32)
+            elif isinstance(stats[key]["mean"], torch.Tensor):
+                # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
+                # tensors anywhere (for example, when we use the same stats for normalization and
+                # unnormalization). See the logic here
+                # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
+                if norm_mode is NormalizationMode.MEAN_STD:
+                    buffer["mean"].data = stats[key]["mean"].clone().to(dtype=torch.float32)
+                    buffer["std"].data = stats[key]["std"].clone().to(dtype=torch.float32)
+                elif norm_mode is NormalizationMode.MIN_MAX:
+                    buffer["min"].data = stats[key]["min"].clone().to(dtype=torch.float32)
+                    buffer["max"].data = stats[key]["max"].clone().to(dtype=torch.float32)
+            else:
+                type_ = type(stats[key]["mean"])
+                raise ValueError(f"np.ndarray or torch.Tensor expected, but type is '{type_}' instead.")
+
+        stats_buffers[key] = buffer
+    return stats_buffers
+
+
+def _no_stats_error_str(name: str) -> str:
+    """Returns an error message string for missing statistics.
+
+    Args:
+        name: Name of the statistic (e.g., "mean", "std").
+
+    Returns:
+        str: The error message string.
+    """
+    return (
+        f"`{name}` is infinity. You should either initialize with `stats` as an argument, or use a "
+        "pretrained model."
+    )
+
+
+class Normalize(nn.Module):
+    """Normalizes data (e.g. "observation.image") for more stable and faster convergence during training."""
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        """Initializes the Normalize module.
+
+        Args:
+            features: A dictionary where keys are input modalities (e.g. "observation.image") and values
+                are their PolicyFeature definitions.
+            norm_map: A dictionary where keys are feature types and values are their normalization modes.
+            stats: A dictionary where keys are output modalities (e.g. "observation.image")
+                and values are dictionaries of statistic types and their values (e.g.
+                `{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
+                training the model for the first time, these statistics will overwrite the default buffers. If
+                not provided, as expected for finetuning or evaluation, the default buffers should to be
+                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
+                dataset is not needed to get the stats, since they are already in the policy state_dict.
+        """
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+        self.stats = stats
+        stats_buffers = create_stats_buffers(features, norm_map, stats)
+        for key, buffer in stats_buffers.items():
+            setattr(self, "buffer_" + key.replace(".", "_"), buffer)
+
+    # TODO(rcadene): should we remove torch.no_grad?
+    @torch.no_grad
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Normalizes the batch data.
+
+        Args:
+            batch: Dictionary containing the data to normalize.
+
+        Returns:
+            dict[str, Tensor]: The normalized batch data.
+
+        Raises:
+            ValueError: If an unknown normalization mode is encountered.
+        """
+        warn_missing_keys(self.features, batch, "normalization")
+        batch = dict(batch)  # shallow copy avoids mutating the input batch
+        for key, ft in self.features.items():
+            if key not in batch:
+                # FIXME(aliberts, rcadene): This might lead to silent fail!
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+
+            if batch[key].numel() == 0:  # skip empty tensors, which won't broadcast well
+                continue
+
+            buffer = getattr(self, "buffer_" + key.replace(".", "_"))
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = buffer["mean"]
+                std = buffer["std"]
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                batch[key] = (batch[key] - mean) / (std + EPS)
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                min = buffer["min"]
+                max = buffer["max"]
+                assert not torch.isinf(min).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max).any(), _no_stats_error_str("max")
+                batch[key] = (batch[key] - min) / (max - min + EPS)
+                # normalize to [-1, 1]
+                batch[key] = batch[key] * 2 - 1
+            else:
+                raise ValueError(norm_mode)
+        return batch
+
+
+class Unnormalize(nn.Module):
+    """
+    Similar to `Normalize` but unnormalizes output data (e.g. `{"action": torch.randn(b,c)}`) in their
+    original range used by the environment.
+    """
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        """Initializes the Unnormalize module.
+
+        Args:
+            features: A dictionary where keys are input modalities (e.g. "observation.image") and values
+                are their PolicyFeature definitions.
+            norm_map: A dictionary where keys are feature types and values are their normalization modes.
+            stats: A dictionary where keys are output modalities (e.g. "observation.image")
+                and values are dictionaries of statistic types and their values. If provided,
+                these statistics will overwrite the default buffers.
+        """
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+        self.stats = stats
+        # `self.buffer_observation_state["mean"]` contains `torch.tensor(state_dim)`
+        stats_buffers = create_stats_buffers(features, norm_map, stats)
+        for key, buffer in stats_buffers.items():
+            setattr(self, "buffer_" + key.replace(".", "_"), buffer)
+
+    # TODO(rcadene): should we remove torch.no_grad?
+    @torch.no_grad
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Unnormalizes the batch data.
+
+        Args:
+            batch: Dictionary containing the data to unnormalize.
+
+        Returns:
+            dict[str, Tensor]: The unnormalized batch data.
+
+        Raises:
+            ValueError: If an unknown normalization mode is encountered.
+        """
+        warn_missing_keys(self.features, batch, "unnormalization")
+        batch = dict(batch)  # shallow copy avoids mutating the input batch
+        for key, ft in self.features.items():
+            if key not in batch:
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+
+            if batch[key].numel() == 0:  # skip empty tensors, which won't broadcast well
+                continue
+
+            buffer = getattr(self, "buffer_" + key.replace(".", "_"))
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = buffer["mean"]
+                std = buffer["std"]
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                batch[key] = batch[key] * (std + EPS) + mean
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                min = buffer["min"]
+                max = buffer["max"]
+                assert not torch.isinf(min).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max).any(), _no_stats_error_str("max")
+                batch[key] = (batch[key] + 1) / 2
+                batch[key] = batch[key] * (max - min + EPS) + min
+            else:
+                raise ValueError(norm_mode)
+        return batch

opentau/policies/pi0/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+"""PI0 Policy Module.
+
+This module implements the π0 (Pi0) Vision-Language-Action Flow Model policy,
+designed for general robot control. It includes the policy definition,
+configuration, and model architecture.
+"""

opentau/policies/pi0/configuration_pi0.py

@@ -0,0 +1,250 @@
+# 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 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.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Literal
+
+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,
+)
+
+
+@PreTrainedConfig.register_subclass("pi0")
+@dataclass
+class PI0Config(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.
+
+    Args:
+        n_obs_steps: Number of observation steps to use. Defaults to 1.
+        chunk_size: Size of the action chunk. The upper bound for n_action_steps. Defaults to 50.
+        n_action_steps: Number of action steps to predict. 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".
+        init_strategy: Initialization strategy. One of "no_init", "full_he_init", "expert_only_he_init".
+            Defaults to "full_he_init".
+        use_cache: Whether to use KV cache during inference. Defaults to True.
+        attention_implementation: Attention implementation to use ("eager" or "fa2"). Defaults 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.
+    """
+
+    # Input / output structure.
+    n_obs_steps: int = 1
+    chunk_size: int = 50
+    n_action_steps: int = 50
+    safety_buffer: int = 0
+
+    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 pi0_aloha_sim which adds the empty
+    # left and right wrist cameras in addition to the top camera.
+    empty_cameras: int = 0
+
+    # Tokenizer
+    tokenizer_max_length: int = 48
+
+    # Projector
+    proj_width: int = 1024
+
+    # Dropout
+    dropout: float = 0.1
+
+    # Decoding
+    num_steps: int = 10
+
+    # advantage binarization threshold for AWR
+    advantage_threshold: float = 0.0
+
+    # When set to "use", the advantage values provided in the dataset will be used.
+    # When set to "ignore", no advantage conditioning will be applied.
+    # When set to "on", the advantage will always be True.
+    # This should only be "on" when training on expert demonstrations or interventions.
+    advantage: Literal["ignore", "on", "use"] = "use"
+
+    # Initialization strategy
+    init_strategy: Literal["no_init", "full_he_init", "expert_only_he_init"] = "full_he_init"
+
+    # Attention utils
+    use_cache: bool = True
+    attention_implementation: str = "eager"  # or fa2
+
+    # Finetuning settings
+    freeze_vision_encoder: bool = True
+    train_expert_only: bool = False
+    train_state_proj: bool = True
+
+    # 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
+
+    # TODO: Add EMA
+
+    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 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.n_obs_steps != 1:
+            raise ValueError(
+                f"Multiple observation steps not handled yet. Got `nobs_steps={self.n_obs_steps}`"
+            )
+
+        assert self.init_strategy in ["no_init", "full_he_init", "expert_only_he_init"], (
+            f"Invalid init strategy: {self.init_strategy} must be one of ['no_init', 'full_he_init', 'expert_only_he_init']"
+        )
+
+        if self.init_strategy == "expert_only_he_init" and self.pretrained_path == "lerobot/pi0":
+            raise ValueError(
+                "You cannot load pretrained PI0 model when init_strategy is 'expert_only_he_init' due to differences in PaliGemma tokenizer vocab sizes."
+            )
+
+        if self.pretrained_path is not None and self.pretrained_path != "lerobot/pi0":
+            logging.info("Setting init_strategy to 'no_init' because we are resuming from a checkpoint.")
+            self.init_strategy = "no_init"
+
+    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.
+        """
+        # TODO: implement value error
+        # if not self.image_features and not self.env_state_feature:
+        #     raise ValueError("You must provide at least one image or the environment state among the inputs.")
+
+        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
+
+    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,
+        )
+
+    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