opentau 0.1.0__py3-none-any.whl

opentau/utils/train_utils.py

@@ -0,0 +1,198 @@
+#!/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 training checkpoint management and state persistence.
+
+This module provides functions for saving and loading training checkpoints,
+managing checkpoint directories, and pruning old checkpoints.
+"""
+
+import logging
+import shutil
+from pathlib import Path
+
+from termcolor import colored
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+from opentau.configs.train import TrainPipelineConfig
+from opentau.constants import (
+    CHECKPOINTS_DIR,
+    LAST_CHECKPOINT_LINK,
+    TRAINING_STEP,
+)
+from opentau.datasets.utils import load_json, write_json
+from opentau.utils.random_utils import load_rng_state, save_rng_state
+
+
+def log_output_dir(out_dir):
+    """Log the output directory path with colored formatting.
+
+    Args:
+        out_dir: Output directory path to log.
+    """
+    logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {out_dir}")
+
+
+def get_step_identifier(step: int, total_steps: int) -> str:
+    """Generate a zero-padded step identifier string.
+
+    Args:
+        step: Current step number.
+        total_steps: Total number of steps (used to determine padding width).
+
+    Returns:
+        Zero-padded string representation of the step number.
+    """
+    num_digits = max(6, len(str(total_steps)))
+    return f"{step:0{num_digits}d}"
+
+
+def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Path:
+    """Returns the checkpoint sub-directory corresponding to the step number."""
+    step_identifier = get_step_identifier(step, total_steps)
+    return output_dir / CHECKPOINTS_DIR / step_identifier
+
+
+def save_training_step(step: int, save_dir: Path) -> None:
+    """Save the current training step number to a file.
+
+    Args:
+        step: Current training step number.
+        save_dir: Directory where the step file will be saved.
+    """
+    write_json({"step": step}, save_dir / TRAINING_STEP)
+
+
+def load_training_step(save_dir: Path) -> int:
+    """Load the training step number from a file.
+
+    Args:
+        save_dir: Directory containing the step file.
+
+    Returns:
+        Training step number.
+    """
+    training_step = load_json(save_dir / TRAINING_STEP)
+    return training_step["step"]
+
+
+def update_last_checkpoint(checkpoint_dir: Path) -> Path:
+    """Update the symlink pointing to the last checkpoint.
+
+    Args:
+        checkpoint_dir: Path to the checkpoint directory to link.
+
+    Returns:
+        Path to the symlink that was created or updated.
+    """
+    last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK
+    if last_checkpoint_dir.is_symlink():
+        last_checkpoint_dir.unlink()
+    relative_target = checkpoint_dir.relative_to(checkpoint_dir.parent)
+    last_checkpoint_dir.symlink_to(relative_target)
+
+
+def save_checkpoint(
+    checkpoint_dir: Path,
+    step: int,
+    cfg: TrainPipelineConfig,
+) -> None:
+    """Save training checkpoint including config and RNG state.
+
+    Note: accelerate saves the model and training run. This method saves all
+    other auxiliary objects such as configs and RNG state.
+
+    Args:
+        checkpoint_dir: Directory where the checkpoint will be saved.
+        step: Current training step number.
+        cfg: The training config used for this run.
+    """
+    cfg.save_pretrained(checkpoint_dir)
+    save_training_step(step, checkpoint_dir)
+    save_rng_state(checkpoint_dir)
+
+
+def load_training_state(checkpoint_dir: Path) -> tuple[int, Optimizer, LRScheduler | None]:
+    """Load training state including step, optimizer, scheduler, and RNG state.
+
+    This is used to resume a training run. Note: optimizer and scheduler states
+    are loaded by accelerate, not by this function.
+
+    Args:
+        checkpoint_dir: The checkpoint directory. Should contain a 'training_state' dir.
+
+    Returns:
+        Tuple containing the training step number. Note: optimizer and scheduler
+        are loaded separately by accelerate.
+
+    Raises:
+        NotADirectoryError: If checkpoint_dir doesn't exist or is not a directory.
+    """
+    if not checkpoint_dir.is_dir():
+        raise NotADirectoryError(checkpoint_dir)
+
+    load_rng_state(checkpoint_dir)
+    step = load_training_step(checkpoint_dir)
+
+    return step
+
+
+def prune_old_checkpoints(latest_checkpoint_path: str) -> None:
+    """Delete all checkpoint directories except the specified one.
+
+    Recursively deletes all checkpoint directories in a parent folder except
+    for the specified one. This function is designed to clean up old model
+    checkpoints, preserving only the most recent one. It includes safety checks
+    to ensure it only deletes directories and handles potential filesystem errors.
+
+    Args:
+        latest_checkpoint_path: The full path to the checkpoint directory
+            that should be kept.
+    """
+    try:
+        latest_checkpoint = Path(latest_checkpoint_path).resolve()
+        parent_dir = latest_checkpoint.parent
+
+        if not parent_dir.is_dir():
+            logging.error(f"Parent directory '{parent_dir.resolve()}' does not exist. Aborting cleanup.")
+            return
+
+        if not latest_checkpoint.is_dir():
+            logging.warning(
+                f"Checkpoint '{latest_checkpoint.resolve()}' is not a valid directory. Aborting cleanup."
+            )
+            return
+
+        logging.info(
+            f"Starting cleanup in '{parent_dir.resolve()}'. Keeping checkpoint: '{latest_checkpoint.name}'"
+        )
+
+        # Iterate and delete other directories
+        for item in parent_dir.iterdir():
+            # Skip the checkpoint we want to keep and any files
+            if item.resolve() == latest_checkpoint.resolve() or not item.is_dir():
+                continue
+
+            try:
+                logging.info(f"Deleting old checkpoint directory: {item.name}")
+                shutil.rmtree(item)
+                logging.info(f"Successfully deleted {item.name}")
+            except OSError as e:
+                logging.error(f"Failed to delete '{item.name}'. Error: {e}")
+
+    except Exception as e:
+        logging.critical(f"An unexpected error occurred during checkpoint pruning setup: {e}")

opentau/utils/utils.py

@@ -0,0 +1,471 @@
+#!/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.
+"""General utility functions for device management, logging, and common operations.
+
+This module provides utilities for device selection, logging initialization,
+number formatting, platform-specific operations, and various helper functions
+used throughout the OpenTau codebase.
+"""
+
+import enum
+import inspect
+import logging
+import os
+import platform
+import warnings
+from copy import copy
+from dataclasses import fields, is_dataclass
+from datetime import datetime, timezone
+from functools import wraps
+from typing import Any, Callable
+
+import accelerate
+import numpy as np
+import torch
+
+
+def inside_slurm() -> bool:
+    """Check whether the Python process was launched through SLURM.
+
+    Returns:
+        True if the process is running in a SLURM environment, False otherwise.
+    """
+    # TODO(rcadene): return False for interactive mode `--pty bash`
+    return "SLURM_JOB_ID" in os.environ
+
+
+def auto_torch_device() -> torch.device:
+    """Automatically select the best available torch device.
+
+    Returns:
+        torch.device instance. Priority: CUDA > MPS > CPU.
+    """
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def get_safe_torch_device(try_device: str, log: bool = False, accelerator: Callable = None) -> torch.device:
+    """Given a string, return a torch.device with checks on whether the device is available."""
+    match try_device:
+        case "cuda":
+            assert torch.cuda.is_available()
+            device = accelerator.device if accelerator else torch.device("cuda")
+        case "mps":
+            assert torch.backends.mps.is_available()
+            device = torch.device("mps")
+        case "cpu":
+            device = torch.device("cpu")
+            if log:
+                logging.warning("Using CPU, this will be slow.")
+        case _:
+            device = torch.device(try_device)
+            if log:
+                logging.warning(f"Using custom {try_device} device.")
+
+    return device
+
+
+def get_safe_dtype(dtype: torch.dtype, device: str | torch.device) -> torch.dtype:
+    """Get a dtype that is compatible with the given device.
+
+    MPS is currently not compatible with float64, so this function converts
+    float64 to float32 for MPS devices.
+
+    Args:
+        dtype: Desired dtype.
+        device: Target device (string or torch.device).
+
+    Returns:
+        Compatible dtype for the device.
+    """
+    if isinstance(device, torch.device):
+        device = device.type
+    if device == "mps" and dtype == torch.float64:
+        return torch.float32
+    else:
+        return dtype
+
+
+def is_torch_device_available(try_device: str) -> bool:
+    """Check if a torch device is available.
+
+    Args:
+        try_device: Device name to check ('cuda', 'mps', or 'cpu').
+
+    Returns:
+        True if the device is available, False otherwise.
+
+    Raises:
+        ValueError: If try_device is not one of the recognized device types.
+    """
+    if try_device == "cuda":
+        return torch.cuda.is_available()
+    elif try_device == "mps":
+        return torch.backends.mps.is_available()
+    elif try_device == "cpu":
+        return True
+    else:
+        raise ValueError(f"Unknown device '{try_device}.")
+
+
+def is_amp_available(device: str) -> bool:
+    """Check if automatic mixed precision (AMP) is available for a device.
+
+    Args:
+        device: Device name to check ('cuda', 'mps', or 'cpu').
+
+    Returns:
+        True if AMP is available for the device, False otherwise.
+
+    Raises:
+        ValueError: If device is not one of the recognized device types.
+    """
+    if device in ["cuda", "cpu"]:
+        return True
+    elif device == "mps":
+        return False
+    else:
+        raise ValueError(f"Unknown device '{device}.")
+
+
+# Global variable to ensure logging is initialized only once
+_logging_init_stack = ""
+
+
+def _format_stack(stack: list[inspect.FrameInfo]) -> str:
+    return "\n".join(
+        f"  File '{frame.filename}', line {frame.lineno}, in {frame.function}"
+        for frame in stack[1:]  # skip the current frame
+    )
+
+
+def init_logging(accelerator: accelerate.Accelerator | None = None, level=logging.INFO) -> None:
+    """Initialize logging configuration with custom formatter.
+
+    This function sets up logging with a custom formatter that includes
+    timestamp, filename, and line number. It can only be initialized once
+    per process and will warn if called multiple times.
+
+    Args:
+        accelerator: Optional Accelerator instance. If provided, logging level
+            is set to WARNING on non-main processes to avoid duplicate logs.
+        level: Logging level to use. Defaults to logging.INFO.
+    """
+    global _logging_init_stack
+    stack = inspect.stack()
+
+    if _logging_init_stack:
+        warnings.warn(
+            f"""Logging was already initialized through the following stack:
+            {_logging_init_stack}
+            Not initializing again through the following stack:
+            {_format_stack(stack)}""",
+            stacklevel=2,
+        )
+    else:
+        _logging_init_stack = _format_stack(stack)
+
+    class CustomFormatter(logging.Formatter):
+        def format(self, record):
+            dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            fnameline = f"{record.pathname}:{record.lineno}"
+            return f"{record.levelname} {dt} {fnameline[-15:]:>15} {record.getMessage()}"
+
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(CustomFormatter())
+
+    logging.basicConfig(level=level, force=True, handlers=[console_handler])
+
+    if accelerator and not accelerator.is_main_process:
+        # Disable duplicate logging on non-main processes
+        logging.info(f"Setting logging level on non-main process {accelerator.process_index} to WARNING.")
+        logging.getLogger().setLevel(logging.WARNING)
+
+
+def format_big_number(num: float | int, precision: int = 0) -> str:
+    """Format a large number with appropriate suffix (K, M, B, T, Q).
+
+    Args:
+        num: Number to format.
+        precision: Number of decimal places. Defaults to 0.
+
+    Returns:
+        Formatted string with suffix (e.g., "1.5K", "2.3M").
+    """
+    suffixes = ["", "K", "M", "B", "T", "Q"]
+    divisor = 1000.0
+
+    for suffix in suffixes:
+        if abs(num) < divisor:
+            return f"{num:.{precision}f}{suffix}"
+        num /= divisor
+
+    return num
+
+
+def capture_timestamp_utc() -> datetime:
+    """Capture the current UTC timestamp.
+
+    Returns:
+        datetime object representing the current UTC time.
+    """
+    return datetime.now(timezone.utc)
+
+
+def say(text: str, blocking: bool = False) -> None:
+    """Use text-to-speech to speak text (platform-specific).
+
+    Args:
+        text: Text to speak.
+        blocking: If True, wait for speech to complete before returning.
+            Defaults to False.
+    """
+    # Check if mac, linux, or windows.
+    if platform.system() == "Darwin":
+        cmd = f'say "{text}"'
+        if not blocking:
+            cmd += " &"
+    elif platform.system() == "Linux":
+        cmd = f'spd-say "{text}"'
+        if blocking:
+            cmd += "  --wait"
+    elif platform.system() == "Windows":
+        # TODO(rcadene): Make blocking option work for Windows
+        cmd = (
+            'PowerShell -Command "Add-Type -AssemblyName System.Speech; '
+            f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')\""
+        )
+
+    os.system(cmd)  # nosec: B605
+
+
+def log_say(text: str, play_sounds: bool, blocking: bool = False) -> None:
+    """Log text and optionally speak it using text-to-speech.
+
+    Args:
+        text: Text to log and optionally speak.
+        play_sounds: If True, also speak the text using text-to-speech.
+        blocking: If True, wait for speech to complete. Defaults to False.
+    """
+    logging.info(text)
+
+    if play_sounds:
+        say(text, blocking)
+
+
+def get_channel_first_image_shape(image_shape: tuple) -> tuple:
+    """Convert image shape from HWC to CHW format if needed.
+
+    Args:
+        image_shape: Image shape tuple, either (H, W, C) or (C, H, W).
+
+    Returns:
+        Image shape in CHW format (C, H, W).
+
+    Raises:
+        ValueError: If the input shape is not in a recognized format.
+    """
+    shape = copy(image_shape)
+    if shape[2] < shape[0] and shape[2] < shape[1]:  # (h, w, c) -> (c, h, w)
+        shape = (shape[2], shape[0], shape[1])
+    elif not (shape[0] < shape[1] and shape[0] < shape[2]):
+        raise ValueError(image_shape)
+
+    return shape
+
+
+def has_method(cls: object, method_name: str) -> bool:
+    """Check if a class or object has a specific method.
+
+    Args:
+        cls: Class or object to check.
+        method_name: Name of the method to check for.
+
+    Returns:
+        True if the method exists and is callable, False otherwise.
+    """
+    return hasattr(cls, method_name) and callable(getattr(cls, method_name))
+
+
+def is_valid_numpy_dtype_string(dtype_str: str) -> bool:
+    """Check if a string can be converted to a numpy dtype.
+
+    Args:
+        dtype_str: String to check.
+
+    Returns:
+        True if the string is a valid numpy dtype, False otherwise.
+    """
+    try:
+        # Attempt to convert the string to a numpy dtype
+        np.dtype(dtype_str)
+        return True
+    except TypeError:
+        # If a TypeError is raised, the string is not a valid dtype
+        return False
+
+
+def is_launched_with_accelerate() -> bool:
+    """Check if the process was launched with accelerate.
+
+    Returns:
+        True if ACCELERATE_MIXED_PRECISION is in the environment, False otherwise.
+    """
+    return "ACCELERATE_MIXED_PRECISION" in os.environ
+
+
+def attempt_torch_compile(fn: callable, device_hint=None) -> callable:
+    """Attempt to compile a PyTorch function using torch.compile.
+
+    The argument device_hint is used to check if torch.compile works reliably
+    on the device. Compilation is skipped if the device is MPS (Metal Performance
+    Shaders) as it is experimental.
+
+    Args:
+        fn: Function to compile.
+        device_hint: Optional device hint to check compatibility.
+
+    Returns:
+        Compiled function if compilation succeeds, otherwise the original function.
+    """
+    if device_hint and "mps" in str(device_hint):
+        logging.warning("torch.compile is experimental on MPS devices. Compilation skipped.")
+        return fn
+
+    if hasattr(torch, "compile"):
+        logging.info("Attempting to compile the policy with torch.compile()...")
+        try:
+            # Other options: "default", "max-autotune" (longer compile time)
+            fn = torch.compile(fn)
+            logging.info("Policy compiled successfully.")
+        except Exception as e:
+            logging.warning(f"torch.compile failed with error: {e}. Proceeding without compilation.")
+    else:
+        logging.warning(
+            "torch.compile is not available. Requires PyTorch 2.0+. Proceeding without compilation."
+        )
+
+    return fn
+
+
+def create_dummy_observation(cfg, device, dtype=torch.bfloat16) -> dict:
+    """Create a dummy observation dictionary for testing or initialization.
+
+    Args:
+        cfg: Configuration object with num_cams, resolution, max_state_dim,
+            and action_chunk attributes.
+        device: Device to create tensors on.
+        dtype: Data type for tensors. Defaults to torch.bfloat16.
+
+    Returns:
+        Dictionary containing dummy camera observations, state, prompt, and
+        padding flags.
+    """
+    camera_observations = {
+        f"camera{i}": torch.zeros((1, 3, *cfg.resolution), dtype=dtype, device=device)
+        for i in range(cfg.num_cams)
+    }
+    return {
+        **camera_observations,
+        "state": torch.zeros((1, cfg.max_state_dim), dtype=dtype, device=device),
+        "prompt": ["Pick up yellow lego block and put it in the bin"],
+        "img_is_pad": torch.zeros((1, cfg.num_cams), dtype=torch.bool, device=device),
+        "action_is_pad": torch.zeros((1, cfg.action_chunk), dtype=torch.bool, device=device),
+    }
+
+
+def encode_accelerator_state_dict(obj) -> Any:
+    """Encode an object into a JSON/YAML-compatible primitive type.
+
+    Args:
+        obj: Object to encode (can be Enum, dict, list, tuple, dataclass, etc.).
+
+    Returns:
+        Encoded object with all nested structures converted to primitives.
+    """
+    if isinstance(obj, enum.Enum):
+        return encode_accelerator_state_dict(obj.value)
+    elif isinstance(obj, (str, int, float, bool)) or obj is None:
+        return obj
+    elif isinstance(obj, (list, tuple)):
+        return [encode_accelerator_state_dict(item) for item in obj]
+    elif isinstance(obj, dict):
+        return {key.replace(".", "_"): encode_accelerator_state_dict(value) for key, value in obj.items()}
+    elif is_dataclass(obj):
+        return {f.name: encode_accelerator_state_dict(getattr(obj, f.name)) for f in fields(obj)}
+    else:
+        return str(obj)  # Fallback to string representation for unsupported types
+
+
+def on_accelerate_main_proc(*, local=False, _sync=False):
+    r"""Returns a decorator to run a function only on the main process when using `accelerate`.
+
+    If `local` is True (defaults to False), the function will run on the main process of each node
+    (useful for multi-node setups).
+    If `_sync` is True (defaults to False), the output of the function will be broadcasted to all processes.
+    If `_sync` is True, you must ensure that all processes call the decorated function, otherwise it will deadlock.
+
+    YOU SHOULD BE EXTREMELY CAREFUL WHEN USING THIS DECORATOR with _sync=True. Consider the following example::
+
+        @on_accelerate_main_proc()
+        def f():
+            return g()
+
+        @on_accelerate_main_proc(_sync=True)
+        def g():
+            return 42
+
+    In this case, if f() is called on all processes, they will deadlock at g() because child processes don't even
+    enter f(), hence never call g(), and thus won't reach the broadcast.
+
+    Another example::
+
+        @on_accelerate_main_proc(_sync=cond())
+        def f():
+            print("hi")
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            state = accelerate.state.PartialState()
+            if not is_launched_with_accelerate() or not state.use_distributed:
+                return func(*args, **kwargs)
+
+            output, exception = None, None
+            flag = state.is_local_main_process if local else state.is_main_process
+            if flag:
+                try:
+                    output = func(*args, **kwargs)
+                except Exception as e:
+                    exception = e
+
+            if _sync:
+                payload = [output, exception]
+                accelerate.utils.broadcast_object_list(payload, from_process=0)
+                output, exception = payload
+
+            if exception is not None:
+                raise RuntimeError("An exception occurred in the main process.") from exception
+            return output
+
+        return wrapper
+
+    return decorator