torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/_utils.py

@@ -0,0 +1,1431 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import collections
+import functools
+import inspect
+import logging
+import math
+import os
+import pickle
+import sys
+import threading
+import time
+import traceback
+import warnings
+from collections.abc import Callable
+from contextlib import nullcontext
+from functools import wraps
+from textwrap import indent
+from typing import Any, cast, TypeVar
+
+import numpy as np
+import torch
+
+from pyvers import implement_for  # noqa: F401
+from tensordict import unravel_key
+from tensordict.utils import NestedKey
+from torch import multiprocessing as mp, Tensor
+from torch.autograd.profiler import record_function
+
+try:
+    from torch.compiler import is_compiling
+except ImportError:
+    from torch._dynamo import is_compiling
+
+
+def _get_default_mp_start_method() -> str:
+    """Returns TorchRL's preferred multiprocessing start method.
+
+    If the user has explicitly set a global start method via ``mp.set_start_method()``,
+    that method is returned. Otherwise, defaults to ``"spawn"`` for improved safety
+    across backends and to avoid known issues with ``fork`` in multi-threaded programs.
+    """
+    # Check if user has explicitly set a global start method
+    try:
+        current = mp.get_start_method(allow_none=True)
+        if current is not None:
+            return current
+    except (TypeError, RuntimeError):
+        pass
+    return "spawn"
+
+
+def _get_mp_ctx(start_method: str | None = None):
+    """Return a multiprocessing context with TorchRL's preferred start method.
+
+    This is intentionally context-based (instead of relying on global
+    ``mp.set_start_method``) so that TorchRL components can consistently allocate
+    primitives (Queue/Pipe/Lock/Process) with a matching context.
+    """
+    if start_method is None:
+        start_method = _get_default_mp_start_method()
+    try:
+        return mp.get_context(start_method)
+    except ValueError:
+        # Best effort fallback if a start method isn't supported on this platform.
+        return mp.get_context("spawn")
+
+
+def _set_mp_start_method_if_unset(start_method: str | None = None) -> str | None:
+    """Set the global start method only if it hasn't been set yet.
+
+    Returns the (possibly pre-existing) start method, or ``None`` if it cannot be
+    determined.
+    """
+    if start_method is None:
+        start_method = _get_default_mp_start_method()
+
+    current = None
+    try:
+        current = mp.get_start_method(allow_none=True)
+    except TypeError:
+        # Older python/torch wrappers may not accept allow_none.
+        try:
+            current = mp.get_start_method()
+        except Exception:
+            current = None
+    except Exception:
+        current = None
+
+    if current is None:
+        try:
+            mp.set_start_method(start_method, force=False)
+            current = start_method
+        except Exception:
+            # If another library already touched the context, we should not
+            # override it here.
+            pass
+    return current
+
+
+@implement_for("torch", None, "2.8")
+def _mp_sharing_strategy_for_spawn() -> str | None:
+    # On older torch stacks, pickling Process objects for "spawn" can end up
+    # passing file descriptors for shared storages; using "file_system" reduces
+    # FD passing and avoids spawn-time failures on some old Python versions.
+    return "file_system"
+
+
+@implement_for("torch", "2.8")
+def _mp_sharing_strategy_for_spawn() -> str | None:  # noqa: F811
+    return None
+
+
+def strtobool(val: Any) -> bool:
+    """Convert a string representation of truth to a boolean.
+
+    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'.
+    Raises ValueError if 'val' is anything else.
+    """
+    val = val.lower()
+    if val in ("y", "yes", "t", "true", "on", "1"):
+        return True
+    if val in ("n", "no", "f", "false", "off", "0"):
+        return False
+    raise ValueError(f"Invalid truth value {val!r}")
+
+
+LOGGING_LEVEL = os.environ.get("RL_LOGGING_LEVEL", "INFO")
+logger = logging.getLogger("torchrl")
+logger.setLevel(LOGGING_LEVEL)
+logger.propagate = False
+# Clear existing handlers
+while logger.hasHandlers():
+    logger.removeHandler(logger.handlers[0])
+stream_handlers = {
+    "stdout": sys.stdout,
+    "stderr": sys.stderr,
+}
+TORCHRL_CONSOLE_STREAM = os.getenv("TORCHRL_CONSOLE_STREAM")
+stream_handler = stream_handlers.get(TORCHRL_CONSOLE_STREAM, sys.stdout)
+
+
+# Create colored handler
+class _CustomFormatter(logging.Formatter):
+    def format(self, record):
+        # Format the initial part in green
+        green_format = "\033[92m%(asctime)s [%(name)s][%(levelname)s]\033[0m"
+        # Format the message part
+        message_format = "%(message)s"
+        # End marker in green
+        end_marker = "\033[92m [END]\033[0m"
+        # Combine all parts
+        formatted_message = logging.Formatter(
+            green_format + indent(message_format, " " * 4) + end_marker
+        ).format(record)
+
+        return formatted_message
+
+
+console_handler = logging.StreamHandler(stream=stream_handler)
+console_handler.setFormatter(_CustomFormatter())
+logger.addHandler(console_handler)
+
+console_handler.setLevel(LOGGING_LEVEL)
+logger.debug(f"Logging level: {logger.getEffectiveLevel()}")
+
+VERBOSE = strtobool(os.environ.get("VERBOSE", str(logger.isEnabledFor(logging.DEBUG))))
+_os_is_windows = sys.platform == "win32"
+RL_WARNINGS = strtobool(os.environ.get("RL_WARNINGS", "1"))
+if RL_WARNINGS:
+    warnings.filterwarnings("once", category=DeprecationWarning, module="torchrl")
+
+BATCHED_PIPE_TIMEOUT = float(os.environ.get("BATCHED_PIPE_TIMEOUT", "10000.0"))
+WEIGHT_SYNC_TIMEOUT = float(os.environ.get("WEIGHT_SYNC_TIMEOUT", "60.0"))
+
+_TORCH_DTYPES = (
+    torch.bfloat16,
+    torch.bool,
+    torch.complex128,
+    torch.complex32,
+    torch.complex64,
+    torch.float16,
+    torch.float32,
+    torch.float64,
+    torch.int16,
+    torch.int32,
+    torch.int64,
+    torch.int8,
+    torch.qint32,
+    torch.qint8,
+    torch.quint4x2,
+    torch.quint8,
+    torch.uint8,
+)
+if hasattr(torch, "uint16"):
+    _TORCH_DTYPES = _TORCH_DTYPES + (torch.uint16,)
+if hasattr(torch, "uint32"):
+    _TORCH_DTYPES = _TORCH_DTYPES + (torch.uint32,)
+if hasattr(torch, "uint64"):
+    _TORCH_DTYPES = _TORCH_DTYPES + (torch.uint64,)
+_STR_DTYPE_TO_DTYPE = {str(dtype): dtype for dtype in _TORCH_DTYPES}
+_STRDTYPE2DTYPE = _STR_DTYPE_TO_DTYPE
+_DTYPE_TO_STR_DTYPE = {
+    dtype: str_dtype for str_dtype, dtype in _STR_DTYPE_TO_DTYPE.items()
+}
+_DTYPE2STRDTYPE = _STR_DTYPE_TO_DTYPE
+
+
+class timeit:
+    """A dirty but easy to use decorator for profiling code.
+
+    Args:
+        name (str): The name of the timer.
+
+    Examples:
+        >>> from torchrl import timeit
+        >>> @timeit("my_function")
+        >>> def my_function():
+            ...
+        >>> my_function()
+        >>> with timeit("my_other_function"):
+        ...     my_other_function()
+        >>> timeit.print()  # prints the state of the timer for each function
+
+        The timer can also be queried mid-execution using the :meth:`elapsed` method:
+
+        >>> with timeit("my_function") as timer:
+        ...     # do some work
+        ...     print(f"Elapsed so far: {timer.elapsed():.3f}s")
+        ...     # do more work
+
+        For long-running processes where a context manager isn't practical,
+        use the :meth:`start` method:
+
+        >>> timer = timeit("long_process").start()
+        >>> for i in range(100):
+        ...     # do work
+        ...     print(f"Elapsed: {timer.elapsed():.3f}s")
+    """
+
+    _REG = {}
+
+    def __init__(self, name):
+        self.name = name
+
+    def __call__(self, fn: Callable) -> Callable:
+        @wraps(fn)
+        def decorated_fn(*args, **kwargs):
+            with self:
+                out = fn(*args, **kwargs)
+                return out
+
+        return decorated_fn
+
+    def __enter__(self) -> timeit:
+        self.t0 = time.time()
+        return self
+
+    def start(self) -> timeit:
+        """Starts the timer without using a context manager.
+
+        This is useful when you need to track elapsed time over a long-running
+        loop or process where a context manager isn't practical.
+
+        Returns:
+            timeit: Returns self for method chaining.
+
+        Examples:
+            >>> timer = timeit("my_long_process").start()
+            >>> for i in range(100):
+            ...     # do work
+            ...     if i % 10 == 0:
+            ...         print(f"Elapsed: {timer.elapsed():.3f}s")
+        """
+        self.t0 = time.time()
+        return self
+
+    def elapsed(self) -> float:
+        """Returns the elapsed time in seconds since the timer was started.
+
+        This can be called during execution to query the current elapsed time.
+
+        Returns:
+            float: Elapsed time in seconds.
+
+        Examples:
+            >>> with timeit("my_function") as timer:
+            ...     # do some work
+            ...     print(f"Elapsed so far: {timer.elapsed():.3f}s")
+            ...     # do more work
+        """
+        return time.time() - self.t0
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        t = self.elapsed()
+        val = self._REG.setdefault(self.name, [0.0, 0.0, 0])
+
+        count = val[2]
+        N = count + 1
+        val[0] = val[0] * (count / N) + t / N
+        val[1] += t
+        val[2] = N
+
+    @staticmethod
+    def print(prefix: str | None = None) -> str:  # noqa: T202
+        """Prints the state of the timer.
+
+        Args:
+            prefix (str): The prefix to add to the keys. If `None`, no prefix is added.
+
+        Returns:
+            the string printed using the logger.
+        """
+        keys = list(timeit._REG)
+        keys.sort()
+        string = []
+        for name in keys:
+            strings = []
+            if prefix:
+                strings.append(prefix)
+            strings.append(
+                f"{name} took {timeit._REG[name][0] * 1000:4.4f} msec (total = {timeit._REG[name][1]: 4.4f} sec since last reset)."
+            )
+            string.append(" -- ".join(strings))
+            logger.info(string[-1])
+        return "\n".join(string)
+
+    _printevery_count = 0
+
+    @classmethod
+    def printevery(
+        cls,
+        num_prints: int,
+        total_count: int,
+        *,
+        prefix: str | None = None,
+        erase: bool = False,
+    ) -> None:
+        """Prints the state of the timer at regular intervals.
+
+        Args:
+            num_prints (int): The number of times to print the state of the timer, given the total_count.
+            total_count (int): The total number of times to print the state of the timer.
+            prefix (str): The prefix to add to the keys. If `None`, no prefix is added.
+            erase (bool): If True, erase the timer after printing. Default is `False`.
+
+        """
+        interval = max(1, total_count // num_prints)
+        if cls._printevery_count % interval == 0:
+            cls.print(prefix=prefix)
+            if erase:
+                cls.erase()
+        cls._printevery_count += 1
+
+    @classmethod
+    def todict(
+        cls, percall: bool = True, prefix: str | None = None
+    ) -> dict[str, float]:
+        """Convert the timer to a dictionary.
+
+        Args:
+            percall (bool): If True, return the average time per call.
+            prefix (str): The prefix to add to the keys.
+        """
+
+        def _make_key(key):
+            if prefix:
+                return f"{prefix}/{key}"
+            return key
+
+        if percall:
+            return {_make_key(key): val[0] for key, val in cls._REG.items()}
+        return {_make_key(key): val[1] for key, val in cls._REG.items()}
+
+    @staticmethod
+    def erase():
+        """Erase the timer.
+
+        .. seealso:: :meth:`reset`
+        """
+        for k in timeit._REG:
+            timeit._REG[k] = [0.0, 0.0, 0]
+
+    @classmethod
+    def reset(cls):
+        """Reset the timer.
+
+        .. seealso:: :meth:`erase`
+        """
+        cls.erase()
+
+
+# Global flag to enable detailed profiling instrumentation.
+# When False (default), _maybe_record_function returns nullcontext() immediately
+# to avoid overhead in hot code paths.
+_PROFILING_ENABLED = False
+
+# Singleton nullcontext to avoid repeated object creation
+_NULL_CONTEXT = nullcontext()
+
+
+def set_profiling_enabled(enabled: bool) -> None:
+    """Enable or disable detailed profiling instrumentation.
+
+    When disabled (default), `_maybe_record_function` and `_maybe_timeit`
+    return immediately with minimal overhead. Enable only when actively
+    profiling to avoid performance regression.
+
+    Args:
+        enabled: If True, enable profiling instrumentation.
+    """
+    global _PROFILING_ENABLED
+    _PROFILING_ENABLED = enabled
+
+
+def _maybe_timeit(name):
+    """Return timeit context if not compiling, nullcontext otherwise.
+
+    torch.compiler.is_compiling() returns True when inside a compiled region,
+    and timeit uses time.time() which dynamo cannot trace.
+    """
+    if is_compiling():
+        return _NULL_CONTEXT
+    return timeit(name)
+
+
+def _maybe_record_function(name):
+    """Return record_function context if profiling enabled and not compiling.
+
+    When _PROFILING_ENABLED is False (default), returns immediately with
+    minimal overhead to avoid performance regression in hot code paths.
+    """
+    if not _PROFILING_ENABLED:
+        return _NULL_CONTEXT
+    if is_compiling():
+        return _NULL_CONTEXT
+
+    return record_function(name)
+
+
+def _maybe_record_function_decorator(name: str) -> Callable[[Callable], Callable]:
+    """Decorator version of :func:`_maybe_record_function`.
+
+    This is preferred over sprinkling many context managers in hot code paths,
+    as it reduces Python overhead while keeping a useful profiler structure.
+
+    When _PROFILING_ENABLED is False (default), the decorator is a no-op.
+    """
+
+    def decorator(fn: Callable) -> Callable:
+        @wraps(fn)
+        def wrapped(*args, **kwargs):
+            if not _PROFILING_ENABLED:
+                return fn(*args, **kwargs)
+            with _maybe_record_function(name):
+                return fn(*args, **kwargs)
+
+        return wrapped
+
+    return decorator
+
+
+def _check_for_faulty_process(processes):
+    terminate = False
+    for p in processes:
+        if not p._closed and not p.is_alive():
+            terminate = True
+            for _p in processes:
+                _p: mp.Process
+                if not _p._closed and _p.is_alive():
+                    try:
+                        _p.terminate()
+                    except Exception:
+                        _p.kill()
+                    finally:
+                        time.sleep(0.1)
+                        _p.close()
+            if terminate:
+                break
+    if terminate:
+        raise RuntimeError(
+            "At least one process failed. Check for more infos in the log."
+        )
+
+
+def seed_generator(seed):
+    """A seed generator function.
+
+    Given a seeding integer, generates a deterministic next seed to be used in a
+    seeding sequence.
+
+    Args:
+        seed (int): initial seed.
+
+    Returns: Next seed of the chain.
+
+    """
+    max_seed_val = (
+        2**32 - 1
+    )  # https://discuss.pytorch.org/t/what-is-the-max-seed-you-can-set-up/145688
+    rng = np.random.default_rng(seed)
+    seed = int.from_bytes(rng.bytes(8), "big")
+    return seed % max_seed_val
+
+
+class KeyDependentDefaultDict(collections.defaultdict):
+    """A key-dependent default dict.
+
+    Examples:
+        >>> my_dict = KeyDependentDefaultDict(lambda key: "foo_" + key)
+        >>> print(my_dict["bar"])
+        foo_bar
+    """
+
+    def __init__(self, fun):
+        self.fun = fun
+        super().__init__()
+
+    def __missing__(self, key):
+        value = self.fun(key)
+        self[key] = value
+        return value
+
+
+def prod(sequence):
+    """General prod function, that generalised usage across math and np.
+
+    Created for multiple python versions compatibility).
+
+    """
+    if hasattr(math, "prod"):
+        return math.prod(sequence)
+    else:
+        return int(np.prod(sequence))
+
+
+def get_binary_env_var(key):
+    """Parses and returns the binary environment variable value.
+
+    If not present in environment, it is considered `False`.
+
+    Args:
+        key (str): name of the environment variable.
+    """
+    val = os.environ.get(key, "False")
+    if val in ("0", "False", "false"):
+        val = False
+    elif val in ("1", "True", "true"):
+        val = True
+    else:
+        raise ValueError(
+            f"Environment variable {key} should be in 'True', 'False', '0' or '1'. "
+            f"Got {val} instead."
+        )
+    return val
+
+
+class _Dynamic_CKPT_BACKEND:
+    """Allows CKPT_BACKEND to be changed on-the-fly."""
+
+    backends = ["torch", "torchsnapshot"]
+
+    def _get_backend(self):
+        backend = os.environ.get("CKPT_BACKEND", "torch")
+        if backend == "torchsnapshot":
+            try:
+                import torchsnapshot  # noqa: F401
+            except ImportError as err:
+                raise ImportError(
+                    f"torchsnapshot not found, but the backend points to this library. "
+                    f"Consider installing torchsnapshot or choose another backend (available backends: {self.backends})"
+                ) from err
+        return backend
+
+    def __getattr__(self, item):
+        return getattr(self._get_backend(), item)
+
+    def __eq__(self, other):
+        return self._get_backend() == other
+
+    def __ne__(self, other):
+        return self._get_backend() != other
+
+    def __repr__(self):
+        return self._get_backend()
+
+
+_CKPT_BACKEND = _Dynamic_CKPT_BACKEND()
+
+
+def accept_remote_rref_invocation(func):
+    """Decorator that allows a method to be invoked remotely.
+
+    Passes the `rpc.RRef` associated with the remote object construction as first argument in place of the object reference.
+
+    """
+
+    @wraps(func)
+    def unpack_rref_and_invoke_function(self, *args, **kwargs):
+        # windows does not know torch._C._distributed_rpc.PyRRef
+        if not _os_is_windows and isinstance(self, torch._C._distributed_rpc.PyRRef):
+            self = self.local_value()
+        return func(self, *args, **kwargs)
+
+    return unpack_rref_and_invoke_function
+
+
+def accept_remote_rref_udf_invocation(decorated_class):
+    """Class decorator that applies `accept_remote_rref_invocation` to all public methods."""
+    # ignores private methods
+    for name in dir(decorated_class):
+        method = getattr(decorated_class, name, None)
+        if method is None:
+            continue
+        if callable(method) and not name.startswith("_"):
+            setattr(decorated_class, name, accept_remote_rref_invocation(method))
+    return decorated_class
+
+
+# We copy this from torch as older versions do not have it
+# see torch.utils._contextlib
+
+# Extra utilities for working with context managers that should have been
+# in the standard library but are not
+
+# Used for annotating the decorator usage of _DecoratorContextManager (e.g.,
+# 'no_grad' and 'enable_grad').
+# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators
+FuncType = Callable[..., Any]
+F = TypeVar("F", bound=FuncType)
+
+
+def _wrap_generator(ctx_factory, func):
+    """Wrap each generator invocation with the context manager factory.
+
+    The input should be a function that returns a context manager,
+    not a context manager itself, to handle one-shot context managers.
+    """
+
+    @functools.wraps(func)
+    def generator_context(*args, **kwargs):
+        gen = func(*args, **kwargs)
+
+        # Generators are suspended and unsuspended at `yield`, hence we
+        # make sure the grad mode is properly set every time the execution
+        # flow returns into the wrapped generator and restored when it
+        # returns through our `yield` to our caller (see PR #49017).
+        try:
+            # Issuing `None` to a generator fires it up
+            with ctx_factory():
+                response = gen.send(None)
+
+            while True:
+                try:
+                    # Forward the response to our caller and get its next request
+                    request = yield response
+
+                except GeneratorExit:
+                    # Inform the still active generator about its imminent closure
+                    with ctx_factory():
+                        gen.close()
+                    raise
+
+                except BaseException:
+                    # Propagate the exception thrown at us by the caller
+                    with ctx_factory():
+                        response = gen.throw(*sys.exc_info())
+
+                else:
+                    # Pass the last request to the generator and get its response
+                    with ctx_factory():
+                        response = gen.send(request)
+
+        # We let the exceptions raised above by the generator's `.throw` or
+        # `.send` methods bubble up to our caller, except for StopIteration
+        except StopIteration as e:
+            # The generator informed us that it is done: take whatever its
+            # returned value (if any) was and indicate that we're done too
+            # by returning it (see docs for python's return-statement).
+            return e.value
+
+    return generator_context
+
+
+def context_decorator(ctx, func):
+    """Context decorator.
+
+    Like contextlib.ContextDecorator, but:
+
+    1. Is done by wrapping, rather than inheritance, so it works with context
+       managers that are implemented from C and thus cannot easily inherit from
+       Python classes
+    2. Wraps generators in the intuitive way (c.f. https://bugs.python.org/issue37743)
+    3. Errors out if you try to wrap a class, because it is ambiguous whether
+       or not you intended to wrap only the constructor
+
+    The input argument can either be a context manager (in which case it must
+    be a multi-shot context manager that can be directly invoked multiple times)
+    or a callable that produces a context manager.
+    """
+    if callable(ctx) and hasattr(ctx, "__enter__"):
+        raise RuntimeError(
+            f"Passed in {ctx} is both callable and also a valid context manager "
+            "(has __enter__), making it ambiguous which interface to use.  If you "
+            "intended to pass a context manager factory, rewrite your call as "
+            "context_decorator(lambda: ctx()); if you intended to pass a context "
+            "manager directly, rewrite your call as context_decorator(lambda: ctx)"
+        )
+
+    if not callable(ctx):
+
+        def ctx_factory():
+            return ctx
+
+    else:
+        ctx_factory = ctx
+
+    if inspect.isclass(func):
+        raise RuntimeError(
+            "Cannot decorate classes; it is ambiguous whether only the "
+            "constructor or all methods should have the context manager applied; "
+            "additionally, decorating a class at definition-site will prevent "
+            "use of the identifier as a conventional type.  "
+            "To specify which methods to decorate, decorate each of them "
+            "individually."
+        )
+
+    if inspect.isgeneratorfunction(func):
+        return _wrap_generator(ctx_factory, func)
+
+    @functools.wraps(func)
+    def decorate_context(*args, **kwargs):
+        with ctx_factory():
+            return func(*args, **kwargs)
+
+    return decorate_context
+
+
+class _DecoratorContextManager:
+    """Allow a context manager to be used as a decorator."""
+
+    def __call__(self, orig_func: F) -> F:
+        if inspect.isclass(orig_func):
+            warnings.warn(
+                "Decorating classes is deprecated and will be disabled in "
+                "future versions. You should only decorate functions or methods. "
+                "To preserve the current behavior of class decoration, you can "
+                "directly decorate the `__init__` method and nothing else."
+            )
+            func = cast(F, lambda *args, **kwargs: orig_func(*args, **kwargs))
+        else:
+            func = orig_func
+
+        return cast(F, context_decorator(self.clone, func))
+
+    def __enter__(self) -> None:
+        raise NotImplementedError
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        raise NotImplementedError
+
+    def clone(self):
+        # override this method if your children class takes __init__ parameters
+        return self.__class__()
+
+
+def get_trace():
+    """A simple debugging util to spot where a function is being called."""
+    traceback.print_stack()
+
+
+def _make_process_no_warn_cls(ctx=None):
+    """Create a _ProcessNoWarn class that inherits from the appropriate Process class.
+
+    When using multiprocessing contexts (e.g., fork or spawn), the Process class
+    used must match the context to ensure synchronization primitives like locks
+    work correctly. This factory function creates a _ProcessNoWarn class that
+    inherits from the context's Process class.
+
+    Args:
+        ctx: A multiprocessing context (e.g., from mp.get_context('fork')).
+            If None, uses the default mp.Process.
+
+    Returns:
+        A _ProcessNoWarn class that inherits from the appropriate Process base.
+
+    .. note::
+        For the "spawn" start method, this returns pre-defined module-level classes
+        to ensure they can be pickled correctly.
+    """
+    if ctx is None:
+        return _ProcessNoWarn
+
+    start_method = ctx.get_start_method()
+    if start_method == "fork":
+        return _ProcessNoWarnFork
+    elif start_method == "spawn":
+        return _ProcessNoWarnSpawn
+    elif start_method == "forkserver":
+        return _ProcessNoWarnForkserver
+    else:
+        # For unknown start methods, fall back to default
+        return _ProcessNoWarn
+
+
+# Keep the old class name as a default for backwards compatibility
+class _ProcessNoWarn(mp.Process):
+    """A private Process class that shuts down warnings on the subprocess and controls the number of threads in the subprocess.
+
+    .. note::
+        When using multiprocessing contexts with synchronization primitives (locks, etc.),
+        use :func:`_make_process_no_warn_cls` with the context to ensure compatibility.
+    """
+
+    @wraps(mp.Process.__init__)
+    def __init__(self, *args, num_threads=None, _start_method=None, **kwargs):
+        import torchrl
+
+        self.filter_warnings_subprocess = torchrl.filter_warnings_subprocess
+        self.num_threads = num_threads
+        if _start_method is not None:
+            self._start_method = _start_method
+        super().__init__(*args, **kwargs)
+
+    def run(self, *args, **kwargs):
+        if self.num_threads is not None:
+            torch.set_num_threads(self.num_threads)
+        if self.filter_warnings_subprocess:
+            import warnings
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                return mp.Process.run(self, *args, **kwargs)
+        return mp.Process.run(self, *args, **kwargs)
+
+
+# Pre-defined _ProcessNoWarn classes for different multiprocessing start methods.
+# These must be defined at module level to be picklable with the "spawn" start method.
+#
+# We use a mixin pattern to avoid code duplication while still having
+# distinct module-level classes that can be pickled.
+
+
+class _ProcessNoWarnMixin:
+    """Mixin class providing the common functionality for _ProcessNoWarn variants."""
+
+    def _init_process_no_warn(self, num_threads=None, _start_method=None):
+        import torchrl
+
+        self.filter_warnings_subprocess = torchrl.filter_warnings_subprocess
+        self.num_threads = num_threads
+        if _start_method is not None:
+            self._start_method = _start_method
+
+    def run(self, *args, **kwargs):
+        if self.num_threads is not None:
+            torch.set_num_threads(self.num_threads)
+        if self.filter_warnings_subprocess:
+            import warnings
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                return super().run(*args, **kwargs)
+        return super().run(*args, **kwargs)
+
+
+# Spawn-specific class (for macOS default and Windows)
+try:
+    _spawn_ctx = mp.get_context("spawn")
+
+    class _ProcessNoWarnSpawn(_ProcessNoWarnMixin, _spawn_ctx.Process):
+        """_ProcessNoWarn for the 'spawn' multiprocessing context."""
+
+        def __init__(self, *args, num_threads=None, _start_method=None, **kwargs):
+            self._init_process_no_warn(num_threads, _start_method)
+            super().__init__(*args, **kwargs)
+
+except ValueError:
+    _ProcessNoWarnSpawn = _ProcessNoWarn
+
+
+# Fork-specific class (for Linux default, not available on Windows)
+try:
+    _fork_ctx = mp.get_context("fork")
+
+    class _ProcessNoWarnFork(_ProcessNoWarnMixin, _fork_ctx.Process):
+        """_ProcessNoWarn for the 'fork' multiprocessing context."""
+
+        def __init__(self, *args, num_threads=None, _start_method=None, **kwargs):
+            self._init_process_no_warn(num_threads, _start_method)
+            super().__init__(*args, **kwargs)
+
+except ValueError:
+    _ProcessNoWarnFork = _ProcessNoWarn
+
+
+# Forkserver-specific class (not available on Windows)
+try:
+    _forkserver_ctx = mp.get_context("forkserver")
+
+    class _ProcessNoWarnForkserver(_ProcessNoWarnMixin, _forkserver_ctx.Process):
+        """_ProcessNoWarn for the 'forkserver' multiprocessing context."""
+
+        def __init__(self, *args, num_threads=None, _start_method=None, **kwargs):
+            self._init_process_no_warn(num_threads, _start_method)
+            super().__init__(*args, **kwargs)
+
+except ValueError:
+    _ProcessNoWarnForkserver = _ProcessNoWarn
+
+
+def print_directory_tree(path, indent="", display_metadata=True):
+    """Prints the directory tree starting from the specified path.
+
+    Args:
+        path (str): The path of the directory to print.
+        indent (str): The current indentation level for formatting.
+        display_metadata (bool): if ``True``, metadata of the dir will be
+            displayed too.
+
+    """
+    if display_metadata:
+
+        def get_directory_size(path="."):
+            total_size = 0
+
+            for dirpath, _, filenames in os.walk(path):
+                for filename in filenames:
+                    file_path = os.path.join(dirpath, filename)
+                    total_size += os.path.getsize(file_path)
+
+            return total_size
+
+        def format_size(size):
+            # Convert size to a human-readable format
+            for unit in ["B", "KB", "MB", "GB", "TB"]:
+                if size < 1024.0:
+                    return f"{size:.2f} {unit}"
+                size /= 1024.0
+
+        total_size_bytes = get_directory_size(path)
+        formatted_size = format_size(total_size_bytes)
+        logger.info(f"Directory size: {formatted_size}")
+
+    if os.path.isdir(path):
+        logger.info(indent + os.path.basename(path) + "/")
+        indent += "    "
+        for item in os.listdir(path):
+            print_directory_tree(
+                os.path.join(path, item), indent=indent, display_metadata=False
+            )
+    else:
+        logger.info(indent + os.path.basename(path))
+
+
+def _ends_with(key, match):
+    if isinstance(key, str):
+        return key == match
+    return key[-1] == match
+
+
+def _replace_last(key: NestedKey, new_ending: str) -> NestedKey:
+    if isinstance(key, str):
+        return new_ending
+    else:
+        return key[:-1] + (new_ending,)
+
+
+def _append_last(key: NestedKey, new_suffix: str) -> NestedKey:
+    key = unravel_key(key)
+    if isinstance(key, str):
+        return key + new_suffix
+    else:
+        return key[:-1] + (key[-1] + new_suffix,)
+
+
+class _rng_decorator(_DecoratorContextManager):
+    """Temporarily sets the seed and sets back the rng state when exiting."""
+
+    def __init__(self, seed, device=None):
+        self.seed = seed
+        self.device = device
+        self.has_cuda = torch.cuda.is_available()
+
+    def __enter__(self):
+        self._get_state()
+        torch.manual_seed(self.seed)
+
+    def _get_state(self):
+        if self.has_cuda:
+            if self.device is None:
+                self._state = (torch.random.get_rng_state(), torch.cuda.get_rng_state())
+            else:
+                self._state = (
+                    torch.random.get_rng_state(),
+                    torch.cuda.get_rng_state(self.device),
+                )
+
+        else:
+            self._state = torch.random.get_rng_state()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.has_cuda:
+            torch.random.set_rng_state(self._state[0])
+            if self.device is not None:
+                torch.cuda.set_rng_state(self._state[1], device=self.device)
+            else:
+                torch.cuda.set_rng_state(self._state[1])
+
+        else:
+            torch.random.set_rng_state(self._state)
+
+
+def _can_be_pickled(obj):
+    try:
+        pickle.dumps(obj)
+        return True
+    except (pickle.PickleError, AttributeError, TypeError):
+        return False
+
+
+def _make_ordinal_device(device: torch.device):
+    if device is None:
+        return device
+    device = torch.device(device)
+    if device.type == "cuda" and device.index is None:
+        return torch.device("cuda", index=torch.cuda.current_device())
+    if device.type == "mps" and device.index is None:
+        return torch.device("mps", index=0)
+    return device
+
+
+def get_available_device(return_str: bool = False) -> torch.device | str:
+    """Return the available accelerator device, or CPU if none is found.
+
+    Checks for accelerator availability in the following order: CUDA, NPU, MPS.
+    Returns the first available accelerator, or CPU if none are present.
+
+    .. note::
+        PyTorch generally assumes a single accelerator type per system.
+        Running with multiple accelerator types (e.g., both CUDA and NPU)
+        is not officially supported. This function simply returns the first
+        available accelerator it finds.
+
+    Args:
+        return_str: If ``True``, returns a string representation of the device
+            instead of a :class:`~torch.device` object. Defaults to ``False``.
+
+    Returns:
+        The available accelerator device as a :class:`~torch.device` object,
+        or as a string if ``return_str`` is ``True``. Falls back to CPU if
+        no accelerator is available.
+
+    Examples:
+        >>> from torchrl._utils import get_available_device
+        >>> device = get_available_device()
+        >>> # Use with config fallback:
+        >>> device = cfg.device or get_available_device()
+    """
+    if torch.cuda.is_available():
+        device = "cuda:0"
+    elif hasattr(torch, "npu") and torch.npu.is_available():
+        device = "npu:0"
+    elif torch.backends.mps.is_available():
+        device = "mps:0"
+    else:
+        device = "cpu"
+    if return_str:
+        return device
+    return torch.device(device)
+
+
+class _ContextManager:
+    def __init__(self):
+        self._mode: Any | None = None
+        self._lock = threading.Lock()
+
+    def get_mode(self) -> Any | None:
+        cm = self._lock if not is_compiling() else nullcontext()
+        with cm:
+            return self._mode
+
+    def set_mode(self, type: Any | None) -> None:
+        cm = self._lock if not is_compiling() else nullcontext()
+        with cm:
+            self._mode = type
+
+
+def _standardize(
+    input: Tensor,
+    exclude_dims: tuple[int] = (),
+    mean: Tensor | None = None,
+    std: Tensor | None = None,
+    eps: float | None = None,
+):
+    """Standardizes the input tensor with the possibility of excluding specific dims from the statistics.
+
+    Useful when processing multi-agent data to keep the agent dimensions independent.
+
+    Args:
+        input (Tensor): the input tensor to be standardized.
+        exclude_dims (Tuple[int]): dimensions to exclude from the statistics, can be negative. Default: ().
+        mean (Tensor): a mean to be used for standardization. Must be of shape broadcastable to input. Default: None.
+        std (Tensor): a standard deviation to be used for standardization. Must be of shape broadcastable to input. Default: None.
+        eps (:obj:`float`): epsilon to be used for numerical stability. Default: float32 resolution.
+
+    """
+    if eps is None:
+        if input.dtype.is_floating_point:
+            eps = torch.finfo(torch.float).resolution
+        else:
+            eps = 1e-6
+
+    len_exclude_dims = len(exclude_dims)
+    if not len_exclude_dims:
+        if mean is None:
+            mean = input.mean()
+        else:
+            # Assume dtypes are compatible
+            mean = torch.as_tensor(mean, device=input.device)
+        if std is None:
+            std = input.std()
+        else:
+            # Assume dtypes are compatible
+            std = torch.as_tensor(std, device=input.device)
+        return (input - mean) / std.clamp_min(eps)
+
+    input_shape = input.shape
+    exclude_dims = [
+        d if d >= 0 else d + len(input_shape) for d in exclude_dims
+    ]  # Make negative dims positive
+
+    if len(set(exclude_dims)) != len_exclude_dims:
+        raise ValueError("Exclude dims has repeating elements")
+    if any(dim < 0 or dim >= len(input_shape) for dim in exclude_dims):
+        raise ValueError(
+            f"exclude_dims={exclude_dims} provided outside bounds for input of shape={input_shape}"
+        )
+    if len_exclude_dims == len(input_shape):
+        warnings.warn(
+            "_standardize called but all dims were excluded from the statistics, returning unprocessed input"
+        )
+        return input
+
+    included_dims = tuple(d for d in range(len(input_shape)) if d not in exclude_dims)
+    if mean is None:
+        mean = torch.mean(input, keepdim=True, dim=included_dims)
+    if std is None:
+        std = torch.std(input, keepdim=True, dim=included_dims)
+    return (input - mean) / std.clamp_min(eps)
+
+
+@wraps(torch.compile)
+def compile_with_warmup(*args, warmup: int = 1, **kwargs):
+    """Compile a model with warm-up.
+
+    This function wraps :func:`~torch.compile` to add a warm-up phase. During the warm-up phase,
+    the original model is used. After the warm-up phase, the model is compiled using
+    `torch.compile`.
+
+    Args:
+        *args: Arguments to be passed to `torch.compile`.
+        warmup (int): Number of calls to the model before compiling it. Defaults to 1.
+        **kwargs: Keyword arguments to be passed to `torch.compile`.
+
+    Returns:
+        A callable that wraps the original model. If no model is provided, returns a
+        lambda function that takes a model as input and returns the wrapped model.
+
+    Notes:
+        If no model is provided, this function returns a lambda function that can be
+        used to wrap a model later. This allows for delayed compilation of the model.
+
+    Example:
+        >>> model = torch.nn.Linear(5, 3)
+        >>> compiled_model = compile_with_warmup(model, warmup=10)
+        >>> # First 10 calls use the original model
+        >>> # After 10 calls, the model is compiled and used
+    """
+    if len(args):
+        model = args[0]
+        args = ()
+    else:
+        model = kwargs.pop("model", None)
+    if model is None:
+        return lambda model: compile_with_warmup(model, warmup=warmup, **kwargs)
+    else:
+        count = -1
+        compiled_model = model
+
+        @wraps(model)
+        def count_and_compile(*model_args, **model_kwargs):
+            nonlocal count
+            nonlocal compiled_model
+            count += 1
+            if count == warmup:
+                compiled_model = torch.compile(model, *args, **kwargs)
+            return compiled_model(*model_args, **model_kwargs)
+
+        return count_and_compile
+
+
+# auto unwrap control
+_DEFAULT_AUTO_UNWRAP = True
+_AUTO_UNWRAP = os.environ.get("AUTO_UNWRAP_TRANSFORMED_ENV")
+
+
+class set_auto_unwrap_transformed_env(_DecoratorContextManager):
+    """A context manager or decorator to control whether TransformedEnv should automatically unwrap nested TransformedEnv instances.
+
+    Args:
+        mode (bool): Whether to automatically unwrap nested :class:`~torchrl.envs.TransformedEnv`
+            instances. If ``False``, :class:`~torchrl.envs.TransformedEnv` will not unwrap nested instances.
+            Defaults to ``True``.
+
+    .. note:: Until v0.9, this will raise a warning if :class:`~torchrl.envs.TransformedEnv` are nested
+        and the value is not set explicitly (`auto_unwrap=True` default behavior).
+        You can set the value of :func:`~torchrl.envs.auto_unwrap_transformed_env`
+        through:
+
+        - The ``AUTO_UNWRAP_TRANSFORMED_ENV`` environment variable;
+        - By setting ``torchrl.set_auto_unwrap_transformed_env(val: bool).set()`` at the
+          beginning of your script;
+        - By using ``torchrl.set_auto_unwrap_transformed_env(val: bool)`` as a context
+          manager or a decorator.
+
+    .. seealso:: :class:`~torchrl.envs.TransformedEnv`
+
+    Examples:
+        >>> with set_auto_unwrap_transformed_env(False):
+        ...     env = TransformedEnv(TransformedEnv(env))
+        ...     assert not isinstance(env.base_env, TransformedEnv)
+        >>> @set_auto_unwrap_transformed_env(False)
+        ... def my_function():
+        ...     env = TransformedEnv(TransformedEnv(env))
+        ...     assert not isinstance(env.base_env, TransformedEnv)
+        ...     return env
+
+    """
+
+    def __init__(self, mode: bool) -> None:
+        super().__init__()
+        self.mode = mode
+
+    def clone(self) -> set_auto_unwrap_transformed_env:
+        # override this method if your children class takes __init__ parameters
+        return type(self)(self.mode)
+
+    def __enter__(self) -> None:
+        self.set()
+
+    def set(self) -> None:
+        global _AUTO_UNWRAP
+        self._old_mode = _AUTO_UNWRAP
+        _AUTO_UNWRAP = bool(self.mode)
+        # we do this such that sub-processes see the same lazy op than the main one
+        os.environ["AUTO_UNWRAP_TRANSFORMED_ENV"] = str(_AUTO_UNWRAP)
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        global _AUTO_UNWRAP
+        _AUTO_UNWRAP = self._old_mode
+        os.environ["AUTO_UNWRAP_TRANSFORMED_ENV"] = str(_AUTO_UNWRAP)
+
+
+def auto_unwrap_transformed_env(allow_none=False):
+    """Get the current setting for automatically unwrapping TransformedEnv instances.
+
+    Args:
+        allow_none (bool, optional): If True, returns ``None`` if no setting has been
+            specified. Otherwise, returns the default setting. Defaults to ``False``.
+
+    seealso: :func:`~torchrl.set_auto_unwrap_transformed_env`
+
+    Returns:
+        bool or None: The current setting for automatically unwrapping TransformedEnv
+            instances.
+    """
+    global _AUTO_UNWRAP  # noqa: F824
+    if _AUTO_UNWRAP is None and allow_none:
+        return None
+    elif _AUTO_UNWRAP is None:
+        return _DEFAULT_AUTO_UNWRAP
+    return strtobool(_AUTO_UNWRAP) if isinstance(_AUTO_UNWRAP, str) else _AUTO_UNWRAP
+
+
+def safe_is_current_stream_capturing():
+    """A safe proxy to torch.cuda.is_current_stream_capturing."""
+    if not torch.cuda.is_available():
+        return False
+    try:
+        return torch.cuda.is_current_stream_capturing()
+    except Exception as error:
+        warnings.warn(
+            f"torch.cuda.is_current_stream_capturing() exited unexpectedly with the error message {error=}. "
+            f"Returning False by default."
+        )
+        return False
+
+
+@classmethod
+def as_remote(cls, remote_config: dict[str, Any] | None = None):
+    """Creates an instance of a remote ray class.
+
+    Args:
+        cls (Python Class): class to be remotely instantiated.
+        remote_config (dict): the quantity of CPU cores to reserve for this class.
+
+    Returns:
+        A function that creates ray remote class instances.
+    """
+    import ray
+
+    if remote_config is None:
+        remote_config = {}
+
+    remote_collector = ray.remote(**remote_config)(cls)
+    remote_collector.is_remote = True
+    return remote_collector
+
+
+def get_ray_default_runtime_env() -> dict[str, Any]:
+    """Get the default Ray runtime environment configuration for TorchRL.
+
+    This function returns a runtime environment configuration that excludes
+    large directories and files that should not be uploaded to Ray workers.
+    This helps prevent issues with Ray's working_dir size limits (512MB default).
+
+    Returns:
+        dict: A dictionary containing the default runtime_env configuration with
+            excludes for common large directories.
+
+    Examples:
+        >>> import ray
+        >>> from torchrl._utils import get_ray_default_runtime_env
+        >>> ray_init_config = {"num_cpus": 4}
+        >>> ray_init_config["runtime_env"] = get_ray_default_runtime_env()
+        >>> ray.init(**ray_init_config)
+
+    Note:
+        The excludes list includes:
+        - Virtual environments (.venv/, venv/, etc.)
+        - Test files and caches
+        - Documentation builds
+        - Benchmarks
+        - Examples and tutorials
+        - CI/CD configurations
+        - IDE configurations
+
+    """
+    return {
+        "excludes": [
+            ".venv/",
+            "venv/",
+            "env/",
+            "ENV/",
+            "env.bak/",
+            "venv.bak/",
+            "test/",
+            "tests/",
+            "docs/",
+            "benchmarks/",
+            "tutorials/",
+            "examples/",
+            ".github/",
+            ".pytest_cache/",
+            ".mypy_cache/",
+            ".ruff_cache/",
+            "__pycache__/",
+            "*.pyc",
+            "*.pyo",
+            "*.egg-info/",
+            ".idea/",
+            ".vscode/",
+            "dev/",
+            "main/",
+            "*.html",
+        ]
+    }
+
+
+def merge_ray_runtime_env(ray_init_config: dict[str, Any]) -> dict[str, Any]:
+    """Merge user-provided ray_init_config with default runtime_env excludes.
+
+    This function ensures that the default TorchRL runtime_env excludes are applied
+    to prevent large directories from being uploaded to Ray workers, while preserving
+    any user-provided configuration.
+
+    Args:
+        ray_init_config (dict): The ray init configuration dictionary to merge.
+
+    Returns:
+        dict: The merged configuration with default runtime_env excludes applied.
+
+    Examples:
+        >>> from torchrl._utils import merge_ray_runtime_env
+        >>> ray_init_config = {"num_cpus": 4}
+        >>> ray_init_config = merge_ray_runtime_env(ray_init_config)
+        >>> ray.init(**ray_init_config)
+
+    """
+    default_runtime_env = get_ray_default_runtime_env()
+    runtime_env = ray_init_config.get("runtime_env")
+
+    # Handle None or missing runtime_env
+    if runtime_env is None:
+        runtime_env = {}
+        ray_init_config["runtime_env"] = runtime_env
+    elif not isinstance(runtime_env, dict):
+        runtime_env = dict(runtime_env)
+        ray_init_config["runtime_env"] = runtime_env
+
+    # Merge excludes lists
+    excludes = runtime_env.get("excludes", [])
+    runtime_env["excludes"] = list(set(default_runtime_env["excludes"] + excludes))
+
+    # Ensure env_vars exists
+    if "env_vars" not in runtime_env:
+        runtime_env["env_vars"] = {}
+    elif not isinstance(runtime_env["env_vars"], dict):
+        runtime_env["env_vars"] = dict(runtime_env["env_vars"])
+
+    return ray_init_config
+
+
+def rl_warnings():
+    """Checks the status of the RL_WARNINGS env varioble."""
+    return RL_WARNINGS