opentau.utils.utils

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.

Functions

`attempt_torch_compile`(fn[, device_hint])	Attempt to compile a PyTorch function using torch.compile.
`auto_torch_device`()	Automatically select the best available torch device.
`capture_timestamp_utc`()	Capture the current UTC timestamp.
`create_dummy_observation`(cfg, device[, dtype])	Create a dummy observation dictionary for testing or initialization.
`encode_accelerator_state_dict`(obj)	Encode an object into a JSON/YAML-compatible primitive type.
`format_big_number`(num[, precision])	Format a large number with appropriate suffix (K, M, B, T, Q).
`get_channel_first_image_shape`(image_shape)	Convert image shape from HWC to CHW format if needed.
`get_safe_dtype`(dtype, device)	Get a dtype that is compatible with the given device.
`get_safe_torch_device`(try_device[, log, ...])	Given a string, return a torch.device with checks on whether the device is available.
`has_method`(cls, method_name)	Check if a class or object has a specific method.
`init_logging`([accelerator, level])	Initialize logging configuration with custom formatter.
`inside_slurm`()	Check whether the Python process was launched through SLURM.
`is_amp_available`(device)	Check if automatic mixed precision (AMP) is available for a device.
`is_launched_with_accelerate`()	Check if the process was launched with accelerate.
`is_torch_device_available`(try_device)	Check if a torch device is available.
`is_valid_numpy_dtype_string`(dtype_str)	Check if a string can be converted to a numpy dtype.
`log_say`(text, play_sounds[, blocking])	Log text and optionally speak it using text-to-speech.
`on_accelerate_main_proc`(*[, local, _sync])	Returns a decorator to run a function only on the main process when using accelerate.
`say`(text[, blocking])	Use text-to-speech to speak text (platform-specific).

opentau.utils.utils.attempt_torch_compile(fn: callable, device_hint=None) → callable[source]

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.

Parameters:

fn – Function to compile.
device_hint – Optional device hint to check compatibility.

Returns:

Compiled function if compilation succeeds, otherwise the original function.

opentau.utils.utils.auto_torch_device() → device[source]

Automatically select the best available torch device.

Returns:: CUDA > MPS > CPU.
Return type:: torch.device instance. Priority

opentau.utils.utils.capture_timestamp_utc() → datetime[source]

Capture the current UTC timestamp.

Returns:: datetime object representing the current UTC time.

opentau.utils.utils.create_dummy_observation(cfg, device, dtype=torch.bfloat16) → dict[source]

Create a dummy observation dictionary for testing or initialization.

Parameters:

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.

opentau.utils.utils.encode_accelerator_state_dict(obj) → Any[source]

Encode an object into a JSON/YAML-compatible primitive type.

Parameters:: obj – Object to encode (can be Enum, dict, list, tuple, dataclass, etc.).
Returns:: Encoded object with all nested structures converted to primitives.

opentau.utils.utils.format_big_number(num: float | int, precision: int = 0) → str[source]

Format a large number with appropriate suffix (K, M, B, T, Q).

Parameters:

num – Number to format.
precision – Number of decimal places. Defaults to 0.

Returns:

Formatted string with suffix (e.g., “1.5K”, “2.3M”).

opentau.utils.utils.get_channel_first_image_shape(image_shape: tuple) → tuple[source]

Convert image shape from HWC to CHW format if needed.

Parameters:: 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.

opentau.utils.utils.get_safe_dtype(dtype: dtype, device: str | device) → dtype[source]

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.

Parameters:

dtype – Desired dtype.
device – Target device (string or torch.device).

Returns:

Compatible dtype for the device.

opentau.utils.utils.get_safe_torch_device(try_device: str, log: bool = False, accelerator: Callable | None = None) → device[source]: Given a string, return a torch.device with checks on whether the device is available.

opentau.utils.utils.has_method(cls: object, method_name: str) → bool[source]

Check if a class or object has a specific method.

Parameters:

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.

opentau.utils.utils.init_logging(accelerator: Accelerator | None = None, level=20) → None[source]

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.

Parameters:

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.

opentau.utils.utils.inside_slurm() → bool[source]

Check whether the Python process was launched through SLURM.

Returns:: True if the process is running in a SLURM environment, False otherwise.

opentau.utils.utils.is_amp_available(device: str) → bool[source]

Check if automatic mixed precision (AMP) is available for a device.

Parameters:: 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.

opentau.utils.utils.is_launched_with_accelerate() → bool[source]

Check if the process was launched with accelerate.

Returns:: True if ACCELERATE_MIXED_PRECISION is in the environment, False otherwise.

opentau.utils.utils.is_torch_device_available(try_device: str) → bool[source]

Check if a torch device is available.

Parameters:: 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.

opentau.utils.utils.is_valid_numpy_dtype_string(dtype_str: str) → bool[source]

Check if a string can be converted to a numpy dtype.

Parameters:: dtype_str – String to check.
Returns:: True if the string is a valid numpy dtype, False otherwise.

opentau.utils.utils.log_say(text: str, play_sounds: bool, blocking: bool = False) → None[source]

Log text and optionally speak it using text-to-speech.

Parameters:

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.

opentau.utils.utils.on_accelerate_main_proc(*, local=False, _sync=False)[source]

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")

opentau.utils.utils.say(text: str, blocking: bool = False) → None[source]

Use text-to-speech to speak text (platform-specific).

Parameters:

text – Text to speak.
blocking – If True, wait for speech to complete before returning. Defaults to False.