torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/collectors/utils.py

@@ -0,0 +1,433 @@
+# 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 contextlib
+from collections.abc import Callable, Sequence
+
+import torch
+from pyvers import implement_for
+
+from tensordict import NestedKey, pad, set_lazy_legacy, TensorDict, TensorDictBase
+from tensordict.utils import Buffer
+from torch import multiprocessing as mp, nn as nn
+from torch.nn import Parameter
+
+_NON_NN_POLICY_WEIGHTS = (
+    "The policy is not an nn.Module. TorchRL will assume that the parameter set is empty and "
+    "update_policy_weights_ will be a no-op. Consider passing a local/weight_updater object "
+    "to your collector to handle the weight updates."
+)
+
+
+def _stack_output(fun) -> Callable:
+    def stacked_output_fun(*args, **kwargs):
+        out = fun(*args, **kwargs)
+        return tuple(torch.stack(_o, 0) for _o in out)
+
+    return stacked_output_fun
+
+
+def _stack_output_zip(fun) -> Callable:
+    def stacked_output_fun(*args, **kwargs):
+        out = fun(*args, **kwargs)
+        return tuple(torch.stack(_o, 0) for _o in zip(*out))
+
+    return stacked_output_fun
+
+
+@set_lazy_legacy(False)
+def split_trajectories(
+    rollout_tensordict: TensorDictBase,
+    *,
+    prefix=None,
+    trajectory_key: NestedKey | None = None,
+    done_key: NestedKey | None = None,
+    as_nested: bool = False,
+) -> TensorDictBase:
+    """A util function for trajectory separation.
+
+    Takes a tensordict with a key traj_ids that indicates the id of each trajectory.
+
+    From there, builds a B x T x ... zero-padded tensordict with B batches on max duration T
+
+    Args:
+        rollout_tensordict (TensorDictBase): a rollout with adjacent trajectories
+            along the last dimension.
+
+    Keyword Args:
+        prefix (NestedKey, optional): the prefix used to read and write meta-data,
+            such as ``"traj_ids"`` (the optional integer id of each trajectory)
+            and the ``"mask"`` entry indicating which data are valid and which
+            aren't. Defaults to ``"collector"`` if the input has a ``"collector"``
+            entry, ``()`` (no prefix) otherwise.
+            ``prefix`` is kept as a legacy feature and will be deprecated eventually.
+            Prefer ``trajectory_key`` or ``done_key`` whenever possible.
+        trajectory_key (NestedKey, optional): the key pointing to the trajectory
+            ids. Supersedes ``done_key`` and ``prefix``. If not provided, defaults
+            to ``(prefix, "traj_ids")``.
+        done_key (NestedKey, optional): the key pointing to the ``"done""`` signal,
+            if the trajectory could not be directly recovered. Defaults to ``"done"``.
+        as_nested (bool or torch.layout, optional): whether to return the results as nested
+            tensors. Defaults to ``False``. If a ``torch.layout`` is provided, it will be used
+            to construct the nested tensor, otherwise the default layout will be used.
+
+            .. note:: Using ``split_trajectories(tensordict, as_nested=True).to_padded_tensor(mask=mask_key)``
+                should result in the exact same result as ``as_nested=False``. Since this is an experimental
+                feature and relies on nested_tensors, which API may change in the future, we made this
+                an optional feature. The runtime should be faster with ``as_nested=True``.
+
+            .. note:: Providing a layout lets the user control whether the nested tensor is to be used
+                with ``torch.strided`` or ``torch.jagged`` layout. While the former has slightly more
+                capabilities at the time of writing, the second will be the main focus of the PyTorch team
+                in the future due to its better compatibility with :func:`~torch.compile`.
+
+    Returns:
+        A new tensordict with a leading dimension corresponding to the trajectory.
+        A ``"mask"`` boolean entry sharing the ``trajectory_key`` prefix
+        and the tensordict shape is also added. It indicated the valid elements of the tensordict,
+        as well as a ``"traj_ids"`` entry if ``trajectory_key`` could not be found.
+
+    Examples:
+        >>> from tensordict import TensorDict
+        >>> import torch
+        >>> from torchrl.collectors.utils import split_trajectories
+        >>> obs = torch.cat([torch.arange(10), torch.arange(5)])
+        >>> obs_ = torch.cat([torch.arange(1, 11), torch.arange(1, 6)])
+        >>> done = torch.zeros(15, dtype=torch.bool)
+        >>> done[9] = True
+        >>> trajectory_id = torch.cat([torch.zeros(10, dtype=torch.int32),
+        ...     torch.ones(5, dtype=torch.int32)])
+        >>> data = TensorDict({"obs": obs, ("next", "obs"): obs_, ("next", "done"): done, "trajectory": trajectory_id}, batch_size=[15])
+        >>> data_split = split_trajectories(data, done_key="done")
+        >>> print(data_split)
+        TensorDict(
+            fields={
+                mask: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
+                        obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False)},
+                    batch_size=torch.Size([2, 10]),
+                    device=None,
+                    is_shared=False),
+                obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False),
+                traj_ids: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False),
+                trajectory: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int32, is_shared=False)},
+            batch_size=torch.Size([2, 10]),
+            device=None,
+            is_shared=False)
+        >>> # check that split_trajectories got the trajectories right with the done signal
+        >>> assert (data_split["traj_ids"] == data_split["trajectory"]).all()
+        >>> print(data_split["mask"])
+        tensor([[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],
+                [ True,  True,  True,  True,  True, False, False, False, False, False]])
+        >>> data_split = split_trajectories(data, trajectory_key="trajectory")
+        >>> print(data_split)
+        TensorDict(
+            fields={
+                mask: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.bool, is_shared=False),
+                        obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False)},
+                    batch_size=torch.Size([2, 10]),
+                    device=None,
+                    is_shared=False),
+                obs: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False),
+                trajectory: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int32, is_shared=False)},
+            batch_size=torch.Size([2, 10]),
+            device=None,
+            is_shared=False)
+
+    """
+    mask_key = None
+    if trajectory_key is not None:
+        from torchrl.envs.utils import _replace_last
+
+        traj_ids_key = trajectory_key
+        mask_key = _replace_last(trajectory_key, "mask")
+    else:
+        if prefix is None and "collector" in rollout_tensordict.keys():
+            prefix = "collector"
+        if prefix is None:
+            traj_ids_key = "traj_ids"
+            mask_key = "mask"
+        else:
+            traj_ids_key = (prefix, "traj_ids")
+            mask_key = (prefix, "mask")
+
+    rollout_tensordict = rollout_tensordict.copy()
+    traj_ids = rollout_tensordict.get(traj_ids_key, None)
+    if traj_ids is None:
+        if done_key is None:
+            done_key = "done"
+        done_key = ("next", done_key)
+        done = rollout_tensordict.get(done_key)
+        idx = (slice(None),) * (rollout_tensordict.ndim - 1) + (slice(None, -1),)
+        done_sel = done[idx]
+        pads = [1, 0]
+        pads = [0, 0] * (done.ndim - rollout_tensordict.ndim) + pads
+        done_sel = torch.nn.functional.pad(done_sel, pads)
+        if done_sel.shape != done.shape:
+            raise RuntimeError(
+                f"done and done_sel have different shape {done.shape} - {done_sel.shape} "
+            )
+        traj_ids = done_sel.cumsum(rollout_tensordict.ndim - 1)
+        traj_ids = traj_ids.squeeze(-1)
+        if rollout_tensordict.ndim > 1:
+            for i in range(1, rollout_tensordict.shape[0]):
+                traj_ids[i] += traj_ids[i - 1].max() + 1
+        rollout_tensordict.set(traj_ids_key, traj_ids)
+
+    splits = traj_ids.reshape(-1)
+    splits = [(splits == i).sum().item() for i in splits.unique_consecutive()]
+    # if all splits are identical then we can skip this function
+    if len(set(splits)) == 1 and splits[0] == traj_ids.shape[-1]:
+        rollout_tensordict.set(
+            mask_key,
+            torch.ones(
+                rollout_tensordict.shape,
+                device=rollout_tensordict.device,
+                dtype=torch.bool,
+            ),
+        )
+        if rollout_tensordict.ndimension() == 1:
+            rollout_tensordict = rollout_tensordict.unsqueeze(0)
+        return rollout_tensordict
+
+    out_splits = rollout_tensordict.reshape(-1)
+
+    if as_nested:
+        if hasattr(torch, "_nested_compute_contiguous_strides_offsets"):
+
+            def nest(x, splits=splits):
+                # Convert splits into shapes
+                shape = torch.tensor([[int(split), *x.shape[1:]] for split in splits])
+                return torch._nested_view_from_buffer(
+                    x.reshape(-1),
+                    shape,
+                    *torch._nested_compute_contiguous_strides_offsets(shape),
+                )
+
+            return out_splits._fast_apply(
+                nest,
+                batch_size=[len(splits), -1],
+            )
+        else:
+            out_splits = out_splits.split(splits, 0)
+
+            layout = as_nested if as_nested is not bool else None
+
+            if torch.__version__ < "2.4":
+                # Layout must be True, there is no other layout available
+                if layout not in (True,):
+                    raise RuntimeError(
+                        f"layout={layout} is only available for torch>=v2.4"
+                    )
+
+                def nest(*x):
+                    return torch.nested.nested_tensor(list(x))
+
+            else:
+
+                def nest(*x):
+                    return torch.nested.nested_tensor(list(x), layout=layout)
+
+            return out_splits[0]._fast_apply(
+                nest,
+                *out_splits[1:],
+                batch_size=[len(out_splits), *out_splits[0].batch_size[:-1], -1],
+            )
+
+    out_splits = out_splits.split(splits, 0)
+
+    for out_split in out_splits:
+        out_split.set(
+            mask_key,
+            torch.ones(
+                out_split.shape,
+                dtype=torch.bool,
+                device=out_split.device,
+            ),
+        )
+    if len(out_splits) > 1:
+        MAX = max(*[out_split.shape[0] for out_split in out_splits])
+    else:
+        MAX = out_splits[0].shape[0]
+    td = torch.stack(
+        [pad(out_split, [0, MAX - out_split.shape[0]]) for out_split in out_splits], 0
+    )
+    return td
+
+
+@implement_for("torch", "2.5.0")
+def _cast(
+    p: nn.Parameter | torch.Tensor,
+    param_maybe_buffer: nn.Parameter | torch.Tensor | None = None,
+) -> nn.Parameter | torch.Tensor:
+    if param_maybe_buffer is None:
+        param_maybe_buffer = p
+        p = p.data
+    if isinstance(param_maybe_buffer, Parameter):
+        # Create parameter without gradients to avoid serialization issues
+        return Parameter(p, requires_grad=False)
+    if isinstance(param_maybe_buffer, Buffer):
+        return Buffer(p)
+    if p.requires_grad:
+        raise RuntimeError(f"Cannot cast tensor {p} with gradients")
+    return p
+
+
+def _make_meta_policy(policy: nn.Module):
+    """Create context manager that temporarily puts policy parameters on meta device.
+
+    This is used with weight sync schemes to send policy structure without weights.
+    The actual weights are distributed by the schemes.
+
+    Args:
+        policy: Policy module to temporarily modify.
+
+    Returns:
+        A context manager that temporarily replaces policy parameters with meta device versions.
+        On exit, the original parameters are restored to the policy.
+    """
+    param_and_buf = TensorDict.from_module(policy, as_module=True)
+    return param_and_buf.data.to("meta").apply(_cast, param_and_buf).to_module(policy)
+
+
+@implement_for("torch", None, "2.8")
+def _make_meta_policy_cm(
+    policy: nn.Module, *, mp_start_method: str
+) -> contextlib.AbstractContextManager:
+    """Return the context manager used to make a policy 'stateless' for worker pickling.
+
+    On older PyTorch versions (<2.8), pickling meta-device storages when using the
+    ``spawn`` start method may fail (e.g., triggering ``_share_filename_: only available on CPU``).
+    In that case, we avoid converting parameters/buffers to meta and simply return a no-op
+    context manager.
+    """
+    if mp_start_method == "spawn":
+        return contextlib.nullcontext()
+    return _make_meta_policy(policy)
+
+
+@implement_for("torch", "2.8")
+def _make_meta_policy_cm(  # noqa: F811
+    policy: nn.Module, *, mp_start_method: str
+) -> contextlib.AbstractContextManager:
+    """Return the context manager used to make a policy 'stateless' for worker pickling.
+
+    On PyTorch >= 2.8, meta-device policy structures can be pickled reliably under ``spawn``.
+    """
+    return _make_meta_policy(policy)
+
+
+@implement_for("torch", None, "2.5.0")
+def _cast(  # noqa
+    p: nn.Parameter | torch.Tensor,
+    param_maybe_buffer: nn.Parameter | torch.Tensor | None = None,
+) -> nn.Parameter | torch.Tensor:
+    if param_maybe_buffer is None:
+        param_maybe_buffer = p
+        p = p.data
+    if isinstance(param_maybe_buffer, Parameter):
+        # Create parameter without gradients to avoid serialization issues
+        return Parameter(p, requires_grad=False)
+    if p.requires_grad:
+        raise RuntimeError(f"Cannot cast tensor {p} with gradients")
+    return p
+
+
+def _map_to_cpu_if_needed(x):
+    """Map tensors on exotic devices (MPS, NPU, etc.) to CPU.
+
+    CPU and CUDA tensors are kept as-is since they can be shared across processes.
+    Only exotic devices that don't support multiprocessing are mapped to CPU.
+    """
+    if isinstance(x, torch.Tensor):
+        # CPU and CUDA can be shared across processes
+        if x.device.type in ("cpu", "cuda"):
+            return x
+        # Exotic devices (MPS, NPU, etc.) need to be mapped to CPU
+        return x.cpu()
+    return x
+
+
+def _make_meta_params(param):
+    is_param = isinstance(param, Parameter)
+
+    pd = param.detach().to("meta")
+
+    if is_param:
+        pd = Parameter(pd, requires_grad=False)
+    return pd
+
+
+class _TrajectoryPool:
+    def __init__(self, ctx=None, lock: bool = False):
+        self.ctx = ctx
+        self._traj_id = torch.zeros((), device="cpu", dtype=torch.int)
+        # Only use shared memory when multiprocessing context is provided
+        # This avoids issues with shared memory when the mp subsystem is in a bad state
+        if ctx is not None:
+            self._traj_id = self._traj_id.share_memory_()
+        if ctx is None:
+            self.lock = contextlib.nullcontext() if not lock else mp.RLock()
+        else:
+            self.lock = contextlib.nullcontext() if not lock else ctx.RLock()
+
+    def get_traj_and_increment(self, n=1, device=None):
+        with self.lock:
+            v = self._traj_id.item()
+            out = torch.arange(v, v + n).to(device)
+            self._traj_id.copy_(1 + out[-1].item())
+        return out
+
+
+def _map_weight(
+    weight,
+    policy_device,
+):
+
+    is_param = isinstance(weight, Parameter)
+    is_buffer = isinstance(weight, Buffer)
+    weight = weight.data
+    if weight.device != policy_device:
+        weight = weight.to(policy_device)
+    elif weight.device.type in ("cpu",):
+        weight = weight.share_memory_()
+    if is_param:
+        weight = Parameter(weight, requires_grad=False)
+    elif is_buffer:
+        weight = Buffer(weight)
+    return weight
+
+
+def _make_policy_factory(
+    *, policy: Callable, policy_factory, weight_sync_scheme, worker_idx, pipe=None
+):
+    has_policy_factory = policy_factory is not None and (
+        (isinstance(policy_factory, Sequence) and any(policy_factory))
+        or not isinstance(policy_factory, Sequence)
+    )
+    if policy is not None and has_policy_factory:
+        raise ValueError("policy cannot be used with policy_factory")
+    elif has_policy_factory:
+        if isinstance(policy_factory, Sequence):
+            return policy_factory
+        else:
+            policy = policy_factory()
+
+    if weight_sync_scheme is not None:
+        # Initialize the receiver on the worker side
+        weight_sync_scheme.init_on_receiver(
+            model=policy,
+            model_id="policy",
+            worker_idx=worker_idx,
+        )
+        # Synchronize initial weights
+        weight_sync_scheme.connect(worker_idx=worker_idx)
+    return policy