opentau 0.1.0__py3-none-any.whl

opentau/policies/pretrained.py

@@ -0,0 +1,315 @@
+# 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.
+
+"""Base class for pre-trained policies in OpenTau.
+
+This module defines the abstract base class `PreTrainedPolicy` which handles
+loading, saving, and basic interface requirements for all policy implementations
+in the OpenTau library. It integrates with Hugging Face Hub for model sharing
+and safetensors for efficient serialization.
+"""
+
+import abc
+import logging
+import os
+from pathlib import Path
+from typing import Type, TypeVar
+
+import packaging
+import safetensors
+import torch
+from huggingface_hub import hf_hub_download
+from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
+from huggingface_hub.errors import HfHubHTTPError
+from safetensors.torch import load_model as load_model_as_safetensor
+from safetensors.torch import save_model as save_model_as_safetensor
+from torch import Tensor, nn
+
+from opentau.configs.policies import PreTrainedConfig
+from opentau.policies.utils import log_model_loading_keys
+from opentau.utils.hub import HubMixin
+
+T = TypeVar("T", bound="PreTrainedPolicy")
+
+DEFAULT_POLICY_CARD = """
+---
+# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
+# Doc / guide: https://huggingface.co/docs/hub/model-cards
+{{ card_data }}
+---
+
+This policy has been pushed to the Hub using [OpenTau](https://github.com/TensorAuto/OpenTau):
+- Docs: {{ docs_url | default("[More Information Needed]", true) }}
+"""
+
+
+class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
+    """Base class for all policy models in OpenTau.
+
+    This class extends `nn.Module` and `HubMixin` to provide common functionality
+    for policy models, including configuration management, model loading/saving,
+    and abstract methods that all policies must implement.
+
+    Attributes:
+        config: The configuration instance for this policy.
+    """
+
+    config_class: None
+    """The configuration class associated with this policy. Must be defined in subclasses."""
+
+    name: None
+    """The name of the policy. Must be defined in subclasses."""
+
+    def __init__(self, config: PreTrainedConfig, *inputs, **kwargs):
+        """Initializes the PreTrainedPolicy.
+
+        Args:
+            config: The configuration object for the policy.
+            *inputs: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Raises:
+            ValueError: If `config` is not an instance of `PreTrainedConfig`.
+        """
+        super().__init__()
+        if not isinstance(config, PreTrainedConfig):
+            raise ValueError(
+                f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class "
+                "`PreTrainedConfig`. To create a model from a pretrained model use "
+                f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`"
+            )
+        self.config = config
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if not getattr(cls, "config_class", None):
+            raise TypeError(f"Class {cls.__name__} must define 'config_class'")
+        if not getattr(cls, "name", None):
+            raise TypeError(f"Class {cls.__name__} must define 'name'")
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        """Saves the policy and its configuration to a directory.
+
+        Args:
+            save_directory: The directory to save the policy to.
+        """
+        self.config._save_pretrained(save_directory)
+        model_to_save = self.module if hasattr(self, "module") else self
+        save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE))
+
+    @classmethod
+    def from_pretrained(
+        cls: Type[T],
+        pretrained_name_or_path: str | Path,
+        *,
+        config: PreTrainedConfig | None = None,
+        force_download: bool = False,
+        resume_download: bool | None = None,
+        proxies: dict | None = None,
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        revision: str | None = None,
+        strict: bool = False,
+        **kwargs,
+    ) -> T:
+        """Loads a pretrained policy from a local path or the Hugging Face Hub.
+
+        The policy is set in evaluation mode by default using `policy.eval()`
+        (dropout modules are deactivated). To train it, you should first set it
+        back in training mode with `policy.train()`.
+
+        Args:
+            pretrained_name_or_path: The name or path of the pretrained model.
+            config: Optional configuration object. If None, it will be loaded from the
+                pretrained model.
+            force_download: Whether to force download the model weights.
+            resume_download: Whether to resume an interrupted download.
+            proxies: Proxy configuration for downloading.
+            token: Hugging Face token for authentication.
+            cache_dir: Directory to cache downloaded files.
+            local_files_only: Whether to only look for local files.
+            revision: The specific model version to use (branch, tag, or commit hash).
+            strict: Whether to strictly enforce matching keys in state_dict.
+            **kwargs: Additional keyword arguments passed to the constructor.
+
+        Returns:
+            T: An instance of the loaded policy.
+
+        Raises:
+            FileNotFoundError: If the model file is not found.
+        """
+        if config is None:
+            config = PreTrainedConfig.from_pretrained(
+                pretrained_name_or_path=pretrained_name_or_path,
+                force_download=force_download,
+                resume_download=resume_download,
+                proxies=proxies,
+                token=token,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+                revision=revision,
+                **kwargs,
+            )
+        model_id = str(pretrained_name_or_path)
+        instance = cls(config, **kwargs)
+        if os.path.isdir(model_id):
+            print("Loading weights from local directory")
+            model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
+            policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
+        else:
+            try:
+                model_file = hf_hub_download(
+                    repo_id=model_id,
+                    filename=SAFETENSORS_SINGLE_FILE,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    token=token,
+                    local_files_only=local_files_only,
+                )
+                policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
+            except HfHubHTTPError as e:
+                raise FileNotFoundError(
+                    f"{SAFETENSORS_SINGLE_FILE} not found on the HuggingFace Hub in {model_id}"
+                ) from e
+
+        policy.eval()
+        return policy
+
+    def _tile_linear_input_weight(self, state_dict_to_load: dict):
+        """Modifies the `state_dict_to_load` in-place by tiling linear layer input weights.
+
+        This ensures compatibility with the model architecture when weight dimensions don't match exactly,
+        typically used for expanding input layers.
+
+        Args:
+            state_dict_to_load: The state dictionary to modify.
+        """
+        for name, submodule in self.named_modules():
+            if not isinstance(submodule, torch.nn.Linear):
+                continue
+            weight_name = f"{name}.weight"
+            if weight_name not in state_dict_to_load:
+                continue
+            weight = state_dict_to_load[weight_name]
+            assert len(weight.shape) == 2, f"Shape of {weight_name} must be 2D, got {weight.shape}"
+            out_dim, in_dim = weight.shape
+            assert submodule.out_features == out_dim, (
+                f"Output of {name} = {submodule.out_features} does not match loaded weight output dim {out_dim}"
+            )
+            if submodule.in_features == in_dim:
+                continue
+
+            logging.warning(f"Tiling {weight_name} from shape {weight.shape} to {submodule.weight.shape}")
+            repeat, remainder = divmod(submodule.in_features, in_dim)
+            weight = torch.cat([weight] * repeat + [weight[:, :remainder]], dim=1)
+            state_dict_to_load[weight_name] = weight
+
+    @classmethod
+    def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
+        """Loads model weights from a safetensors file.
+
+        Args:
+            model: The model instance to load weights into.
+            model_file: Path to the safetensors file.
+            map_location: Device to map the weights to.
+            strict: Whether to enforce strict key matching.
+
+        Returns:
+            T: The model with loaded weights.
+        """
+        # Create base kwargs
+        kwargs = {"strict": strict}
+
+        # Add device parameter for newer versions that support it
+        if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
+            kwargs["device"] = map_location
+
+        # Load the model with appropriate kwargs
+        missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
+        log_model_loading_keys(missing_keys, unexpected_keys)
+
+        # For older versions, manually move to device if needed
+        if "device" not in kwargs and map_location != "cpu":
+            logging.warning(
+                "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
+                " This means that the model is loaded on 'cpu' first and then copied to the device."
+                " This leads to a slower loading time."
+                " Please update safetensors to version 0.4.3 or above for improved performance."
+            )
+            model.to(map_location)
+        return model
+
+    # def generate_model_card(self, *args, **kwargs) -> ModelCard:
+    #     card = ModelCard.from_template(
+    #         card_data=self._hub_mixin_info.model_card_data,
+    #         template_str=self._hub_mixin_info.model_card_template,
+    #         repo_url=self._hub_mixin_info.repo_url,
+    #         docs_url=self._hub_mixin_info.docs_url,
+    #         **kwargs,
+    #     )
+    #     return card
+
+    @abc.abstractmethod
+    def get_optim_params(self) -> dict:
+        """Returns the policy-specific parameters dict to be passed on to the optimizer.
+
+        Returns:
+            dict: A dictionary of parameters to optimize.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def reset(self):
+        """Resets the policy state.
+
+        This method should be called whenever the environment is reset.
+        It handles tasks like clearing caches or resetting internal states for stateful policies.
+        """
+        raise NotImplementedError
+
+    # TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
+    @abc.abstractmethod
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
+        """Performs a forward pass of the policy.
+
+        Args:
+            batch: A dictionary of input tensors.
+
+        Returns:
+            tuple[Tensor, dict | None]: A tuple containing:
+                - The loss tensor.
+                - An optional dictionary of metrics or auxiliary outputs.
+                  Apart from the loss, items should be logging-friendly native Python types.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+        """Selects an action based on the input batch.
+
+        This method handles action selection during inference, including
+        caching for stateful policies (e.g. RNNs, Transformers).
+
+        Args:
+            batch: A dictionary of observation tensors.
+
+        Returns:
+            Tensor: The selected action(s).
+        """
+        raise NotImplementedError

opentau/policies/utils.py

@@ -0,0 +1,123 @@
+#!/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.
+
+"""Utility functions for policy implementations in OpenTau.
+
+This module provides helper functions for managing data queues, inspecting model
+properties (device, dtype), determining output shapes, and logging model loading
+information.
+"""
+
+import logging
+from collections import deque
+
+import torch
+from torch import nn
+
+
+def populate_queues(
+    queues: dict[str, deque], batch: dict[str, torch.Tensor], exclude_keys: list[str] | None = None
+) -> dict[str, deque]:
+    """Populates queues with batch data.
+
+    If a queue is not full (e.g. at the start of an episode), it is filled by repeating
+    the first observation. Otherwise, the latest observation is appended.
+
+    Args:
+        queues: A dictionary of deques to be populated.
+        batch: A dictionary containing the data to add to the queues.
+        exclude_keys: A list of keys to exclude from population. Defaults to None.
+
+    Returns:
+        dict[str, deque]: The updated dictionary of queues.
+    """
+    if exclude_keys is None:
+        exclude_keys = []
+    for key in batch:
+        # Ignore keys not in the queues already (leaving the responsibility to the caller to make sure the
+        # queues have the keys they want).
+        if key not in queues or key in exclude_keys:
+            continue
+        if len(queues[key]) != queues[key].maxlen:
+            # initialize by copying the first observation several times until the queue is full
+            while len(queues[key]) != queues[key].maxlen:
+                queues[key].append(batch[key])
+        else:
+            # add latest observation to the queue
+            queues[key].append(batch[key])
+    return queues
+
+
+def get_device_from_parameters(module: nn.Module) -> torch.device:
+    """Get a module's device by checking one of its parameters.
+
+    Note:
+        Assumes that all parameters have the same device.
+
+    Args:
+        module: The PyTorch module to inspect.
+
+    Returns:
+        torch.device: The device of the module's parameters.
+    """
+    return next(iter(module.parameters())).device
+
+
+def get_dtype_from_parameters(module: nn.Module) -> torch.dtype:
+    """Get a module's parameter dtype by checking one of its parameters.
+
+    Note:
+        Assumes that all parameters have the same dtype.
+
+    Args:
+        module: The PyTorch module to inspect.
+
+    Returns:
+        torch.dtype: The data type of the module's parameters.
+    """
+    return next(iter(module.parameters())).dtype
+
+
+def get_output_shape(module: nn.Module, input_shape: tuple) -> tuple:
+    """Calculates the output shape of a PyTorch module given an input shape.
+
+    Args:
+        module: A PyTorch module.
+        input_shape: A tuple representing the input shape, e.g., (batch_size, channels, height, width).
+
+    Returns:
+        tuple: The output shape of the module.
+    """
+    dummy_input = torch.zeros(size=input_shape)
+    with torch.inference_mode():
+        output = module(dummy_input)
+    return tuple(output.shape)
+
+
+def log_model_loading_keys(missing_keys: list[str], unexpected_keys: list[str]) -> None:
+    """Log missing and unexpected keys when loading a model.
+
+    Args:
+        missing_keys: Keys that were expected but not found.
+        unexpected_keys: Keys that were found but not expected.
+    """
+    if missing_keys:
+        # DO NOT UPDATE THIS MESSAGE WITHOUT UPDATING THE REGEX IN .gitlab/scripts/check_pi0_state_keys.py
+        logging.warning(f"Missing key(s) when loading model: {missing_keys}")
+    if unexpected_keys:
+        # DO NOT UPDATE THIS MESSAGE WITHOUT UPDATING THE REGEX IN .gitlab/scripts/check_pi0_state_keys.py
+        logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")

opentau/policies/value/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+"""Value Policy Module.
+
+This module implements value-based policies for robot control. It includes
+configurations, models, and reward functions for value estimation.
+"""

opentau/policies/value/configuration_value.py

@@ -0,0 +1,170 @@
+# Copyright 2025 Physical Intelligence and 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 Value policy.
+
+This module defines the `ValueConfig` class, which handles the configuration parameters
+for the Value policy. 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.reward import RewardConfig
+from opentau.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from opentau.optim.optimizers import AdamWConfig
+from opentau.optim.schedulers import (
+    CosineDecayWithWarmupSchedulerConfig,
+    LRSchedulerConfig,
+)
+
+
+@PreTrainedConfig.register_subclass("value")
+@dataclass
+class ValueConfig(PreTrainedConfig):
+    """Configuration class for the Value policy.
+
+    Args:
+        n_obs_steps: Number of observation steps to be used.
+        chunk_size: The chunk size for the policy.
+        normalization_mapping: Mapping of feature types to normalization modes.
+        max_state_dim: Maximum dimension for state vectors.
+        resize_imgs_with_padding: Tuple indicating the size to resize images with padding.
+        empty_cameras: Number of empty cameras to add.
+        tokenizer_max_length: Maximum length for the tokenizer.
+        reward_config: Configuration for the reward.
+        optimizer_lr: Learning rate for the optimizer.
+        optimizer_betas: Betas for the optimizer.
+        optimizer_eps: Epsilon for the optimizer.
+        optimizer_weight_decay: Weight decay for the optimizer.
+        scheduler_warmup_steps: Number of warmup steps for the scheduler.
+        scheduler_decay_steps: Number of decay steps for the scheduler.
+        scheduler_decay_lr: Decay learning rate for the scheduler.
+    """
+
+    # Input / output structure.
+    n_obs_steps: int = 1
+    chunk_size: int = 50
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MEAN_STD,
+            "VALUE": NormalizationMode.MEAN_STD,
+        }
+    )
+
+    # Shorter state vectors will be padded
+    max_state_dim: int = 32
+
+    # Image preprocessing
+    resize_imgs_with_padding: tuple[int, int] = (224, 224)
+
+    # Add empty images.
+    empty_cameras: int = 0
+
+    # Tokenizer
+    tokenizer_max_length: int = 48
+
+    # Reward config
+    reward_config: RewardConfig = field(default_factory=RewardConfig)
+
+    # 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
+
+    def __post_init__(self):
+        """Input validation (not exhaustive)."""
+        super().__post_init__()
+
+        if self.n_obs_steps != 1:
+            raise ValueError(
+                f"Multiple observation steps not handled yet. Got `nobs_steps={self.n_obs_steps}`"
+            )
+
+        max_episode_length = self.reward_config.reward_normalizer + self.reward_config.C_neg
+        assert max_episode_length < abs(self.reward_config.C_neg), (
+            "Max episode length should be less than the absolute value of C_neg for proper separation of failed and successful episodes"
+        )
+
+    def validate_features(self) -> None:
+        """Validates features and adds empty cameras if specified."""
+        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 optimizer preset configuration.
+
+        Returns:
+            AdamWConfig: The optimizer configuration.
+        """
+        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 scheduler preset configuration.
+
+        Returns:
+            CosineDecayWithWarmupSchedulerConfig: The scheduler configuration.
+        """
+        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:
+        """Returns the observation delta indices.
+
+        Returns:
+            None: Always returns None.
+        """
+        return None
+
+    @property
+    def action_delta_indices(self) -> list:
+        """Returns the action delta indices.
+
+        Returns:
+            list: List of indices from 0 to chunk_size.
+        """
+        return list(range(self.chunk_size))
+
+    @property
+    def reward_delta_indices(self) -> None:
+        """Returns the reward delta indices.
+
+        Returns:
+            None: Always returns None.
+        """
+        return None