Update (base update)

[ghstack-poisoned]

Author: vmoens

torchrl/__init__.py

@@ -29,10 +29,6 @@
 
 set_lazy_legacy(False).set()
 
-if torch.cuda.device_count() > 1:
-    n = torch.cuda.device_count() - 1
-    os.environ["MUJOCO_EGL_DEVICE_ID"] = str(1 + (os.getpid() % n))
-
 from ._extension import _init_extension  # noqa: E402
 
 __version__ = None  # type: ignore

torchrl/_utils.py

@@ -333,6 +333,47 @@ def reset(cls):
         cls.erase()
 
 
+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 nullcontext()
+    return timeit(name)
+
+
+def _maybe_record_function(name):
+    """Return record_function context if not compiling, nullcontext otherwise.
+
+    torch.autograd.profiler.record_function cannot be used inside compiled regions.
+    """
+    from torch.autograd.profiler import record_function
+
+    if is_compiling():
+        return nullcontext()
+    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.
+    """
+
+    def decorator(fn: Callable) -> Callable:
+        @wraps(fn)
+        def wrapped(*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:

torchrl/collectors/_constants.py

@@ -28,6 +28,7 @@ def cudagraph_mark_step_begin():
     "INSTANTIATE_TIMEOUT",
     "_MIN_TIMEOUT",
     "_MAX_IDLE_COUNT",
+    "WEIGHT_SYNC_TIMEOUT",
     "DEFAULT_EXPLORATION_TYPE",
     "_is_osx",
     "_Interruptor",
@@ -38,6 +39,9 @@ def cudagraph_mark_step_begin():
 _TIMEOUT = 1.0
 INSTANTIATE_TIMEOUT = 20
 _MIN_TIMEOUT = 1e-3  # should be several orders of magnitude inferior wrt time spent collecting a trajectory
+# Timeout for weight synchronization during collector init.
+# Increase this when using many collectors across different CUDA devices.
+WEIGHT_SYNC_TIMEOUT = float(os.environ.get("TORCHRL_WEIGHT_SYNC_TIMEOUT", 120.0))
 # MAX_IDLE_COUNT is the maximum number of times a Dataloader worker can timeout with his queue.
 _MAX_IDLE_COUNT = int(os.environ.get("MAX_IDLE_COUNT", torch.iinfo(torch.int64).max))
 

torchrl/collectors/collectors.py

@@ -18,6 +18,7 @@
     cudagraph_mark_step_begin,
     DEFAULT_EXPLORATION_TYPE,
     INSTANTIATE_TIMEOUT,
+    WEIGHT_SYNC_TIMEOUT,
 )
 
 from torchrl.collectors._multi_async import MultiAsyncCollector, MultiaSyncDataCollector
@@ -50,6 +51,7 @@
     # Constants
     "_TIMEOUT",
     "INSTANTIATE_TIMEOUT",
+    "WEIGHT_SYNC_TIMEOUT",
     "_MIN_TIMEOUT",
     "_MAX_IDLE_COUNT",
     "DEFAULT_EXPLORATION_TYPE",

torchrl/data/replay_buffers/samplers.py

@@ -1556,14 +1556,30 @@ def _get_index(
             out_of_traj = relative_starts < 0
             if out_of_traj.any():
                 # a negative start means sampling fewer elements
+                # Convert seq_length to tensor to avoid torch.compile inductor C++ codegen
+                # bug with mixed scalar/tensor int64 in blendv operations (see PyTorch #xyz)
+                seq_length_t = torch.as_tensor(
+                    seq_length,
+                    dtype=relative_starts.dtype,
+                    device=relative_starts.device,
+                )
                 seq_length = torch.where(
-                    ~out_of_traj, seq_length, seq_length + relative_starts
+                    ~out_of_traj, seq_length_t, seq_length_t + relative_starts
+                )
+                relative_starts = torch.where(
+                    ~out_of_traj, relative_starts, torch.zeros_like(relative_starts)
                 )
-                relative_starts = torch.where(~out_of_traj, relative_starts, 0)
         if self.span[1]:
             out_of_traj = relative_starts + seq_length > lengths[traj_idx]
             if out_of_traj.any():
                 # a negative start means sampling fewer elements
+                # Convert seq_length to tensor if it's still a scalar
+                if not isinstance(seq_length, torch.Tensor):
+                    seq_length = torch.as_tensor(
+                        seq_length,
+                        dtype=relative_starts.dtype,
+                        device=relative_starts.device,
+                    )
                 seq_length = torch.minimum(
                     seq_length, lengths[traj_idx] - relative_starts
                 )

torchrl/envs/model_based/common.py

@@ -117,11 +117,13 @@ def __init__(
         device: DEVICE_TYPING = "cpu",
         batch_size: torch.Size | None = None,
         run_type_checks: bool = False,
+        allow_done_after_reset: bool = False,
     ):
         super().__init__(
             device=device,
             batch_size=batch_size,
             run_type_checks=run_type_checks,
+            allow_done_after_reset=allow_done_after_reset,
         )
         self.world_model = world_model.to(self.device)
         self.world_model_params = params
@@ -161,12 +163,13 @@ def _step(
         else:
             tensordict_out = self.world_model(tensordict_out)
         # done can be missing, it will be filled by `step`
-        tensordict_out = tensordict_out.select(
-            *self.observation_spec.keys(),
-            *self.full_done_spec.keys(),
-            *self.full_reward_spec.keys(),
-            strict=False,
+        # Convert to list for torch.compile compatibility (dynamo can't unpack _CompositeSpecKeysView)
+        keys_to_select = (
+            list(self.observation_spec.keys())
+            + list(self.full_done_spec.keys())
+            + list(self.full_reward_spec.keys())
         )
+        tensordict_out = tensordict_out.select(*keys_to_select, strict=False)
         return tensordict_out
 
     @abc.abstractmethod

torchrl/modules/distributions/continuous.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import weakref
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from numbers import Number
 
 import numpy as np
@@ -76,7 +76,7 @@ class IndependentNormal(D.Independent):
     def __init__(
         self,
         loc: torch.Tensor,
-        scale: torch.Tensor,
+        scale: torch.Tensor | float | Callable[[torch.Tensor], torch.Tensor],
         upscale: float = 5.0,
         tanh_loc: bool = False,
         event_dim: int = 1,
@@ -86,11 +86,25 @@ def __init__(
         self.upscale = upscale
         self._event_dim = event_dim
         self._kwargs = kwargs
+        # Support callable scale (e.g., torch.ones_like) for compile-friendliness
+        if callable(scale) and not isinstance(scale, torch.Tensor):
+            scale = scale(loc)
+        elif not isinstance(scale, torch.Tensor):
+            scale = torch.as_tensor(scale, device=loc.device, dtype=loc.dtype)
+        elif scale.device != loc.device:
+            scale = scale.to(loc.device, non_blocking=loc.device.type == "cuda")
         super().__init__(D.Normal(loc, scale, **kwargs), event_dim)
 
     def update(self, loc, scale):
         if self.tanh_loc:
             loc = self.upscale * (loc / self.upscale).tanh()
+        # Support callable scale (e.g., torch.ones_like) for compile-friendliness
+        if callable(scale) and not isinstance(scale, torch.Tensor):
+            scale = scale(loc)
+        elif not isinstance(scale, torch.Tensor):
+            scale = torch.as_tensor(scale, device=loc.device, dtype=loc.dtype)
+        elif scale.device != loc.device:
+            scale = scale.to(loc.device, non_blocking=loc.device.type == "cuda")
         super().__init__(D.Normal(loc, scale, **self._kwargs), self._event_dim)
 
     @property
@@ -343,7 +357,7 @@ class TanhNormal(FasterTransformedDistribution):
     def __init__(
         self,
         loc: torch.Tensor,
-        scale: torch.Tensor,
+        scale: torch.Tensor | float | Callable[[torch.Tensor], torch.Tensor],
         upscale: torch.Tensor | Number = 5.0,
         low: torch.Tensor | Number = -1.0,
         high: torch.Tensor | Number = 1.0,
@@ -353,8 +367,14 @@ def __init__(
     ):
         if not isinstance(loc, torch.Tensor):
             loc = torch.as_tensor(loc, dtype=torch.get_default_dtype())
-        if not isinstance(scale, torch.Tensor):
-            scale = torch.as_tensor(scale, dtype=torch.get_default_dtype())
+        _non_blocking = loc.device.type == "cuda"
+        # Support callable scale (e.g., torch.ones_like) for compile-friendliness
+        if callable(scale) and not isinstance(scale, torch.Tensor):
+            scale = scale(loc)
+        elif not isinstance(scale, torch.Tensor):
+            scale = torch.as_tensor(scale, device=loc.device, dtype=loc.dtype)
+        elif scale.device != loc.device:
+            scale = scale.to(loc.device, non_blocking=_non_blocking)
         if event_dims is None:
             event_dims = min(1, loc.ndim)
 
@@ -373,11 +393,11 @@ def __init__(
         if not isinstance(high, torch.Tensor):
             high = torch.as_tensor(high, device=loc.device)
         elif high.device != loc.device:
-            high = high.to(loc.device)
+            high = high.to(loc.device, non_blocking=_non_blocking)
         if not isinstance(low, torch.Tensor):
             low = torch.as_tensor(low, device=loc.device)
         elif low.device != loc.device:
-            low = low.to(loc.device)
+            low = low.to(loc.device, non_blocking=_non_blocking)
         if not is_compiling() and not safe_is_current_stream_capturing():
             self.non_trivial_max = (high != 1.0).any()
             self.non_trivial_min = (low != -1.0).any()
@@ -391,10 +411,10 @@ def __init__(
         self.upscale = (
             upscale
             if not isinstance(upscale, torch.Tensor)
-            else upscale.to(self.device)
+            else upscale.to(self.device, non_blocking=_non_blocking)
         )
 
-        low = low.to(loc.device)
+        low = low.to(loc.device, non_blocking=_non_blocking)
         self.low = low
         self.high = high
 
@@ -434,6 +454,13 @@ def update(self, loc: torch.Tensor, scale: torch.Tensor) -> None:
             # loc must be rescaled if tanh_loc
             if is_compiling() or (self.non_trivial_max or self.non_trivial_min):
                 loc = loc + (self.high - self.low) / 2 + self.low
+        # Support callable scale (e.g., torch.ones_like) for compile-friendliness
+        if callable(scale) and not isinstance(scale, torch.Tensor):
+            scale = scale(loc)
+        elif not isinstance(scale, torch.Tensor):
+            scale = torch.as_tensor(scale, device=loc.device, dtype=loc.dtype)
+        elif scale.device != loc.device:
+            scale = scale.to(loc.device, non_blocking=loc.device.type == "cuda")
         self.loc = loc
         self.scale = scale
 

torchrl/modules/distributions/truncated_normal.py

@@ -35,8 +35,9 @@ class TruncatedStandardNormal(Distribution):
 
     def __init__(self, a, b, validate_args=None, device=None):
         self.a, self.b = broadcast_all(a, b)
-        self.a = self.a.to(device)
-        self.b = self.b.to(device)
+        _non_blocking = device is not None and torch.device(device).type == "cuda"
+        self.a = self.a.to(device, non_blocking=_non_blocking)
+        self.b = self.b.to(device, non_blocking=_non_blocking)
         if isinstance(a, Number) and isinstance(b, Number):
             batch_shape = torch.Size()
         else:
@@ -146,8 +147,9 @@ class TruncatedNormal(TruncatedStandardNormal):
     def __init__(self, loc, scale, a, b, validate_args=None, device=None):
         scale = scale.clamp_min(self.eps)
         self.loc, self.scale, a, b = broadcast_all(loc, scale, a, b)
-        a = a.to(device)
-        b = b.to(device)
+        _non_blocking = device is not None and torch.device(device).type == "cuda"
+        a = a.to(device, non_blocking=_non_blocking)
+        b = b.to(device, non_blocking=_non_blocking)
         self._non_std_a = a
         self._non_std_b = b
         a = (a - self.loc) / self.scale

torchrl/modules/distributions/utils.py

@@ -16,21 +16,23 @@
 
 def _cast_device(elt: torch.Tensor | float, device) -> torch.Tensor | float:
     if isinstance(elt, torch.Tensor):
-        return elt.to(device)
+        _non_blocking = device is not None and torch.device(device).type == "cuda"
+        return elt.to(device, non_blocking=_non_blocking)
     return elt
 
 
 def _cast_transform_device(transform, device):
     if transform is None:
         return transform
-    elif isinstance(transform, d.ComposeTransform):
+    _non_blocking = device is not None and torch.device(device).type == "cuda"
+    if isinstance(transform, d.ComposeTransform):
         for i, t in enumerate(transform.parts):
             transform.parts[i] = _cast_transform_device(t, device)
     elif isinstance(transform, d.Transform):
         for attribute in dir(transform):
             value = getattr(transform, attribute)
             if isinstance(value, torch.Tensor):
-                setattr(transform, attribute, value.to(device))
+                setattr(transform, attribute, value.to(device, non_blocking=_non_blocking))
         return transform
     else:
         raise TypeError(

torchrl/weight_update/_shared.py

@@ -11,6 +11,7 @@
 from torch import multiprocessing as mp, nn
 
 from torchrl._utils import logger as torchrl_logger
+from torchrl.collectors._constants import WEIGHT_SYNC_TIMEOUT
 
 from torchrl.weight_update.utils import _resolve_model
 from torchrl.weight_update.weight_sync_schemes import (
@@ -97,7 +98,7 @@ def setup_connection_and_weights_on_receiver(
         weights: Any = None,
         model: Any = None,
         strategy: Any = None,
-        timeout: float = 10.0,
+        timeout: float = WEIGHT_SYNC_TIMEOUT,
     ) -> TensorDictBase:
         """Receive shared memory buffer reference from sender via their per-worker queues.