opentau 0.1.0__py3-none-any.whl

opentau/configs/parser.py

@@ -0,0 +1,393 @@
+# 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.
+"""Command-line argument parsing and configuration loading utilities.
+
+This module provides utilities for parsing command-line arguments, loading
+configuration files from local paths or the HuggingFace Hub, and handling
+plugin discovery and loading. It extends draccus functionality with support
+for path-based configuration loading and plugin system integration.
+"""
+
+import importlib
+import inspect
+import pkgutil
+import sys
+from argparse import ArgumentError
+from functools import wraps
+from pathlib import Path
+from typing import Sequence
+
+import draccus
+
+from opentau.utils.utils import has_method
+
+PATH_KEY = "path"
+PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
+draccus.set_config_type("json")
+
+
+def get_cli_overrides(field_name: str, args: Sequence[str] | None = None) -> list[str] | None:
+    """Parse arguments from CLI at a given nested attribute level.
+
+    This function extracts command-line arguments that are nested under a specific
+    field name and returns them with the field name prefix removed.
+
+    Args:
+        field_name: The field name to extract nested arguments for.
+        args: Sequence of command-line arguments to parse. If None, uses sys.argv[1:].
+            Defaults to None.
+
+    Returns:
+        List of denested arguments with the field name prefix removed, or None if
+        no matching arguments are found.
+
+    Example:
+        Supposing the main script was called with:
+        ```
+        python myscript.py --arg1=1 --arg2.subarg1=abc --arg2.subarg2=some/path
+        ```
+
+        If called during execution of myscript.py, `get_cli_overrides("arg2")` will
+        return:
+        ```
+        ["--subarg1=abc", "--subarg2=some/path"]
+        ```
+    """
+    if args is None:
+        args = sys.argv[1:]
+    attr_level_args = []
+    detect_string = f"--{field_name}."
+    exclude_strings = (f"--{field_name}.{draccus.CHOICE_TYPE_KEY}=", f"--{field_name}.{PATH_KEY}=")
+    for arg in args:
+        if arg.startswith(detect_string) and not arg.startswith(exclude_strings):
+            denested_arg = f"--{arg.removeprefix(detect_string)}"
+            attr_level_args.append(denested_arg)
+
+    return attr_level_args
+
+
+def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
+    """Parse a single command-line argument value.
+
+    Args:
+        arg_name: Name of the argument to parse (without the '--' prefix).
+        args: Sequence of command-line arguments to parse. If None, uses sys.argv[1:].
+            Defaults to None.
+
+    Returns:
+        The value of the argument if found, or None if not found.
+
+    Example:
+        For command-line arguments `['--batch_size=32', '--lr=0.001']`:
+        - `parse_arg('batch_size')` returns `'32'`
+        - `parse_arg('lr')` returns `'0.001'`
+        - `parse_arg('missing')` returns `None`
+    """
+    if args is None:
+        args = sys.argv[1:]
+    prefix = f"--{arg_name}="
+    for arg in args:
+        if arg.startswith(prefix):
+            return arg[len(prefix) :]
+    return None
+
+
+def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict:
+    """Parse plugin-related arguments from command-line arguments.
+
+    This function extracts arguments from command-line arguments that match a specified suffix pattern.
+    It processes arguments in the format '--key=value' and returns them as a dictionary.
+
+    Args:
+        plugin_arg_suffix (str): The suffix to identify plugin-related arguments.
+        cli_args (Sequence[str]): A sequence of command-line arguments to parse.
+
+    Returns:
+        dict: A dictionary containing the parsed plugin arguments where:
+            - Keys are the argument names (with '--' prefix removed if present)
+            - Values are the corresponding argument values
+
+    Example:
+        >>> args = ['--env.discover_packages_path=my_package',
+        ...         '--other_arg=value']
+        >>> parse_plugin_args('discover_packages_path', args)
+        {'env.discover_packages_path': 'my_package'}
+    """
+    plugin_args = {}
+    for arg in args:
+        if "=" in arg and plugin_arg_suffix in arg:
+            key, value = arg.split("=", 1)
+            # Remove leading '--' if present
+            if key.startswith("--"):
+                key = key[2:]
+            plugin_args[key] = value
+    return plugin_args
+
+
+class PluginLoadError(Exception):
+    """Raised when a plugin fails to load."""
+
+
+def load_plugin(plugin_path: str) -> None:
+    """Load and initialize a plugin from a given Python package path.
+
+    This function attempts to load a plugin by importing its package and any submodules.
+    Plugin registration is expected to happen during package initialization, i.e. when
+    the package is imported the gym environment should be registered and the config classes
+    registered with their parents using the `register_subclass` decorator.
+
+    Args:
+        plugin_path (str): The Python package path to the plugin (e.g. "mypackage.plugins.myplugin")
+
+    Raises:
+        PluginLoadError: If the plugin cannot be loaded due to import errors or if the package path is invalid.
+
+    Examples:
+        >>> load_plugin("external_plugin.core")       # Loads plugin from external package
+
+    Notes:
+        - The plugin package should handle its own registration during import
+        - All submodules in the plugin package will be imported
+        - Implementation follows the plugin discovery pattern from Python packaging guidelines
+
+    See Also:
+        https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/
+    """
+    try:
+        package_module = importlib.import_module(plugin_path, __package__)
+    except (ImportError, ModuleNotFoundError) as e:
+        raise PluginLoadError(
+            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+        ) from e
+
+    def iter_namespace(ns_pkg):
+        return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
+
+    try:
+        for _finder, pkg_name, _ispkg in iter_namespace(package_module):
+            importlib.import_module(pkg_name)
+    except ImportError as e:
+        raise PluginLoadError(
+            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+        ) from e
+
+
+def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
+    """Get the path argument for a given field name.
+
+    This function extracts the path argument for a field, which is typically
+    specified as `--field_name.path=some/path`.
+
+    Args:
+        field_name: The field name to get the path argument for.
+        args: Sequence of command-line arguments to parse. If None, uses sys.argv[1:].
+            Defaults to None.
+
+    Returns:
+        The path value if found, or None if not found.
+
+    Example:
+        For `--policy.path=/path/to/config`, `get_path_arg('policy')` returns
+        `'/path/to/config'`.
+    """
+    return parse_arg(f"{field_name}.{PATH_KEY}", args)
+
+
+def get_type_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
+    """Get the type argument for a given field name.
+
+    This function extracts the type argument for a field, which is typically
+    specified as `--field_name.type=SomeType`.
+
+    Args:
+        field_name: The field name to get the type argument for.
+        args: Sequence of command-line arguments to parse. If None, uses sys.argv[1:].
+            Defaults to None.
+
+    Returns:
+        The type value if found, or None if not found.
+
+    Example:
+        For `--policy.type=Pi0Config`, `get_type_arg('policy')` returns `'Pi0Config'`.
+    """
+    return parse_arg(f"{field_name}.{draccus.CHOICE_TYPE_KEY}", args)
+
+
+def filter_arg(field_to_filter: str, args: Sequence[str] | None = None) -> list[str]:
+    """Filter out arguments matching a specific field name.
+
+    Args:
+        field_to_filter: The field name to filter out (without the '--' prefix).
+        args: Sequence of command-line arguments to filter. If None, uses sys.argv[1:].
+            Defaults to None.
+
+    Returns:
+        List of arguments with the specified field filtered out.
+
+    Example:
+        For `['--batch_size=32', '--lr=0.001', '--batch_size=64']`:
+        `filter_arg('batch_size')` returns `['--lr=0.001']`.
+    """
+    if args is None:
+        args = sys.argv[1:]
+    return [arg for arg in args if not arg.startswith(f"--{field_to_filter}=")]
+
+
+def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | None = None) -> list[str]:
+    """Filter command-line arguments related to fields with specific path arguments.
+
+    This function removes all arguments related to specified fields when a path
+    argument is present for those fields. It also validates that path and type
+    arguments are not both specified for the same field.
+
+    Args:
+        fields_to_filter: A single field name or a list of field names whose
+            arguments need to be filtered.
+        args: The sequence of command-line arguments to be filtered. If None,
+            uses sys.argv[1:]. Defaults to None.
+
+    Returns:
+        A filtered list of arguments, with arguments related to the specified
+        fields removed.
+
+    Raises:
+        ArgumentError: If both a path argument (e.g., `--field_name.path`) and a
+            type argument (e.g., `--field_name.type`) are specified for the same field.
+    """
+    if isinstance(fields_to_filter, str):
+        fields_to_filter = [fields_to_filter]
+
+    filtered_args = args
+    for field in fields_to_filter:
+        if get_path_arg(field, args):
+            if get_type_arg(field, args):
+                raise ArgumentError(
+                    argument=None,
+                    message=f"Cannot specify both --{field}.{PATH_KEY} and --{field}.{draccus.CHOICE_TYPE_KEY}",
+                )
+            filtered_args = [arg for arg in filtered_args if not arg.startswith(f"--{field}.")]
+
+    return filtered_args
+
+
+def filter_distributed_args(args: Sequence[str] | None = None) -> list[str]:
+    """Filter out distributed training arguments.
+
+    This function removes arguments that are automatically injected by distributed
+    training frameworks (e.g., DeepSpeed, torchrun) but not recognized by the
+    custom argument parser.
+
+    Args:
+        args: The sequence of command-line arguments to be filtered. If None,
+            uses sys.argv[1:]. Defaults to None.
+
+    Returns:
+        A filtered list of arguments with distributed training arguments removed.
+
+    Note:
+        Filtered arguments include: local_rank, node_rank, master_addr, master_port,
+        world_size, and rank.
+    """
+    if args is None:
+        args = sys.argv[1:]
+
+    # List of distributed training arguments to filter out
+    distributed_args = [
+        "--local_rank",
+        "--local-rank",
+        "--node_rank",
+        "--node-rank",
+        "--master_addr",
+        "--master-addr",
+        "--master_port",
+        "--master-port",
+        "--world_size",
+        "--world-size",
+        "--rank",
+    ]
+
+    filtered_args = []
+    for arg in args:
+        should_filter = False
+        for distributed_arg in distributed_args:
+            if arg.startswith(f"{distributed_arg}=") or arg == distributed_arg:
+                should_filter = True
+                break
+        if not should_filter:
+            filtered_args.append(arg)
+
+    return filtered_args
+
+
+def wrap(config_path: Path | None = None):
+    """Wrap a function to handle configuration parsing with enhanced features.
+
+    This decorator is similar to `draccus.wrap` but provides three additional features:
+
+    1. Removes '.path' arguments from CLI to process them later
+    2. If a 'config_path' is passed and the main config class has a 'from_pretrained'
+       method, initializes it from there to allow fetching configs from the hub directly
+    3. Loads plugins specified in CLI arguments. These plugins typically register
+       their own subclasses of config classes, so that draccus can find the right
+       class to instantiate from the CLI '.type' arguments
+
+    Args:
+        config_path: Optional path to a configuration file. If provided and the
+            config class supports `from_pretrained`, will load from this path.
+            Defaults to None.
+
+    Returns:
+        A decorator function that wraps the target function with enhanced configuration
+        parsing capabilities.
+
+    Note:
+        This is a HACK wrapper around draccus.wrap to add custom functionality.
+    """
+
+    def wrapper_outer(fn):
+        @wraps(fn)
+        def wrapper_inner(*args, **kwargs):
+            argspec = inspect.getfullargspec(fn)
+            argtype = argspec.annotations[argspec.args[0]]
+            if len(args) > 0 and type(args[0]) is argtype:
+                cfg = args[0]
+                args = args[1:]
+            else:
+                cli_args = sys.argv[1:]
+                # Filter out distributed training arguments first
+                cli_args = filter_distributed_args(cli_args)
+                plugin_args = parse_plugin_args(PLUGIN_DISCOVERY_SUFFIX, cli_args)
+                for plugin_cli_arg, plugin_path in plugin_args.items():
+                    try:
+                        load_plugin(plugin_path)
+                    except PluginLoadError as e:
+                        # add the relevant CLI arg to the error message
+                        raise PluginLoadError(f"{e}\nFailed plugin CLI Arg: {plugin_cli_arg}") from e
+                    cli_args = filter_arg(plugin_cli_arg, cli_args)
+                config_path_cli = parse_arg("config_path", cli_args)
+                if has_method(argtype, "__get_path_fields__"):
+                    path_fields = argtype.__get_path_fields__()
+                    cli_args = filter_path_args(path_fields, cli_args)
+                if has_method(argtype, "from_pretrained") and config_path_cli:
+                    cli_args = filter_arg("config_path", cli_args)
+                    cfg = argtype.from_pretrained(config_path_cli, cli_args=cli_args)
+                else:
+                    cfg = draccus.parse(config_class=argtype, config_path=config_path, args=cli_args)
+            response = fn(cfg, *args, **kwargs)
+            return response
+
+        return wrapper_inner
+
+    return wrapper_outer

opentau/configs/policies.py

@@ -0,0 +1,297 @@
+# 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.
+"""Policy configuration module.
+
+This module provides the base PreTrainedConfig class for policy models, which
+defines the interface and common functionality for all policy configurations.
+It includes support for feature definitions, normalization modes, and loading
+configurations from pretrained models or local paths.
+"""
+
+import abc
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Type, TypeVar
+
+import draccus
+from huggingface_hub import hf_hub_download
+from huggingface_hub.constants import CONFIG_NAME
+from huggingface_hub.errors import HfHubHTTPError
+
+from opentau.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from opentau.optim.optimizers import OptimizerConfig
+from opentau.optim.schedulers import LRSchedulerConfig
+from opentau.utils.hub import HubMixin
+
+# Generic variable that is either PreTrainedConfig or a subclass thereof
+T = TypeVar("T", bound="PreTrainedConfig")
+
+
+@dataclass
+class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
+    """
+    Base configuration class for policy models.
+
+    Args:
+        n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the
+            current step and additional steps going back).
+        input_shapes: A dictionary defining the shapes of the input data for the policy.
+        output_shapes: A dictionary defining the shapes of the output data for the policy.
+        input_normalization_modes: A dictionary with key representing the modality and the value specifies the
+            normalization mode to apply.
+        output_normalization_modes: Similar dictionary as `input_normalization_modes`, but to unnormalize to
+            the original scale.
+    """
+
+    n_obs_steps: int = 1
+    normalization_mapping: dict[str, NormalizationMode] = field(default_factory=dict)
+
+    input_features: dict[str, PolicyFeature] = field(default_factory=dict)
+    output_features: dict[str, PolicyFeature] = field(default_factory=dict)
+
+    device: str | None = None  # cuda | cpu | mps
+    # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
+    # automatic gradient scaling is used.
+    use_amp: bool = False
+    pretrained_path: str | None = None
+
+    # Mean latency of cloud VLM in seconds.
+    cloud_vlm_latency_mean: float = 0.0
+    # Standard deviation of latency of cloud VLM in seconds.
+    cloud_vlm_latency_std: float = 0.0
+    # Lower bound of latency of cloud VLM in seconds.
+    cloud_vlm_latency_lower: float = 0.0
+    # Upper bound of latency of cloud VLM in seconds.
+    cloud_vlm_latency_upper: float = 0.0
+
+    # Mean latency of action decoder in seconds.
+    action_decoder_latency_mean: float = 0.0
+    # Standard deviation of latency of action decoder in seconds.
+    action_decoder_latency_std: float = 0.0
+    # Lower bound of latency of action decoder in seconds.
+    action_decoder_latency_lower: float = 0.0
+    # Upper bound of latency of action decoder in seconds.
+    action_decoder_latency_upper: float = 0.0
+
+    def __post_init__(self):
+        """Initialize post-creation attributes.
+
+        This method can be overridden by subclasses to perform additional
+        initialization after the dataclass is created.
+        """
+        pass
+
+    @property
+    def type(self) -> str:
+        """Get the type name of this configuration.
+
+        Returns:
+            The choice name of this configuration class.
+        """
+        return self.get_choice_name(self.__class__)
+
+    @abc.abstractproperty
+    def observation_delta_indices(self) -> list | 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.
+        """
+        raise NotImplementedError
+
+    @abc.abstractproperty
+    def action_delta_indices(self) -> list | None:
+        """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.
+        """
+        raise NotImplementedError
+
+    @abc.abstractproperty
+    def reward_delta_indices(self) -> list | 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.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_optimizer_preset(self) -> OptimizerConfig:
+        """Get the default optimizer configuration for this policy.
+
+        Returns:
+            An OptimizerConfig instance with default settings for this policy type.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_scheduler_preset(self) -> LRSchedulerConfig | None:
+        """Get the default learning rate scheduler configuration for this policy.
+
+        Returns:
+            An LRSchedulerConfig instance with default settings for this policy type,
+            or None if no scheduler should be used.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def validate_features(self) -> None:
+        """Validate that the feature configuration is correct.
+
+        This method should check that all required features are present and
+        have valid configurations.
+
+        Raises:
+            ValueError: If the feature configuration is invalid.
+        """
+        raise NotImplementedError
+
+    @property
+    def robot_state_feature(self) -> PolicyFeature | None:
+        """Get the robot state feature from input features.
+
+        Returns:
+            The PolicyFeature with type STATE if found, or None otherwise.
+        """
+        for _, ft in self.input_features.items():
+            if ft.type is FeatureType.STATE:
+                return ft
+        return None
+
+    @property
+    def env_state_feature(self) -> PolicyFeature | None:
+        """Get the environment state feature from input features.
+
+        Returns:
+            The PolicyFeature with type ENV if found, or None otherwise.
+        """
+        for _, ft in self.input_features.items():
+            if ft.type is FeatureType.ENV:
+                return ft
+        return None
+
+    @property
+    def image_features(self) -> dict[str, PolicyFeature]:
+        """Get all visual/image features from input features.
+
+        Returns:
+            Dictionary mapping feature names to PolicyFeature instances with
+            type VISUAL.
+        """
+        return {key: ft for key, ft in self.input_features.items() if ft.type is FeatureType.VISUAL}
+
+    @property
+    def action_feature(self) -> PolicyFeature | None:
+        """Get the action feature from output features.
+
+        Returns:
+            The PolicyFeature with type ACTION if found, or None otherwise.
+        """
+        for _, ft in self.output_features.items():
+            if ft.type is FeatureType.ACTION:
+                return ft
+        return None
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        """Save the configuration to a directory.
+
+        Args:
+            save_directory: Directory path where the configuration will be saved.
+        """
+        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
+
+    @classmethod
+    def from_pretrained(
+        cls: Type[T],
+        pretrained_name_or_path: str | Path,
+        *,
+        force_download: bool = False,
+        resume_download: bool = 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,
+        **policy_kwargs,
+    ) -> T:
+        """Load a policy configuration from a pretrained model or local path.
+
+        Args:
+            cls: The class to instantiate.
+            pretrained_name_or_path: Can be either:
+
+                - A string, the model id of a pretrained config hosted inside a model
+                  repo on huggingface.co.
+                - A path to a directory containing a configuration file saved using
+                  the `_save_pretrained` method.
+            force_download: Whether to force (re-)downloading the config files and
+                configuration from the HuggingFace Hub. Defaults to False.
+            resume_download: Whether to resume downloading the config files.
+                Defaults to None.
+            proxies: Dictionary of proxies to use for requests. Defaults to None.
+            token: The token to use as HTTP bearer authorization. If True, will use
+                the token generated when running `huggingface-cli login`. Defaults to None.
+            cache_dir: Path to a directory in which a downloaded pretrained model
+                configuration should be cached. Defaults to None.
+            local_files_only: Whether to only look at local files (i.e., do not try
+                to download the config). Defaults to False.
+            revision: The specific model version to use. It can be a branch name, a
+                tag name, or a commit id. Defaults to None.
+            **policy_kwargs: Additional keyword arguments. May include 'cli_overrides'
+                for command-line argument overrides.
+
+        Returns:
+            An instance of the configuration class loaded from the specified path.
+
+        Raises:
+            FileNotFoundError: If the configuration file is not found on the
+                HuggingFace Hub or in the local path.
+        """
+        model_id = str(pretrained_name_or_path)
+        config_file: str | None = None
+        if Path(model_id).is_dir():
+            if CONFIG_NAME in os.listdir(model_id):
+                config_file = os.path.join(model_id, CONFIG_NAME)
+            else:
+                print(f"{CONFIG_NAME} not found in {Path(model_id).resolve()}")
+        else:
+            try:
+                config_file = hf_hub_download(
+                    repo_id=model_id,
+                    filename=CONFIG_NAME,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    token=token,
+                    local_files_only=local_files_only,
+                )
+            except HfHubHTTPError as e:
+                raise FileNotFoundError(
+                    f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
+                ) from e
+
+        # HACK: this is very ugly, ideally we'd like to be able to do that natively with draccus
+        # something like --policy.path (in addition to --policy.type)
+        cli_overrides = policy_kwargs.pop("cli_overrides", [])
+        return draccus.parse(cls, config_file, args=cli_overrides)