# 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
import platform
import sys
import warnings
from collections.abc import Callable, Sequence
from typing import Literal
import torch
from pyvers import implement_for
from tensordict import (
NestedKey,
NonTensorData,
NonTensorStack,
pad,
set_lazy_legacy,
TensorDict,
TensorDictBase,
)
from tensordict.base import _is_leaf_nontensor, _NESTED_TENSORS_AS_LISTS
from tensordict.utils import Buffer
from torch import multiprocessing as mp, nn as nn
from torch.nn import Parameter
from torchrl._utils import DEFAULT_DONE_KEYS
_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 _log_prob_key_from_sample_key(key: NestedKey) -> NestedKey:
if isinstance(key, tuple):
return (*key[:-1], f"{key[-1]}_log_prob")
return f"{key}_log_prob"
def _ensure_derived_policy_output_keys(policy: Callable | None) -> None:
"""Restore derived policy-output metadata after policy serialization.
Collectors already discover declared policy-produced keys, including
arbitrary non-action metadata, with the initial policy dry-run. Some
modules derive additional output-key metadata at construction time. If
serialization drops that metadata, the first dry-run cannot see those
outputs, so restore the known derived keys before collection starts.
"""
if policy is None:
return
modules = policy.modules() if isinstance(policy, nn.Module) else (policy,)
for module in modules:
if not getattr(module, "return_log_prob", False):
continue
try:
log_prob_keys = module.log_prob_keys
except AttributeError:
continue
if log_prob_keys:
continue
out_keys = getattr(module, "out_keys", None)
if not out_keys:
continue
module.log_prob_keys = [_log_prob_key_from_sample_key(key) for key in out_keys]
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
[docs]
@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.
.. note:: This function splits whatever the input contains: trajectories
spanning several collector batches stay split across the corresponding
calls. To collect batches made of complete trajectories only, pass
``trajs_per_batch`` to the collector instead (see
:ref:`collectors_replay_trajs`).
.. seealso:: This function operates on contiguous *rollout batches* (fresh
collector output). To recover trajectory boundaries from data laid out
in a replay-buffer storage -- where the ring buffer can wrap and the
write cursor acts as an implicit truncation -- use
:func:`~torchrl.data.find_start_stop_traj` instead. Note that the
padded layout this function produces is discouraged for new code
unless explicitly needed; see :ref:`data-layout-split-trajectories`.
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, preserve_module_state=False)
)
@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 _platform_supports_cuda_ipc() -> bool:
"""Whether torch.multiprocessing can share CUDA tensors with other processes.
CUDA tensors are exchanged between processes through CUDA IPC handles,
which the CUDA driver only implements on native Linux. On Windows and on
WSL2 the exchange fails silently: the receiving process observes zeroed
tensors and the sender's CUDA memory can be corrupted as well. See
https://github.com/pytorch/pytorch/issues/149155 and
https://github.com/pytorch/rl/issues/3985.
"""
if sys.platform != "linux":
return False
return "microsoft" not in platform.uname().release.lower()
def _device_shareable_across_processes(device: torch.device | str | int) -> bool:
"""Whether tensors on ``device`` can be sent to another process without a copy.
CPU and meta tensors can be shared on every platform. CUDA tensors rely on
CUDA IPC, which is only available on native Linux. Other backends
(MPS, NPU, ...) cannot be shared across processes at all.
"""
if not isinstance(device, torch.device):
device = torch.device(device)
if device.type in ("cpu", "meta"):
return True
if device.type == "cuda":
return _platform_supports_cuda_ipc()
return False
def _unshareable_devices(
device: torch.device | str | int | None, weights: TensorDictBase | None
) -> list[torch.device]:
"""Return the devices of ``weights`` that cannot cross process boundaries.
``device`` is the target device the weights are about to be moved to; when
it is ``None`` the weights are shipped on their current devices, so each
leaf tensor's device is checked instead.
"""
if device is not None:
if _device_shareable_across_processes(device):
return []
return [torch.device(device)]
if weights is None:
return []
return sorted(
{
tensor.device
for tensor in weights.values(True, True)
if not _device_shareable_across_processes(tensor.device)
},
key=str,
)
def _stage_unshareable_weights_on_cpu(
weights: TensorDictBase | dict | None,
) -> TensorDictBase | dict | None:
"""Return ``weights`` moved to CPU when any leaf cannot cross process boundaries.
Accepts a ``TensorDictBase`` or a state-dict mapping and returns the input
unchanged when every leaf tensor can be shared with other processes.
A warning is emitted when staging happens.
"""
if isinstance(weights, TensorDictBase):
offending = _unshareable_devices(None, weights)
if offending:
_warn_cpu_staging(offending)
weights = weights.to("cpu")
return weights
if isinstance(weights, dict):
offending = sorted(
{
value.device
for value in weights.values()
if isinstance(value, torch.Tensor)
and not _device_shareable_across_processes(value.device)
},
key=str,
)
if offending:
_warn_cpu_staging(offending)
weights = {
key: value.cpu() if isinstance(value, torch.Tensor) else value
for key, value in weights.items()
}
return weights
return weights
def _warn_cpu_staging(devices: Sequence[torch.device]) -> None:
"""Warn that weights bound for ``devices`` are staged through CPU shared memory."""
device_names = sorted({str(d) for d in devices})
warnings.warn(
f"Policy weights on device(s) {device_names} cannot be shared across processes "
"on this platform: sharing CUDA tensors requires CUDA IPC, which is only "
"available on native Linux (not Windows or WSL2), and backends such as MPS "
"cannot be shared at all. Sending such tensors to another process silently "
"corrupts them (they are received as zeros and the sender's memory can be "
"zeroed too, see https://github.com/pytorch/rl/issues/3985). The weights will "
"be staged through CPU shared memory instead and moved back to the policy "
"device inside each worker, adding a host-device copy to every weight sync.",
UserWarning,
)
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 state_dict(self) -> dict[str, torch.Tensor]:
"""Return the next trajectory identifier without exposing shared storage."""
with self.lock:
return {"traj_id": self._traj_id.clone()}
def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
"""Restore the next trajectory identifier in-place."""
with self.lock:
self._traj_id.copy_(state_dict["traj_id"])
class _CollectorProgress:
"""Lock-free collector progress shared by one writer per worker row."""
_KEYS = (
"stepped_frames",
"trajectory_completed_frames",
"trajectory_pending_frames",
"replay_written_frames",
"completed_trajectories",
)
def __init__(self, num_workers: int = 1, ctx=None):
self.num_workers = num_workers
size = num_workers * len(self._KEYS)
self._values = [0] * size if ctx is None else ctx.RawArray("q", size)
def increment_stepped(
self, worker_idx: int, frames: int, *, trajectory_pending: bool
) -> None:
offset = worker_idx * len(self._KEYS)
self._values[offset] += frames
if trajectory_pending:
self._values[offset + 2] += frames
def record_trajectory_completion(
self, worker_idx: int, frames: int, trajectories: int
) -> None:
offset = worker_idx * len(self._KEYS)
self._values[offset + 1] += frames
self._values[offset + 2] -= frames
self._values[offset + 4] += trajectories
def record_trajectory_pending(self, worker_idx: int, frames: int) -> None:
offset = worker_idx * len(self._KEYS)
self._values[offset + 2] += frames
def record_replay_write(self, worker_idx: int, frames: int) -> None:
offset = worker_idx * len(self._KEYS)
self._values[offset + 3] += frames
def clear_pending(self, worker_idx: int | None = None) -> None:
if worker_idx is None:
for row in range(self.num_workers):
self._values[row * len(self._KEYS) + 2] = 0
else:
self._values[worker_idx * len(self._KEYS) + 2] = 0
def snapshot(self, worker_idx: int | None = None) -> dict[str, int]:
if worker_idx is None:
return {
key: sum(
int(self._values[row * len(self._KEYS) + column])
for row in range(self.num_workers)
)
for column, key in enumerate(self._KEYS)
}
offset = worker_idx * len(self._KEYS)
return {
key: int(self._values[offset + column])
for column, key in enumerate(self._KEYS)
}
def load_snapshot(self, worker_idx: int, state: dict[str, int]) -> None:
offset = worker_idx * len(self._KEYS)
for column, key in enumerate(self._KEYS):
# In-flight trajectory data is intentionally not checkpointed.
value = 0 if key == "trajectory_pending_frames" else state.get(key, 0)
self._values[offset + column] = int(value)
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 _traj_chunk_ends_done(chunk: TensorDictBase) -> bool:
"""Return ``True`` if the last step of *chunk* carries a done/terminated signal."""
for leaf in DEFAULT_DONE_KEYS:
signal = chunk.get(("next", leaf), None)
if signal is not None and signal[-1].any().item():
return True
return False
def _maybe_normalize_replay_buffer_tensordict_device(
data: TensorDictBase,
replay_buffer,
) -> TensorDictBase:
"""Align TensorDict root device metadata with replay-buffer storage when safe."""
if not isinstance(data, TensorDictBase):
return data
storage = getattr(replay_buffer, "_storage", None)
if storage is None:
storage = getattr(replay_buffer, "storage", None)
storage_device = getattr(storage, "device", None)
if storage_device is None:
storage_data = getattr(storage, "_storage", None)
storage_device = getattr(storage_data, "device", None)
if storage_device is None:
member_storages = getattr(storage, "_storages", None)
if member_storages:
member_devices = []
for member in member_storages:
member_device = getattr(member, "device", None) or getattr(
getattr(member, "_storage", None), "device", None
)
if member_device is None or member_device == "auto":
member_devices = []
break
member_devices.append(torch.device(member_device))
if member_devices and len(set(member_devices)) == 1:
storage_device = member_devices[0]
if storage_device is None or storage_device == "auto":
return data
storage_device = torch.device(storage_device)
for value in data.values(
include_nested=True,
leaves_only=True,
is_leaf=_NESTED_TENSORS_AS_LISTS,
):
value_device = getattr(value, "device", None)
if value_device is not None and torch.device(value_device) != storage_device:
return data
data = data.copy()
data.clear_device_()
return data.to(storage_device)
def _traj_ingest(
batch: TensorDictBase,
partial_trajs: dict,
complete_trajs: list,
) -> tuple[int, int]:
"""Route steps from *batch* into per-trajectory buffers.
Completed trajectories are moved from *partial_trajs* into *complete_trajs*.
Returns their total frame and trajectory counts.
"""
flat = batch.reshape(-1)
traj_ids = flat.get(("collector", "traj_ids"), None)
if traj_ids is None:
raise KeyError(
"trajs_per_batch requires ('collector', 'traj_ids') in every "
"collector batch. Make sure the collector is initialized with "
"split_trajs=False (the default)."
)
order = torch.argsort(traj_ids.reshape(-1), stable=True)
flat = flat[order]
traj_ids = traj_ids.reshape(-1)[order]
unique_ids, counts = traj_ids.unique_consecutive(return_counts=True)
start = 0
completed_frames = 0
completed_trajectories = 0
for tid_tensor, count in zip(unique_ids, counts):
tid = tid_tensor.item()
stop = start + count.item()
chunk = flat[start:stop]
start = stop
if tid in partial_trajs:
partial_trajs[tid].append(chunk)
else:
partial_trajs[tid] = [chunk]
if _traj_chunk_ends_done(chunk):
chunks = partial_trajs.pop(tid)
complete = torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0]
complete_trajs.append(complete)
completed_frames += complete.numel()
completed_trajectories += 1
return completed_frames, completed_trajectories
def _traj_emit(
complete_trajs: list,
num_trajectories: int,
traj_format: Literal["padded", "cat"] = "padded",
) -> TensorDictBase:
"""Dequeue *num_trajectories* complete trajectories as a single batch.
With ``traj_format="padded"`` (default), trajectories are zero-padded
along time and stacked into a ``(num_trajectories, max_traj_len)`` batch
with a ``("collector", "mask")`` entry marking the valid steps.
With ``traj_format="cat"``, trajectories are concatenated along time into
a flat ``[sum_i T_i]`` batch (no padding, no mask): trajectories are
contiguous, in completion order, with ``("next", "done")`` ``True`` at the
last step of each and ``("collector", "traj_ids")`` constant within each.
"""
trajs = complete_trajs[:num_trajectories]
del complete_trajs[:num_trajectories]
if traj_format == "cat":
return torch.cat(trajs, 0) if len(trajs) > 1 else trajs[0]
max_len = max(t.shape[0] for t in trajs)
padded = []
for traj in trajs:
traj = traj.copy()
traj.set(
("collector", "mask"),
torch.ones(traj.shape[0], dtype=torch.bool, device=traj.device),
)
pad_len = max_len - traj.shape[0]
if not pad_len:
padded.append(traj)
continue
# tensordict.pad drops NonTensor entries (e.g. language
# instructions), valid steps included: pad those separately by
# repeating the last element (the mask marks validity anyway)
non_tensor_keys = [
key
for key in traj.keys(True, True, is_leaf=_is_leaf_nontensor)
if isinstance(traj.get(key), (NonTensorData, NonTensorStack))
]
if non_tensor_keys:
non_tensor = {key: traj.get(key) for key in non_tensor_keys}
padded_traj = pad(traj.exclude(*non_tensor_keys), [0, pad_len])
for key, value in non_tensor.items():
filler = torch.cat([value[-1:]] * pad_len, 0)
padded_traj.set(key, torch.cat([value, filler], 0))
else:
padded_traj = pad(traj, [0, pad_len])
padded.append(padded_traj)
return torch.stack(padded, 0)
def _validate_traj_format(
traj_format: Literal["padded", "cat"] | None,
trajs_per_batch: int | None,
*,
has_replay_buffer: bool = False,
) -> Literal["padded", "cat"]:
"""Validate and resolve the ``traj_format`` / ``trajs_per_batch`` keyword pair.
``None`` resolves to the current default (``"padded"``) and emits a
:class:`FutureWarning` announcing the upcoming default change whenever the
choice matters, i.e. when ``trajs_per_batch`` batches are yielded rather
than written to a replay buffer (replay-buffer writes are always flat).
"""
if traj_format is None:
if trajs_per_batch is not None and not has_replay_buffer:
warnings.warn(
"trajs_per_batch is set but traj_format is not. The current "
"default trajectory layout is 'padded', but it will change "
"to 'cat' in torchrl v0.16. Pass traj_format='padded' to "
"keep the current behavior, or traj_format='cat' "
"(recommended) to opt in to the new flat, unpadded layout.",
FutureWarning,
stacklevel=3,
)
return "padded"
if traj_format not in ("padded", "cat"):
raise ValueError(f"traj_format must be 'padded' or 'cat', got {traj_format!r}.")
if traj_format != "padded" and trajs_per_batch is None:
raise ValueError(
"traj_format has no effect unless trajs_per_batch is set: the "
"collector only assembles whole trajectories when asked to yield "
"them. Set trajs_per_batch to the number of complete "
"trajectories per batch."
)
return traj_format
def _validate_replay_write_mode(
replay_write_mode: Literal["rollout", "trajectory"] | None,
*,
has_replay_buffer: bool,
trajs_per_batch: int | None,
trajs_per_write: int | None,
) -> Literal["rollout", "trajectory"] | None:
"""Validate and resolve replay-buffer write semantics.
``None`` retains the historical coupling where ``trajs_per_batch`` selects
complete-trajectory writes when a replay buffer is present.
"""
if replay_write_mode not in (None, "rollout", "trajectory"):
raise ValueError(
"replay_write_mode must be 'rollout', 'trajectory', or None, got "
f"{replay_write_mode!r}."
)
if replay_write_mode is not None and not has_replay_buffer:
raise ValueError("replay_write_mode requires a replay_buffer.")
if replay_write_mode is not None and trajs_per_batch is not None:
raise ValueError(
"replay_write_mode cannot be combined with trajs_per_batch: "
"trajs_per_batch controls yielded trajectory batches, while "
"replay_write_mode controls replay-buffer writes. Omit "
"replay_write_mode to preserve the legacy replay_buffer + "
"trajs_per_batch behavior."
)
if replay_write_mode == "rollout" and trajs_per_write is not None:
raise ValueError(
"trajs_per_write is only supported with replay_write_mode='trajectory'."
)
if replay_write_mode == "trajectory" and trajs_per_write is not None:
if (
isinstance(trajs_per_write, bool)
or not isinstance(trajs_per_write, int)
or trajs_per_write < 1
):
raise ValueError("trajs_per_write must be a positive integer.")
if not has_replay_buffer:
return None
if replay_write_mode is not None:
return replay_write_mode
if trajs_per_batch is not None:
return "trajectory"
return "rollout"
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):
# Use worker_idx to get the correct factory for this worker
policy = policy_factory[worker_idx]()
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)
_ensure_derived_policy_output_keys(policy)
return policy
