opentau 0.1.0__py3-none-any.whl

opentau/utils/__init__.py

@@ -0,0 +1,20 @@
+# 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 and helpers for the OpenTau package.
+
+This module provides various utility functions for logging, random number generation,
+training utilities, device management, and other common operations used throughout
+the OpenTau codebase.
+"""

opentau/utils/accelerate_utils.py

@@ -0,0 +1,79 @@
+# 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.
+"""Utilities for managing Accelerate accelerator instances.
+
+This module provides functions for setting and getting a global accelerator
+instance, which is useful for accessing accelerator state throughout the codebase.
+"""
+
+import warnings
+
+from accelerate import Accelerator
+
+_acc: Accelerator | None = None
+
+
+def set_proc_accelerator(accelerator: Accelerator, allow_reset: bool = False) -> None:
+    """Set the global accelerator instance for the current process.
+
+    Args:
+        accelerator: Accelerator instance to set.
+        allow_reset: If True, allow resetting an already-set accelerator.
+            Defaults to False.
+
+    Raises:
+        AssertionError: If accelerator is not an Accelerator instance.
+        RuntimeError: If accelerator is already set and allow_reset is False.
+    """
+    global _acc
+
+    assert isinstance(accelerator, Accelerator), (
+        f"Expected an `Accelerator` got {type(accelerator)} with value {accelerator}."
+    )
+    if _acc is not None:
+        if allow_reset:
+            warnings.warn(
+                "Resetting the accelerator. This could have unintended side effects.",
+                UserWarning,
+                stacklevel=2,
+            )
+        else:
+            raise RuntimeError("Accelerator has already been set.")
+    _acc = accelerator
+
+
+def get_proc_accelerator() -> Accelerator:
+    """Get the global accelerator instance for the current process.
+
+    Returns:
+        The accelerator instance, or None if not set.
+    """
+    return _acc
+
+
+def acc_print(*args, **kwargs) -> None:
+    """Print with process index prefix when using accelerate.
+
+    If an accelerator is set, prints with a prefix showing the process index.
+    Otherwise, prints normally.
+
+    Args:
+        *args: Positional arguments to pass to print.
+        **kwargs: Keyword arguments to pass to print.
+    """
+    acc = get_proc_accelerator()
+    if acc is None:
+        print(*args, **kwargs)
+    else:
+        print(f"Acc[{acc.process_index} of {acc.num_processes}]", *args, **kwargs)

opentau/utils/benchmark.py

@@ -0,0 +1,98 @@
+#!/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.
+"""Utilities for benchmarking and timing code execution.
+
+This module provides the TimeBenchmark class for measuring execution time
+using context managers or decorators in a thread-safe manner.
+"""
+
+import threading
+import time
+from contextlib import ContextDecorator
+
+
+class TimeBenchmark(ContextDecorator):
+    """
+    Measures execution time using a context manager or decorator.
+
+    This class supports both context manager and decorator usage, and is thread-safe for multithreaded
+    environments.
+
+    Args:
+        print: If True, prints the elapsed time upon exiting the context or completing the function.
+            Defaults to False.
+
+    Examples:
+
+        Using as a context manager::
+
+            >>> benchmark = TimeBenchmark()
+            >>> with benchmark:
+            ...     time.sleep(1)
+            >>> print(f"Block took {benchmark.result:.4f} seconds")
+            Block took approximately 1.0000 seconds
+
+        Using with multithreading::
+
+            import threading
+
+            benchmark = TimeBenchmark()
+
+            def context_manager_example():
+                with benchmark:
+                    time.sleep(0.01)
+                print(f"Block took {benchmark.result_ms:.2f} milliseconds")
+
+            threads = []
+            for _ in range(3):
+                t1 = threading.Thread(target=context_manager_example)
+                threads.append(t1)
+
+            for t in threads:
+                t.start()
+
+            for t in threads:
+                t.join()
+
+            # Expected output:
+            # Block took approximately 10.00 milliseconds
+            # Block took approximately 10.00 milliseconds
+            # Block took approximately 10.00 milliseconds
+    """
+
+    def __init__(self, print=False):
+        self.local = threading.local()
+        self.print_time = print
+
+    def __enter__(self):
+        self.local.start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, *exc):
+        self.local.end_time = time.perf_counter()
+        self.local.elapsed_time = self.local.end_time - self.local.start_time
+        if self.print_time:
+            print(f"Elapsed time: {self.local.elapsed_time:.4f} seconds")
+        return False
+
+    @property
+    def result(self):
+        return getattr(self.local, "elapsed_time", None)
+
+    @property
+    def result_ms(self):
+        return self.result * 1e3

opentau/utils/fake_tensor.py

@@ -0,0 +1,81 @@
+# 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.
+"""Utilities for working with PyTorch FakeTensor.
+
+This module provides a FakeTensorContext class and decorator for running code
+with FakeTensor mode enabled, which is useful for shape inference and testing
+without actual tensor computations.
+"""
+
+import functools
+
+from torch._subclasses import FakeTensorMode
+from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+from opentau.utils.monkey_patch import (
+    torch_fake_tensor_beta_validate_args_patch,
+    torch_fake_tensor_is_inf_patch,
+    torch_fake_tensor_module_to_patch,
+    torch_fake_tensor_to_numpy_patch,
+)
+
+# Share the ShapeEnv instance across all FakeTensorContext instances
+# Without this, each FakeTensor.item() call would start numbering from 0, which is wrong.
+_shared_shape_env = ShapeEnv()
+
+
+class FakeTensorContext:
+    """Context manager for enabling FakeTensor mode with necessary patches.
+
+    This context manager applies all necessary monkey patches for FakeTensor
+    compatibility and manages the FakeTensorMode lifecycle.
+
+    Args:
+        allow_non_fake_inputs: If True, allow non-fake tensors as inputs.
+            Defaults to True.
+    """
+
+    def __init__(self, allow_non_fake_inputs: bool = True):
+        self.mode = FakeTensorMode(
+            shape_env=_shared_shape_env,
+            allow_non_fake_inputs=allow_non_fake_inputs,
+        )
+        torch_fake_tensor_module_to_patch()
+        torch_fake_tensor_to_numpy_patch()
+        torch_fake_tensor_beta_validate_args_patch()
+        torch_fake_tensor_is_inf_patch()
+
+    def __enter__(self):
+        return self.mode.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self.mode.__exit__(exc_type, exc_val, exc_tb)
+
+
+def run_with_fake_tensor(fn):
+    """Decorator to run a function with FakeTensor enabled.
+
+    Args:
+        fn: Function to wrap.
+
+    Returns:
+        Wrapped function that runs with FakeTensorContext enabled.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        with FakeTensorContext():
+            return fn(*args, **kwargs)
+
+    return wrapper

opentau/utils/hub.py

@@ -0,0 +1,209 @@
+# 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.
+"""Utilities for interacting with the Hugging Face Hub.
+
+This module provides the HubMixin class which enables objects to be saved to
+and loaded from the Hugging Face Hub, similar to ModelHubMixin but with fewer
+assumptions about the object type.
+"""
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, Type, TypeVar
+
+from huggingface_hub import HfApi
+from huggingface_hub.utils import validate_hf_hub_args
+
+T = TypeVar("T", bound="HubMixin")
+
+
+class HubMixin:
+    """
+    A Mixin containing the functionality to push an object to the hub.
+
+    This is similar to huggingface_hub.ModelHubMixin but is lighter and makes less assumptions about its
+    subclasses (in particular, the fact that it's not necessarily a model).
+
+    The inheriting classes must implement '_save_pretrained' and 'from_pretrained'.
+    """
+
+    def save_pretrained(
+        self,
+        save_directory: str | Path,
+        *,
+        repo_id: str | None = None,
+        push_to_hub: bool = False,
+        card_kwargs: dict[str, Any] | None = None,
+        **push_to_hub_kwargs,
+    ) -> str | None:
+        """
+        Save object in local directory.
+
+        Args:
+            save_directory (`str` or `Path`):
+                Path to directory in which the object will be saved.
+            push_to_hub (`bool`, *optional*, defaults to `False`):
+                Whether or not to push your object to the Huggingface Hub after saving it.
+            repo_id (`str`, *optional*):
+                ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if
+                not provided.
+            card_kwargs (`Dict[str, Any]`, *optional*):
+                Additional arguments passed to the card template to customize the card.
+            push_to_hub_kwargs:
+                Additional key word arguments passed along to the [`~HubMixin.push_to_hub`] method.
+        Returns:
+            `str` or `None`: url of the commit on the Hub if `push_to_hub=True`, `None` otherwise.
+        """
+        save_directory = Path(save_directory)
+        save_directory.mkdir(parents=True, exist_ok=True)
+
+        # save object (weights, files, etc.)
+        self._save_pretrained(save_directory)
+
+        # push to the Hub if required
+        if push_to_hub:
+            if repo_id is None:
+                repo_id = save_directory.name  # Defaults to `save_directory` name
+            return self.push_to_hub(repo_id=repo_id, card_kwargs=card_kwargs, **push_to_hub_kwargs)
+        return None
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        """
+        Overwrite this method in subclass to define how to save your object.
+
+        Args:
+            save_directory (`str` or `Path`):
+                Path to directory in which the object files will be saved.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(
+        cls: Type[T],
+        pretrained_name_or_path: str | Path,
+        *,
+        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,
+        **kwargs,
+    ) -> T:
+        """
+        Download the object from the Huggingface Hub and instantiate it.
+
+        Args:
+            pretrained_name_or_path (`str`, `Path`):
+                - Either the `repo_id` (string) of the object hosted on the Hub, e.g. `lerobot/diffusion_pusht`.
+                - Or a path to a `directory` containing the object files saved using `.save_pretrained`,
+                    e.g., `../path/to/my_model_directory/`.
+            revision (`str`, *optional*):
+                Revision on the Hub. Can be a branch name, a git tag or any commit id.
+                Defaults to the latest commit on `main` branch.
+            force_download (`bool`, *optional*, defaults to `False`):
+                Whether to force (re-)downloading the files from the Hub, overriding the existing cache.
+            proxies (`Dict[str, str]`, *optional*):
+                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on every request.
+            token (`str` or `bool`, *optional*):
+                The token to use as HTTP bearer authorization for remote files. By default, it will use the token
+                cached when running `huggingface-cli login`.
+            cache_dir (`str`, `Path`, *optional*):
+                Path to the folder where cached files are stored.
+            local_files_only (`bool`, *optional*, defaults to `False`):
+                If `True`, avoid downloading the file and return the path to the local cached file if it exists.
+            kwargs (`Dict`, *optional*):
+                Additional kwargs to pass to the object during initialization.
+        """
+        raise NotImplementedError
+
+    @validate_hf_hub_args
+    def push_to_hub(
+        self,
+        repo_id: str,
+        *,
+        commit_message: str | None = None,
+        private: bool | None = None,
+        token: str | None = None,
+        branch: str | None = None,
+        create_pr: bool | None = None,
+        allow_patterns: list[str] | str | None = None,
+        ignore_patterns: list[str] | str | None = None,
+        delete_patterns: list[str] | str | None = None,
+        card_kwargs: dict[str, Any] | None = None,
+    ) -> str:
+        """
+        Upload model checkpoint to the Hub.
+
+        Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
+        `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
+        details.
+
+        Args:
+            repo_id (`str`):
+                ID of the repository to push to (example: `"username/my-model"`).
+            commit_message (`str`, *optional*):
+                Message to commit while pushing.
+            private (`bool`, *optional*):
+                Whether the repository created should be private.
+                If `None` (default), the repo will be public unless the organization's default is private.
+            token (`str`, *optional*):
+                The token to use as HTTP bearer authorization for remote files. By default, it will use the token
+                cached when running `huggingface-cli login`.
+            branch (`str`, *optional*):
+                The git branch on which to push the model. This defaults to `"main"`.
+            create_pr (`boolean`, *optional*):
+                Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.
+            allow_patterns (`List[str]` or `str`, *optional*):
+                If provided, only files matching at least one pattern are pushed.
+            ignore_patterns (`List[str]` or `str`, *optional*):
+                If provided, files matching any of the patterns are not pushed.
+            delete_patterns (`List[str]` or `str`, *optional*):
+                If provided, remote files matching any of the patterns will be deleted from the repo.
+            card_kwargs (`Dict[str, Any]`, *optional*):
+                Additional arguments passed to the card template to customize the card.
+
+        Returns:
+            The url of the commit of your object in the given repository.
+        """
+        api = HfApi(token=token)
+        repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
+
+        if commit_message is None:
+            if "Policy" in self.__class__.__name__:
+                commit_message = "Upload policy"
+            elif "Config" in self.__class__.__name__:
+                commit_message = "Upload config"
+            else:
+                commit_message = f"Upload {self.__class__.__name__}"
+
+        # Push the files to the repo in a single commit
+        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
+            saved_path = Path(tmp) / repo_id
+            self.save_pretrained(saved_path, card_kwargs=card_kwargs)
+            return api.upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=saved_path,
+                commit_message=commit_message,
+                revision=branch,
+                create_pr=create_pr,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+                delete_patterns=delete_patterns,
+            )

opentau/utils/import_utils.py

@@ -0,0 +1,79 @@
+#!/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.
+"""Utilities for checking package availability and versions.
+
+This module provides functions to check if packages are installed and optionally
+retrieve their versions without importing them, which is useful for conditional
+imports and dependency checking.
+"""
+
+import importlib
+import logging
+
+
+def is_package_available(pkg_name: str, return_version: bool = False) -> tuple[bool, str] | bool:
+    """Check if a package is available and optionally return its version.
+
+    Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/utils/import_utils.py
+
+    This function checks if the package spec exists and grabs its version to
+    avoid importing a local directory. Note: this doesn't work for all packages.
+
+    Args:
+        pkg_name: Name of the package to check.
+        return_version: If True, return a tuple of (available, version).
+            If False, return only the availability boolean. Defaults to False.
+
+    Returns:
+        If return_version is False, returns a boolean indicating availability.
+        If return_version is True, returns a tuple of (available, version).
+    """
+    package_exists = importlib.util.find_spec(pkg_name) is not None
+    package_version = "N/A"
+    if package_exists:
+        try:
+            # Primary method to get the package version
+            package_version = importlib.metadata.version(pkg_name)
+        except importlib.metadata.PackageNotFoundError:
+            # Fallback method: Only for "torch" and versions containing "dev"
+            if pkg_name == "torch":
+                try:
+                    package = importlib.import_module(pkg_name)
+                    temp_version = getattr(package, "__version__", "N/A")
+                    # Check if the version contains "dev"
+                    if "dev" in temp_version:
+                        package_version = temp_version
+                        package_exists = True
+                    else:
+                        package_exists = False
+                except ImportError:
+                    # If the package can't be imported, it's not available
+                    package_exists = False
+            else:
+                # For packages other than "torch", don't attempt the fallback and set as not available
+                package_exists = False
+        logging.debug(f"Detected {pkg_name} version: {package_version}")
+    if return_version:
+        return package_exists, package_version
+    else:
+        return package_exists
+
+
+_torch_available, _torch_version = is_package_available("torch", return_version=True)
+_gym_xarm_available = is_package_available("gym_xarm")
+_gym_aloha_available = is_package_available("gym_aloha")
+_gym_pusht_available = is_package_available("gym_pusht")