from __future__ import annotations
import abc
import contextlib
import functools
import typing
import warnings
from collections import OrderedDict
from collections.abc import Callable, Iterator
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal, overload
import torch
from tensordict import TensorDict, TensorDictBase
from tensordict.base import NO_DEFAULT
from tensordict.nn import TensorDictModule, TensorDictModuleBase
from torch import nn as nn
from torch.utils.data import IterableDataset
from torchrl.collectors.utils import (
_CollectorProgress,
_map_weight,
_maybe_normalize_replay_buffer_tensordict_device,
_traj_emit,
_traj_ingest,
)
from torchrl.collectors.weight_update import WeightUpdaterBase
from torchrl.weight_update.utils import _resolve_attr, _weight_tensor_signature
from torchrl.weight_update.weight_sync_schemes import WeightSyncScheme
@dataclass
class ProfileConfig:
"""Configuration for profiling collector workers.
This class holds all the settings for profiling collector rollouts
using PyTorch's profiler. It's designed to work across all collector types.
Attributes:
workers: List of worker indices to profile. For single-process collectors
(like Collector), this is ignored. For multi-process collectors
(like MultiSyncCollector, MultiAsyncCollector), only the specified
workers will be profiled. Defaults to [0].
num_rollouts: Total number of rollouts to profile (including warmup).
After this many rollouts, profiling stops. Defaults to 3.
warmup_rollouts: Number of rollouts to skip before starting actual
profiling. This allows JIT/compile warmup. Defaults to 1.
save_path: Path to save the profiling trace. If None, traces are saved
to "./collector_profile_{worker_idx}.json". Supports {worker_idx}
placeholder for worker-specific files.
activities: List of profiler activities. Defaults to CPU and CUDA.
record_shapes: Whether to record tensor shapes. Defaults to True.
profile_memory: Whether to profile memory usage. Defaults to False.
with_stack: Whether to record stack traces. Defaults to True.
with_flops: Whether to compute FLOPS. Defaults to False.
on_trace_ready: Optional callback when trace is ready. If None,
traces are exported to Chrome trace format at save_path.
Example:
>>> from torchrl.collectors import MultiSyncCollector, ProfileConfig
>>> collector = MultiSyncCollector(...)
>>> collector.enable_profile(
... workers=[0],
... num_rollouts=5,
... warmup_rollouts=2,
... save_path="./traces/worker_{worker_idx}.json",
... )
>>> for data in collector:
... # First worker will be profiled for rollouts 2-4
... process(data)
"""
workers: list[int] = field(default_factory=lambda: [0])
num_rollouts: int = 3
warmup_rollouts: int = 1
save_path: str | Path | None = None
activities: list[str] = field(default_factory=lambda: ["cpu", "cuda"])
record_shapes: bool = True
profile_memory: bool = False
with_stack: bool = True
with_flops: bool = False
on_trace_ready: Callable | None = None
def __post_init__(self):
"""Validate configuration after initialization."""
if self.num_rollouts <= self.warmup_rollouts:
raise ValueError(
f"num_rollouts ({self.num_rollouts}) must be greater than "
f"warmup_rollouts ({self.warmup_rollouts})"
)
if self.warmup_rollouts < 0:
raise ValueError(
f"warmup_rollouts must be >= 0, got {self.warmup_rollouts}"
)
def get_save_path(self, worker_idx: int) -> Path:
"""Get the save path for a specific worker.
Args:
worker_idx: The worker index.
Returns:
Path object for the trace file.
"""
if self.save_path is None:
return Path(f"./collector_profile_{worker_idx}.json")
path_str = str(self.save_path).format(worker_idx=worker_idx)
return Path(path_str)
def get_activities(self) -> list:
"""Get PyTorch profiler activity list.
Returns:
List of torch.profiler.ProfilerActivity values.
"""
import torch.profiler
activity_map = {
"cpu": torch.profiler.ProfilerActivity.CPU,
"cuda": torch.profiler.ProfilerActivity.CUDA,
}
result = []
for activity in self.activities:
activity_lower = activity.lower()
if activity_lower in activity_map:
# Only add CUDA if CUDA is available
if activity_lower == "cuda" and not torch.cuda.is_available():
continue
result.append(activity_map[activity_lower])
return result
def should_profile_worker(self, worker_idx: int) -> bool:
"""Check if a specific worker should be profiled.
Args:
worker_idx: The worker index to check.
Returns:
True if this worker should be profiled.
"""
return worker_idx in self.workers
class _ProfilerHook:
"""A ``post_collect_hook`` callable that drives a ``torch.profiler.profile``.
The hook owns a profiler in the process where it lives — the main process
for a single :class:`Collector`, each worker process for a multi-collector,
each remote actor for a Ray collector. It starts the profiler lazily on the
first call, steps once per rollout, and auto-stops + exports after
``config.num_rollouts`` rollouts.
Pickling — the hook is a plain Python object and travels through
:class:`~torchrl.data.utils.CloudpickleWrapper` when pushed to mp workers.
The :class:`torch.profiler.profile` instance is built lazily inside the
target process so it never needs to cross a pickle boundary itself.
"""
def __init__(self, config: ProfileConfig, worker_idx: int = 0):
self.config = config
self.worker_idx = worker_idx
self._profiler = None
self._rollout_count = 0
self._stopped = False
def _build_profiler(self) -> Any | None:
active = self.config.num_rollouts - self.config.warmup_rollouts
schedule = torch.profiler.schedule(
skip_first=self.config.warmup_rollouts,
wait=0,
warmup=0,
active=active,
repeat=1,
)
activities = self.config.get_activities()
if not activities:
return None
if self.config.on_trace_ready is not None:
on_trace_ready = self.config.on_trace_ready
else:
save_path = self.config.get_save_path(self.worker_idx)
save_path.parent.mkdir(parents=True, exist_ok=True)
from torchrl import logger as torchrl_logger
def on_trace_ready(prof, _path=save_path, _idx=self.worker_idx):
prof.export_chrome_trace(str(_path))
torchrl_logger.info(f"Profiler [worker {_idx}]: trace saved to {_path}")
return torch.profiler.profile(
activities=activities,
schedule=schedule,
on_trace_ready=on_trace_ready,
record_shapes=self.config.record_shapes,
profile_memory=self.config.profile_memory,
with_stack=self.config.with_stack,
with_flops=self.config.with_flops,
)
def __call__(self, batch: TensorDictBase | None = None) -> None:
if self._stopped:
return
if self._profiler is None:
self._profiler = self._build_profiler()
if self._profiler is None:
self._stopped = True
return
self._profiler.start()
self._profiler.step()
self._rollout_count += 1
if self._rollout_count >= self.config.num_rollouts:
self.stop()
def stop(self) -> None:
"""Stop the underlying profiler (idempotent)."""
if self._profiler is not None and not self._stopped:
self._profiler.stop()
self._stopped = True
[docs]
class BaseCollector(IterableDataset, metaclass=abc.ABCMeta):
"""Base class for data collectors.
Keyword Args:
trajs_per_batch (int, optional): When set, the collector yields batches
of exactly this many complete trajectories instead of fixed-frame
batches. By default each yielded :class:`~tensordict.TensorDict`
has shape ``(trajs_per_batch, max_traj_len)``, zero-padded along
time, and includes a ``("collector", "mask")`` boolean field
marking valid time steps; with ``traj_format="cat"`` the
trajectories are instead concatenated along time into a flat,
unpadded batch. Trajectories that span multiple internal
collection steps are reassembled automatically.
``frames_per_batch`` still controls how often the environment is
polled internally, but the output batch size is determined by
``trajs_per_batch``.
(:class:`~torchrl.collectors.AsyncBatchedCollector` exposes the
same capability through its ``yield_completed_trajectories``
flag.)
**Replay buffer integration**
For compatibility, combining this argument with a
``replay_buffer`` while leaving ``replay_write_mode=None`` selects
complete-trajectory replay writes. New code should use
``replay_write_mode="trajectory"`` instead.
.. important::
When using a **multi-process** collector with a shared replay
buffer and a :class:`~torchrl.data.SliceSampler`, setting
``replay_write_mode="trajectory"`` is strongly recommended.
Without it,
different workers write batches independently and adjacent
frames in the buffer can come from unrelated episodes without
an intervening ``done`` signal, causing the sampler to draw
slices that cross trajectory boundaries.
**Completeness guarantee**: only trajectories whose last step has
``("next", "done") == True`` are written to the buffer. Partial
trajectories (episodes still in flight) are held internally until
the episode terminates. This means every trajectory in the buffer
is guaranteed to be a complete episode segment.
**Batched environments**: when the environment has a batch size > 1
(e.g. :class:`~torchrl.envs.SerialEnv`), steps are disassembled by
``traj_id`` and each trajectory is written individually as a flat
sequence. The buffer storage should use ``ndim=1`` — ``ndim=2``
is incompatible because variable-length trajectories cannot fill a
fixed second dimension.
**Multi-process and distributed collectors**: trajectory replay
writes are supported for
:class:`~torchrl.collectors.MultiSyncCollector`,
:class:`~torchrl.collectors.MultiAsyncCollector`,
:class:`~torchrl.collectors.distributed.RayCollector`, and
:class:`~torchrl.collectors.distributed.RPCCollector` (through
its remote ``collector_kwargs``).
Trajectory assembly is delegated to each worker's inner collector,
which calls :meth:`_iter_by_trajectories` independently and writes
complete trajectories to the shared replay buffer. Local process
and Ray collectors support both iteration (``for data in
collector``) and asynchronous ``start()``; RPC uses its remote
collector iteration loop.
.. code-block:: python
rb = ReplayBuffer(
storage=LazyTensorStorage(10_000),
sampler=SliceSampler(slice_len=16, end_key=("next", "done")),
shared=True,
)
collector = MultiSyncCollector(
[env_fn] * 4, policy,
replay_buffer=rb,
frames_per_batch=200,
total_frames=-1,
replay_write_mode="trajectory",
)
collector.start() # workers fill rb with complete trajectories
Defaults to ``None`` (fixed-frame batches).
trajs_per_write (int, optional): In trajectory replay-write mode,
write this many completed trajectories to the
buffer per ``extend`` call. Larger values reduce Python overhead
for highly batched environments. For example, if 10 complete
trajectories are queued for replay-buffer insertion,
``trajs_per_write=2`` makes 5 writes, while
``trajs_per_write=10`` or larger makes 1 write. Defaults to
``None`` (write all currently queued completed trajectories).
replay_write_mode (``"rollout"``, ``"trajectory"``, optional): Controls
how a collector writes to ``replay_buffer``. ``"rollout"`` keeps
the fixed-frame rollout layout. ``"trajectory"`` retains
in-flight episodes and writes only completed trajectories as flat
1-D sequences. ``trajs_per_write`` optionally groups completed
trajectories into each ``extend`` call. Defaults to ``None``,
which preserves the legacy behavior: ``replay_buffer`` combined
with ``trajs_per_batch`` selects trajectory writes, and other
replay-buffer configurations select rollout writes.
Explicit replay write modes cannot be combined with
``trajs_per_batch``. The latter controls the number of completed
trajectories in batches yielded without a replay buffer.
traj_format (str, optional): layout of the batches yielded when
``trajs_per_batch`` is set. ``"padded"`` stacks the
trajectories into a ``(trajs_per_batch, max_traj_len)`` batch,
zero-padded along time, with a ``("collector", "mask")`` entry
marking the valid steps. ``"cat"`` concatenates them along time
into a flat, unpadded ``[sum_i T_i]`` batch — no mask, no wasted
memory on padding; trajectories are contiguous, delimited by
``("next", "done")`` (``True`` at the last step of each, by the
completeness guarantee) and ``("collector", "traj_ids")``.
``"cat"`` matches the layout the replay-buffer write path uses
and the one :class:`~torchrl.data.SliceSampler` expects. Has no
effect on replay-buffer writes (always flat); raises if set
without ``trajs_per_batch``. Defaults to ``None``, which
currently resolves to ``"padded"`` and emits a
:class:`FutureWarning` when ``trajs_per_batch`` batches are
yielded without an explicit choice: the default will change to
``"cat"`` in torchrl v0.16.
"""
_task = None
_iterator = None
_iteration_started = False
total_frames: int
requested_frames_per_batch: int
frames_per_batch: int
trust_policy: bool
compiled_policy: bool
cudagraphed_policy: bool
_weight_updater: WeightUpdaterBase | None = None
_weight_sync_schemes: dict[str, WeightSyncScheme] | None = None
verbose: bool = False
_profile_config: ProfileConfig | None = None
trajs_per_batch: int | None = None
trajs_per_write: int | None = None
replay_write_mode: Literal["rollout", "trajectory"] | None = None
traj_format: Literal["padded", "cat"] = "padded"
_pre_collect_hook: Callable[[], None] | None = None
_post_collect_hook: Callable[[TensorDictBase], None] | None = None
def __init__(
self,
*,
pre_collect_hook: Callable[[], None] | None = None,
post_collect_hook: Callable[[TensorDictBase], None] | None = None,
):
self._pre_collect_hook = pre_collect_hook
self._post_collect_hook = post_collect_hook
self._collector_progress = _CollectorProgress()
self._collector_progress_worker_idx = 0
@property
def pre_collect_hook(self) -> Callable[[], None] | None:
"""Get the pre-collection hook.
Returns:
A callable to be executed before each rollout, or None.
"""
return self._pre_collect_hook
@pre_collect_hook.setter
def pre_collect_hook(self, hook: Callable[[], None] | None) -> None:
"""Set the pre-collection hook.
Args:
hook: A callable to be executed before each rollout.
"""
self._pre_collect_hook = hook
@property
def post_collect_hook(self) -> Callable[[TensorDictBase], None] | None:
"""Get the post-collection hook.
Returns:
A callable to be executed after each rollout, receiving the collected
TensorDict as argument, or None.
"""
return self._post_collect_hook
@post_collect_hook.setter
def post_collect_hook(self, hook: Callable[[TensorDictBase], None] | None) -> None:
"""Set the post-collection hook.
Args:
hook: A callable to be executed after each rollout, receiving the
collected TensorDict as argument.
"""
self._post_collect_hook = hook
[docs]
def set_post_collect_hook(
self, hook: Callable[[TensorDictBase], None] | None
) -> None:
"""Method form of the ``post_collect_hook`` setter.
Exposed because Ray actor handles can call methods (`actor.method.remote(...)`)
but cannot directly invoke property setters. Keeping the actual setter
for in-process use and this method for remote-actor use.
"""
self.post_collect_hook = hook
[docs]
def stats(
self,
workers: Literal["aggregate", "per_worker", "both"] = "aggregate",
) -> dict[str, int | float | bool]:
"""Returns a cheap, serializable snapshot of the collector's progress.
The snapshot only contains scalar counters and gauges: it never
includes policy, environment or batch data, does not modify the
collector state and is safe to call while the collector is running.
Cumulative counters such as ``frames`` are meant to be converted into
rates by an external monitor such as
:class:`~torchrl.record.loggers.monitoring.LoggerMonitor`.
Entries are only present when the corresponding state exists on the
collector:
- ``"frames"``: total number of frames delivered so far (the existing
collector-specific semantics are unchanged);
- ``"stepped_frames"``: environment transitions collected, including
frames still held in an unfinished trajectory;
- ``"trajectory_completed_frames"``: frames belonging to trajectories
that have reached a terminal boundary;
- ``"trajectory_pending_frames"``: current in-flight trajectory frames;
- ``"replay_written_frames"``: frames successfully inserted in the
attached replay buffer;
- ``"completed_trajectories"``: trajectories that reached a terminal
boundary;
- ``"batches"``: number of batches delivered so far;
- ``"total_frames"``: requested total frames (absent for endless collectors);
- ``"completed"``: whether the frame budget has been reached;
- ``"requested_frames_per_batch"``: the per-batch frame budget;
- ``"policy_version"``: current policy version, when the collector
tracks it with an integer version.
The progress entries are cumulative except for
``"trajectory_pending_frames"``, which is a gauge. Reset and shutdown
drop in-flight trajectory assembly, so they clear that gauge without
changing the cumulative entries. Checkpoints restore the cumulative
entries but start the gauge at zero because collector checkpoints do
not serialize the environment state or partial trajectory payloads.
Args:
workers (str, optional): controls the worker view. With
``"aggregate"`` (default), only coordinator-side counters are
reported and no worker communication happens. With ``"per_worker"``
or ``"both"``, each worker is queried and its snapshot is
namespaced as ``"worker_<idx>/<metric>"``. For multi-worker
collectors, ``"workers"`` and ``"workers_alive"`` are always
reported. Per-worker queries share the control channel and must
not race with concurrent weight updates or other control calls.
Ray collectors retain their transport-specific timeout and
remote aggregation behavior.
Examples:
>>> from torchrl.collectors import Collector
>>> from torchrl.envs import GymEnv
>>> from torchrl.envs.utils import RandomPolicy
>>> env = GymEnv("Pendulum-v1")
>>> collector = Collector(
... env,
... RandomPolicy(env.action_spec),
... frames_per_batch=10,
... total_frames=20,
... )
>>> for batch in collector:
... print(collector.stats()["frames"])
10
20
"""
if workers not in ("aggregate", "per_worker", "both"):
raise ValueError(
f"workers must be one of 'aggregate', 'per_worker' or 'both', got {workers!r}."
)
stats: dict[str, int | float | bool] = {}
if workers in ("aggregate", "both"):
progress = getattr(self, "_collector_progress", None)
if progress is not None:
stats.update(
progress.snapshot(
None
if getattr(self, "_collector_progress_aggregate", False)
else getattr(self, "_collector_progress_worker_idx", 0)
)
)
frames = getattr(self, "_frames", None)
if frames is not None:
stats["frames"] = int(frames)
iters = getattr(self, "_iter", None)
if iters is not None:
stats["batches"] = int(iters) + 1
total_frames = getattr(self, "total_frames", None)
if isinstance(total_frames, int) and total_frames >= 0:
stats["total_frames"] = total_frames
if frames is not None:
stats["completed"] = bool(frames >= total_frames)
requested = getattr(
self,
"requested_frames_per_batch",
getattr(self, "frames_per_batch", None),
)
if isinstance(requested, int):
stats["requested_frames_per_batch"] = requested
try:
version = self.policy_version
except (AttributeError, RuntimeError):
version = None
if isinstance(version, int):
stats["policy_version"] = version
if hasattr(self, "procs"):
stats["workers"] = int(self.num_workers)
if self.procs:
stats["workers_alive"] = sum(
int(proc.is_alive()) for proc in self.procs
)
if workers in ("per_worker", "both"):
for idx, worker_stats in enumerate(self.map_fn("stats")):
for key, value in worker_stats.items():
stats[f"worker_{idx}/{key}"] = value
return stats
def _record_stepped_frames(self, frames: int) -> None:
self._collector_progress.increment_stepped(
self._collector_progress_worker_idx,
frames,
trajectory_pending=(
self.trajs_per_batch is not None
or getattr(self, "replay_write_mode", None) == "trajectory"
),
)
def _record_trajectory_completion(self, frames: int, trajectories: int) -> None:
self._collector_progress.record_trajectory_completion(
self._collector_progress_worker_idx, frames, trajectories
)
def _record_pending_trajectory_frames(self, frames: int) -> None:
self._collector_progress.record_trajectory_pending(
self._collector_progress_worker_idx, frames
)
def _record_replay_write(self, frames: int) -> None:
self._collector_progress.record_replay_write(
self._collector_progress_worker_idx, frames
)
def _clear_pending_trajectory_progress(self) -> None:
self._collector_progress.clear_pending(self._collector_progress_worker_idx)
def _progress_state_dict(self) -> dict[str, int]:
return self._collector_progress.snapshot(self._collector_progress_worker_idx)
def _load_progress_state_dict(self, state: dict[str, int] | None) -> None:
self._collector_progress.load_snapshot(
self._collector_progress_worker_idx, state or {}
)
[docs]
def enable_profile(
self,
*,
workers: list[int] | None = None,
num_rollouts: int = 3,
warmup_rollouts: int = 1,
save_path: str | Path | None = None,
activities: list[str] | None = None,
record_shapes: bool = True,
profile_memory: bool = False,
with_stack: bool = True,
with_flops: bool = False,
on_trace_ready: Callable | None = None,
) -> None:
"""Enable profiling for collector worker rollouts.
This method configures the collector to profile rollouts using PyTorch's
profiler. For multi-process collectors, profiling happens in the worker
processes. For single-process collectors (Collector), profiling happens
in the main process.
Args:
workers: List of worker indices to profile. Defaults to [0].
For single-process collectors, this is ignored.
num_rollouts: Total number of rollouts to run the profiler for
(including warmup). Profiling stops after this many rollouts.
Defaults to 3.
warmup_rollouts: Number of rollouts to skip before starting actual
profiling. Useful for JIT/compile warmup. The profiler runs
but discards data during warmup. Defaults to 1.
save_path: Path to save the profiling trace. Supports {worker_idx}
placeholder for worker-specific files. If None, traces are
saved to "./collector_profile_{worker_idx}.json".
activities: List of profiler activities ("cpu", "cuda").
Defaults to ["cpu", "cuda"].
record_shapes: Whether to record tensor shapes. Defaults to True.
profile_memory: Whether to profile memory usage. Defaults to False.
with_stack: Whether to record Python stack traces. Defaults to True.
with_flops: Whether to compute FLOPS. Defaults to False.
on_trace_ready: Optional callback when trace is ready. If None,
traces are exported to Chrome trace format at save_path.
Raises:
RuntimeError: If called after iteration has started.
ValueError: If num_rollouts <= warmup_rollouts.
Example:
>>> from torchrl.collectors import MultiSyncCollector
>>> collector = MultiSyncCollector(
... create_env_fn=[make_env] * 4,
... policy=policy,
... frames_per_batch=1000,
... total_frames=100000,
... )
>>> collector.enable_profile(
... workers=[0],
... num_rollouts=5,
... warmup_rollouts=2,
... save_path="./traces/worker_{worker_idx}.json",
... )
>>> # Worker 0 will be profiled for rollouts 2, 3, 4
>>> for data in collector:
... train(data)
>>> collector.shutdown()
Note:
- Profiling adds overhead, so only profile specific workers
- The trace file can be viewed in Chrome's trace viewer
(chrome://tracing) or with PyTorch's TensorBoard plugin
- For multi-process collectors, this must be called BEFORE
iteration starts as it needs to configure workers
"""
if self._iteration_started:
raise RuntimeError(
"Cannot enable profiling after iteration has started. "
"Call enable_profile() before iterating over the collector."
)
if workers is None:
workers = [0]
if activities is None:
activities = ["cpu", "cuda"]
config = ProfileConfig(
workers=workers,
num_rollouts=num_rollouts,
warmup_rollouts=warmup_rollouts,
save_path=save_path,
activities=activities,
record_shapes=record_shapes,
profile_memory=profile_memory,
with_stack=with_stack,
with_flops=with_flops,
on_trace_ready=on_trace_ready,
)
self._profile_config = config
self._install_profile_hooks(config)
def _install_profile_hooks(self, config: ProfileConfig) -> None:
"""Install the per-process profiler hook.
Default implementation handles the single-process :class:`Collector`
by saving the current ``post_collect_hook`` and replacing it with a
:class:`_ProfilerHook`. Multi-process / Ray collectors override this
to fan out per-worker hooks (each worker gets its own ``worker_idx``).
"""
self._saved_post_collect_hook = self._post_collect_hook
self.post_collect_hook = _ProfilerHook(config, worker_idx=0)
[docs]
def disable_profile(self) -> None:
"""Stop any in-flight profiler and restore the prior ``post_collect_hook``.
Safe to call when profiling was never enabled (becomes a no-op). When
the profiler was already self-stopped after ``num_rollouts``, this just
clears the hook and restores any user-set ``post_collect_hook``.
"""
if self._profile_config is None:
return
try:
self._uninstall_profile_hooks(self._profile_config)
finally:
self._profile_config = None
def _uninstall_profile_hooks(self, config: ProfileConfig) -> None:
"""Single-process default uninstall: stop hook locally and restore."""
hook = self._post_collect_hook
if isinstance(hook, _ProfilerHook):
hook.stop()
self._post_collect_hook = getattr(self, "_saved_post_collect_hook", None)
self._saved_post_collect_hook = None
@property
def profile_config(self) -> ProfileConfig | None:
"""Get the profiling configuration.
Returns:
ProfileConfig if profiling is enabled, None otherwise.
"""
return self._profile_config
@property
def weight_updater(self) -> WeightUpdaterBase:
return self._weight_updater
@weight_updater.setter
def weight_updater(self, value: WeightUpdaterBase | None):
if value is not None:
if not isinstance(value, WeightUpdaterBase) and callable(
value
): # Fall back to default constructor
value = value()
value.register_collector(self)
if value.collector is not self:
raise RuntimeError("Failed to register collector.")
self._weight_updater = value
@property
def worker_idx(self) -> int | None:
"""Get the worker index for this collector.
Returns:
The worker index (0-indexed).
Raises:
RuntimeError: If worker_idx has not been set.
"""
if not hasattr(self, "_worker_idx"):
raise RuntimeError(
"worker_idx has not been set. This collector may not have been "
"initialized as a worker in a distributed setup."
)
return self._worker_idx
@worker_idx.setter
def worker_idx(self, value: int | None) -> None:
"""Set the worker index for this collector.
Args:
value: The worker index (0-indexed) or None.
"""
self._worker_idx = value
[docs]
def cascade_execute(self, attr_path: str, *args, **kwargs) -> Any:
"""Execute a method on a nested attribute of this collector.
This method allows remote callers to invoke methods on nested attributes
of the collector without needing to know the full structure. It's particularly
useful for calling methods on weight sync schemes from the sender side.
Args:
attr_path: Full path to the callable, e.g.,
"_receiver_schemes['model_id']._set_dist_connection_info"
*args: Positional arguments to pass to the method.
**kwargs: Keyword arguments to pass to the method.
Returns:
The return value of the method call.
Examples:
>>> collector.cascade_execute(
... "_receiver_schemes['policy']._set_dist_connection_info",
... connection_info_ref,
... worker_idx=0
... )
"""
attr = _resolve_attr(self, attr_path)
if callable(attr):
return attr(*args, **kwargs)
else:
if args or kwargs:
raise ValueError(
f"Arguments and keyword arguments are not supported for non-callable attributes. Got {args} and {kwargs} for {attr_path}"
)
return attr
[docs]
def get_distant_attr(self, attr: str) -> Any:
"""Get a nested attribute of this collector.
This method allows remote callers to retrieve attributes from nested
structures of the collector without needing to know the full structure.
Args:
attr: Full path to the attribute, e.g.,
"_receiver_schemes['model_id'].some_attribute"
Returns:
The value of the attribute.
Examples:
>>> collector.get_distant_attr("_receiver_schemes['policy']._sync_interval")
"""
return _resolve_attr(self, attr)
def _dump_env_transform(self, step: int | None = None) -> None:
"""Dump the environment transform when it supports ``dump``."""
env = getattr(self, "env", None)
transform = getattr(env, "transform", None)
dump = getattr(transform, "dump", None)
if callable(dump):
dump(step=step)
[docs]
def map_fn(
self,
method_name: str,
list_of_args: list[tuple] | None = None,
list_of_kwargs: list[dict] | None = None,
) -> list[Any]:
"""Apply a method to each set of arguments.
This method executes a method on the collector with different arguments,
returning a list of results.
Args:
method_name: Name of the method to call on the collector.
list_of_args: List of positional argument tuples. Each tuple
contains the arguments for one call.
list_of_kwargs: List of keyword argument dicts. Each dict
contains the kwargs for one call.
Returns:
List of return values from each method call.
Examples:
>>> # Call a method with different arguments
>>> collector.map_fn("update_policy_weights_", list_of_args=[(weights1,), (weights2,)])
>>>
>>> # Call with kwargs
>>> collector.map_fn("update_policy_weights_", list_of_kwargs=[{"weights": w1}, {"weights": w2}])
"""
if list_of_args is None:
list_of_args = [()] * len(list_of_kwargs) if list_of_kwargs else [()]
if list_of_kwargs is None:
list_of_kwargs = [{}] * len(list_of_args)
if len(list_of_args) != len(list_of_kwargs):
raise ValueError(
f"list_of_args and list_of_kwargs must have the same length. "
f"Got {len(list_of_args)} and {len(list_of_kwargs)}"
)
method = _resolve_attr(self, method_name)
if not callable(method):
raise AttributeError(f"Attribute {method_name} is not callable.")
results = []
for args, kwargs in zip(list_of_args, list_of_kwargs):
results.append(method(*args, **kwargs))
return results
def _get_policy_and_device(
self,
policy: Callable[[Any], Any] | None = None,
policy_device: Any = NO_DEFAULT,
env_maker: Any | None = None,
env_maker_kwargs: dict[str, Any] | None = None,
) -> tuple[TensorDictModule, None | Callable[[], dict]]:
"""Util method to get a policy and its device given the collector __init__ inputs.
We want to copy the policy and then move the data there, not call policy.to(device).
Args:
policy (TensorDictModule, optional): a policy to be used
policy_device (torch.device, optional): the device where the policy should be placed.
Defaults to self.policy_device
env_maker (a callable or a batched env, optional): the env_maker function for this device/policy pair.
env_maker_kwargs (a dict, optional): the env_maker function kwargs.
"""
if policy_device is NO_DEFAULT:
policy_device = self.policy_device
if not policy_device:
return policy, None
if isinstance(policy, nn.Module):
param_and_buf = TensorDict.from_module(policy, as_module=True)
else:
# Because we want to reach the warning
param_and_buf = TensorDict()
i = -1
for p in param_and_buf.values(True, True):
i += 1
if p.device != policy_device:
# Then we need casting
break
else:
if i == -1 and not self.trust_policy:
# We trust that the policy policy device is adequate
warnings.warn(
"A policy device was provided but no parameter/buffer could be found in "
"the policy. Casting to policy_device is therefore impossible. "
"The collector will trust that the devices match. To suppress this "
"warning, set `trust_policy=True` when building the collector."
)
return policy, None
# Create a stateless policy, then populate this copy with params on device
def get_original_weights(policy=policy):
td = TensorDict.from_module(policy)
return td.data
# We need to use ".data" otherwise buffers may disappear from the `get_original_weights` function
with param_and_buf.data.to("meta").to_module(
policy, preserve_module_state=False
):
policy_new_device = deepcopy(policy)
param_and_buf_new_device = param_and_buf.apply(
functools.partial(_map_weight, policy_device=policy_device),
filter_empty=False,
)
param_and_buf_new_device.to_module(
policy_new_device, preserve_module_state=False
)
# Sanity check
if set(TensorDict.from_module(policy_new_device).keys(True, True)) != set(
get_original_weights().keys(True, True)
):
raise RuntimeError("Failed to map weights. The weight sets mismatch.")
return policy_new_device, get_original_weights
[docs]
def start(self):
"""Starts the collector for asynchronous data collection.
This method initiates the background collection of data, allowing for decoupling of data collection and training.
The collected data is typically stored in a replay buffer passed during the collector's initialization.
.. note:: After calling this method, it's essential to shut down the collector using :meth:`~.async_shutdown`
when you're done with it to free up resources.
.. warning:: Asynchronous data collection can significantly impact training performance due to its decoupled nature.
Ensure you understand the implications for your specific algorithm before using this mode.
Raises:
NotImplementedError: If not implemented by a subclass.
"""
raise NotImplementedError(
f"Collector start() is not implemented for {type(self).__name__}."
)
[docs]
@contextlib.contextmanager
def pause(self):
"""Context manager that pauses the collector if it is running free."""
raise NotImplementedError(
f"Collector pause() is not implemented for {type(self).__name__}."
)
[docs]
def async_shutdown(
self, timeout: float | None = None, close_env: bool = True
) -> None:
"""Shuts down the collector when started asynchronously with the `start` method.
Args:
timeout (float, optional): The maximum time to wait for the collector to shutdown.
close_env (bool, optional): If True, the collector will close the contained environment.
Defaults to `True`.
.. seealso:: :meth:`~.start`
"""
return self.shutdown(timeout=timeout, close_env=close_env)
def _extract_weights_if_needed(self, weights: Any, model_id: str) -> Any:
"""Extract weights from a model if needed.
For the new weight sync scheme system, weight preparation is handled
by the scheme's prepare_weights() method. This method now only handles
legacy weight updater cases.
Args:
weights: Either already-extracted weights or a model to extract from.
model_id: The model identifier for resolving string paths.
Returns:
Extracted weights in the appropriate format.
"""
# New weight sync schemes handle preparation themselves
if self._weight_sync_schemes:
# Just pass through - WeightSender will call scheme.prepare_weights()
return weights
# Legacy weight updater path
return self._legacy_extract_weights(weights, model_id)
def _legacy_extract_weights(self, weights: Any, model_id: str) -> Any:
"""Legacy weight extraction for old weight updater system.
Args:
weights: Either already-extracted weights or a model to extract from.
model_id: The model identifier.
Returns:
Extracted weights.
"""
if weights is None:
if model_id == "policy" and hasattr(self, "policy_weights"):
return self.policy_weights
elif model_id == "policy" and hasattr(self, "_policy_weights_dict"):
policy_device = (
self.policy_device
if not isinstance(self.policy_device, (list, tuple))
else self.policy_device[0]
)
return self._policy_weights_dict.get(policy_device)
return None
return weights
@property
def _legacy_weight_updater(self) -> bool:
return self._weight_updater is not None
# Overloads for update_policy_weights_ to support multiple calling conventions
@overload
def update_policy_weights_(
self,
policy_or_weights: TensorDictBase | TensorDictModuleBase | nn.Module | dict,
/,
) -> None:
...
@overload
def update_policy_weights_(
self,
policy_or_weights: TensorDictBase | TensorDictModuleBase | nn.Module | dict,
/,
*,
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
model_id: str | None = None,
) -> None:
...
@overload
def update_policy_weights_(
self,
*,
weights: TensorDictBase | dict,
model_id: str | None = None,
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
) -> None:
...
@overload
def update_policy_weights_(
self,
*,
policy: TensorDictModuleBase | nn.Module,
model_id: str | None = None,
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
) -> None:
...
@overload
def update_policy_weights_(
self,
*,
weights_dict: dict[
str, TensorDictBase | TensorDictModuleBase | nn.Module | dict
],
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
) -> None:
...
[docs]
def update_policy_weights_(
self,
policy_or_weights: (
TensorDictBase | TensorDictModuleBase | nn.Module | dict | None
) = None,
*,
weights: TensorDictBase | dict | None = None,
policy: TensorDictModuleBase | nn.Module | None = None,
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
model_id: str | None = None,
weights_dict: dict[str, Any] | None = None,
**kwargs,
) -> None:
"""Update policy weights for the data collector.
This method synchronizes the policy weights used by the collector with the latest
trained weights. It supports both local and remote weight updates, depending on
the collector configuration.
The method accepts weights in multiple forms for convenience:
Examples:
>>> # Pass policy module as positional argument
>>> collector.update_policy_weights_(policy_module)
>>>
>>> # Pass TensorDict weights as positional argument
>>> collector.update_policy_weights_(weights_tensordict)
>>>
>>> # Use keyword arguments for clarity
>>> collector.update_policy_weights_(weights=weights_td, model_id="actor")
>>> collector.update_policy_weights_(policy=actor_module, model_id="actor")
>>>
>>> # Update multiple models atomically
>>> collector.update_policy_weights_(weights_dict={
... "actor": actor_weights,
... "critic": critic_weights,
... })
>>>
>>> # Per-worker weight updates (for distinct policy factories)
>>> # Each worker can have independently updated weights
>>> collector.update_policy_weights_({
... 0: worker_0_weights,
... 1: worker_1_weights,
... 2: worker_2_weights,
... })
Args:
policy_or_weights: The weights to update with. Can be:
- ``nn.Module``: A policy module whose weights will be extracted
- ``TensorDictModuleBase``: A TensorDict module whose weights will be extracted
- ``TensorDictBase``: A TensorDict containing weights
- ``dict``: A regular dict containing weights
- ``dict[int, TensorDictBase]``: Per-worker weights where keys are worker indices.
This is used with distinct policy factories where each worker has independent weights.
- ``None``: Will try to get weights from server using ``_get_server_weights()``
Keyword Args:
weights: Alternative to positional argument. A TensorDict or dict containing
weights to update. Cannot be used together with ``policy_or_weights`` or ``policy``.
policy: Alternative to positional argument. An ``nn.Module`` or ``TensorDictModuleBase``
whose weights will be extracted. Cannot be used together with ``policy_or_weights``
or ``weights``.
worker_ids: Identifiers for the workers to update. Relevant when the collector
has multiple workers. Can be int, list of ints, device, or list of devices.
model_id: The model identifier to update (default: ``"policy"``).
Cannot be used together with ``weights_dict``.
weights_dict: Dictionary mapping model_id to weights for updating
multiple models atomically. Keys should match model_ids registered in
``weight_sync_schemes``. Cannot be used together with ``model_id``,
``policy_or_weights``, ``weights``, or ``policy``.
Raises:
TypeError: If ``worker_ids`` is provided but no ``weight_updater`` is configured.
ValueError: If conflicting parameters are provided.
.. note:: Users should extend the ``WeightUpdaterBase`` classes to customize
the weight update logic for specific use cases.
.. seealso:: :class:`~torchrl.collectors.LocalWeightsUpdaterBase` and
:meth:`~torchrl.collectors.RemoteWeightsUpdaterBase`.
"""
# Handle the different keyword argument forms
if weights is not None:
if policy_or_weights is not None:
raise ValueError(
"Cannot specify both positional 'policy_or_weights' and keyword 'weights'"
)
if policy is not None:
raise ValueError("Cannot specify both 'weights' and 'policy'")
policy_or_weights = weights
if policy is not None:
if policy_or_weights is not None:
raise ValueError(
"Cannot specify both positional 'policy_or_weights' and keyword 'policy'"
)
policy_or_weights = policy
if self._legacy_weight_updater:
return self._legacy_weight_update_impl(
policy_or_weights=policy_or_weights,
worker_ids=worker_ids,
model_id=model_id,
weights_dict=weights_dict,
**kwargs,
)
else:
return self._weight_update_impl(
policy_or_weights=policy_or_weights,
worker_ids=worker_ids,
model_id=model_id,
weights_dict=weights_dict,
**kwargs,
)
def _legacy_weight_update_impl(
self,
policy_or_weights: TensorDictBase | TensorDictModuleBase | dict | None = None,
*,
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
model_id: str | None = None,
weights_dict: dict[str, Any] | None = None,
**kwargs,
) -> None:
if weights_dict is not None:
raise ValueError("weights_dict is not supported with legacy weight updater")
if model_id is not None:
raise ValueError("model_id is not supported with legacy weight updater")
# Fall back to old weight updater system
self.weight_updater(
policy_or_weights=policy_or_weights, worker_ids=worker_ids, **kwargs
)
def _weight_update_impl(
self,
policy_or_weights: TensorDictBase | TensorDictModuleBase | dict | None = None,
*,
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
model_id: str | None = None,
weights_dict: dict[str, Any] | None = None,
**kwargs,
) -> None:
if "policy_weights" in kwargs:
warnings.warn(
"`policy_weights` is deprecated. Use `policy_or_weights` instead.",
DeprecationWarning,
)
policy_or_weights = kwargs.pop("policy_weights")
if weights_dict is not None and model_id is not None:
raise ValueError("Cannot specify both 'weights_dict' and 'model_id'")
if weights_dict is not None and policy_or_weights is not None:
raise ValueError(
"Cannot specify both 'weights_dict' and 'policy_or_weights'"
)
if self._weight_sync_schemes:
if model_id is None:
model_id = "policy"
if policy_or_weights is not None and weights_dict is None:
# Use model_id as the key, not hardcoded "policy"
weights_dict = {model_id: policy_or_weights}
elif weights_dict is None:
weights_dict = {model_id: policy_or_weights}
for target_model_id, weights in weights_dict.items():
if target_model_id not in self._weight_sync_schemes:
raise KeyError(
f"Model '{target_model_id}' not found in registered weight sync schemes. "
f"Available models: {list(self._weight_sync_schemes.keys())}"
)
processed_weights = self._extract_weights_if_needed(
weights, target_model_id
)
# Use new send() API with worker_ids support
scheme = self._weight_sync_schemes.get(target_model_id)
if not isinstance(scheme, WeightSyncScheme):
raise TypeError(f"Expected WeightSyncScheme, got {target_model_id}")
self._send_weights_scheme(
scheme=scheme,
processed_weights=processed_weights,
worker_ids=worker_ids,
model_id=target_model_id,
)
elif self._weight_updater is not None:
# unreachable
raise RuntimeError
else:
# No weight updater configured, try fallback
self._maybe_fallback_update(policy_or_weights, model_id=model_id)
def _maybe_fallback_update(
self,
policy_or_weights: TensorDictBase | TensorDictModuleBase | dict | None = None,
*,
model_id: str | None = None,
) -> None:
"""Fallback weight update when no scheme is configured.
Override in subclasses to provide custom fallback behavior.
The base implementation fails rather than silently accepting an update
that cannot reach a policy.
"""
del policy_or_weights, model_id
raise RuntimeError(
"No weight updater, WeightSyncScheme, or concrete local policy "
"update path is configured for this collector."
)
def _send_weights_scheme(self, *, model_id, scheme, processed_weights, worker_ids):
# method to override if the scheme requires an RPC call to receive the weights
scheme.send(weights=processed_weights, worker_ids=worker_ids)
def _receive_weights_scheme(self, model_version: int | None = None):
"""Receive weights for all registered receiver schemes.
scheme.receive() handles both applying weights locally and cascading
to sub-collectors via context.update_policy_weights_().
"""
if not hasattr(self, "_receiver_schemes"):
raise RuntimeError("No receiver schemes registered.")
for scheme in self._receiver_schemes.values():
scheme.receive()
self._set_received_policy_version(model_version)
def _connect_weights_scheme(self, model_version: int | None = None) -> None:
"""Connect all registered receiver schemes for initial publication."""
if not hasattr(self, "_receiver_schemes"):
raise RuntimeError("No receiver schemes registered.")
for scheme in self._receiver_schemes.values():
if not scheme.synchronized_on_receiver:
scheme.connect(worker_idx=self.worker_idx)
self._set_received_policy_version(model_version)
def _set_received_policy_version(self, model_version: int | None) -> None:
"""Apply the sender's semantic policy version after a weight receive."""
if model_version is None:
return
tracker = getattr(self, "policy_version_tracker", None)
if tracker is not None:
tracker.version = int(model_version)
def _weight_sync_signature(
self, model_id: str
) -> tuple[tuple[tuple[str, ...], tuple[int, ...], str], ...]:
"""Return the ordered tensor schema expected by a weight receiver."""
get_model = getattr(self, "get_model", None)
if not callable(get_model):
raise RuntimeError(
f"{type(self).__name__} cannot resolve weight model {model_id!r}."
)
weights = TensorDict.from_module(get_model(model_id))
return _weight_tensor_signature(weights)
# Overloads for receive_weights to support multiple calling conventions
@overload
def receive_weights(self) -> None:
...
@overload
def receive_weights(
self,
policy_or_weights: TensorDictBase | TensorDictModuleBase | nn.Module | dict,
/,
) -> None:
...
@overload
def receive_weights(
self,
*,
weights: TensorDictBase | dict,
) -> None:
...
@overload
def receive_weights(
self,
*,
policy: TensorDictModuleBase | nn.Module,
) -> None:
...
[docs]
def receive_weights(
self,
policy_or_weights: (
TensorDictBase | TensorDictModuleBase | nn.Module | dict | None
) = None,
*,
weights: TensorDictBase | dict | None = None,
policy: TensorDictModuleBase | nn.Module | None = None,
) -> None:
"""Receive and apply weights to the collector's policy.
This method applies weights to the local policy. When receiver schemes are
registered, it delegates to those schemes. Otherwise, it directly applies
the provided weights.
The method accepts weights in multiple forms for convenience:
Examples:
>>> # Receive from registered schemes (distributed collectors)
>>> collector.receive_weights()
>>>
>>> # Apply weights from a policy module (positional)
>>> collector.receive_weights(trained_policy)
>>>
>>> # Apply weights from a TensorDict (positional)
>>> collector.receive_weights(weights_tensordict)
>>>
>>> # Use keyword arguments for clarity
>>> collector.receive_weights(weights=weights_td)
>>> collector.receive_weights(policy=trained_policy)
Args:
policy_or_weights: The weights to apply. Can be:
- ``nn.Module``: A policy module whose weights will be extracted and applied
- ``TensorDictModuleBase``: A TensorDict module whose weights will be extracted
- ``TensorDictBase``: A TensorDict containing weights
- ``dict``: A regular dict containing weights
- ``None``: Receive from registered schemes or mirror from original policy
Keyword Args:
weights: Alternative to positional argument. A TensorDict or dict containing
weights to apply. Cannot be used together with ``policy_or_weights`` or ``policy``.
policy: Alternative to positional argument. An ``nn.Module`` or ``TensorDictModuleBase``
whose weights will be extracted. Cannot be used together with ``policy_or_weights``
or ``weights``.
Raises:
ValueError: If conflicting parameters are provided or if arguments are passed
when receiver schemes are registered.
"""
# Handle the different keyword argument forms
if weights is not None:
if policy_or_weights is not None:
raise ValueError(
"Cannot specify both positional 'policy_or_weights' and keyword 'weights'"
)
if policy is not None:
raise ValueError("Cannot specify both 'weights' and 'policy'")
policy_or_weights = weights
if policy is not None:
if policy_or_weights is not None:
raise ValueError(
"Cannot specify both positional 'policy_or_weights' and keyword 'policy'"
)
policy_or_weights = policy
if getattr(self, "_receiver_schemes", None) is not None:
if policy_or_weights is not None:
raise ValueError(
"Cannot specify 'policy_or_weights' when using 'receiver_schemes'. Schemes should know how to get the weights."
)
self._receive_weights_scheme()
return
# No weight updater configured
# For single-process collectors, apply weights locally if explicitly provided
if policy_or_weights is not None:
from torchrl.weight_update.weight_sync_schemes import WeightStrategy
# Use WeightStrategy to apply weights properly
strategy = WeightStrategy(extract_as="tensordict")
# Extract weights if needed
if isinstance(policy_or_weights, nn.Module):
weights = strategy.extract_weights(policy_or_weights)
else:
weights = policy_or_weights
# Apply to local policy
if hasattr(self, "policy") and isinstance(self.policy, nn.Module):
strategy.apply_weights(self.policy, weights)
# Otherwise, no action needed - policy is local and changes are immediately visible
[docs]
def register_scheme_receiver(
self,
weight_recv_schemes: dict[str, WeightSyncScheme],
*,
synchronize_weights: bool = True,
): # noqa: D417
"""Set up receiver schemes for this collector to receive weights from parent collectors.
This method initializes receiver schemes and stores them in _receiver_schemes
for later use by _receive_weights_scheme() and receive_weights().
Receiver schemes enable cascading weight updates across collector hierarchies:
- Parent collector sends weights via its weight_sync_schemes (senders)
- Child collector receives weights via its weight_recv_schemes (receivers)
- If child is also a parent (intermediate node), it can propagate to its own children
Args:
weight_recv_schemes (dict[str, WeightSyncScheme]): Dictionary of {model_id: WeightSyncScheme} to set up as receivers.
These schemes will receive weights from parent collectors.
Keyword Args:
synchronize_weights (bool, optional): If True, synchronize weights immediately after registering the schemes.
Defaults to `True`.
"""
# Initialize _receiver_schemes if not already present
if not hasattr(self, "_receiver_schemes"):
self._receiver_schemes = {}
# Initialize each scheme on the receiver side
for model_id, scheme in weight_recv_schemes.items():
previous_scheme = self._receiver_schemes.get(model_id)
if previous_scheme is not None and previous_scheme is not scheme:
previous_scheme.shutdown()
if not scheme.initialized_on_receiver:
if scheme.initialized_on_sender:
raise RuntimeError(
"Weight sync scheme cannot be initialized on both sender and receiver."
)
scheme.init_on_receiver(
model_id=model_id,
context=self,
worker_idx=self.worker_idx,
)
elif scheme.context is None:
# The scheme was already initialized on the receiver (e.g. early,
# by _make_policy_factory which has no access to the inner
# collector yet). Now that we *do* have the collector, set it as
# the context so receiver-side bookkeeping (policy version,
# cascading sub-collector updates) can reach it.
scheme.context = self
# Store the scheme for later use in receive_weights()
self._receiver_schemes[model_id] = scheme
# Perform initial synchronization
if synchronize_weights:
for scheme in weight_recv_schemes.values():
if not scheme.synchronized_on_receiver:
scheme.connect(worker_idx=self.worker_idx)
def __iter__(self) -> Iterator[TensorDictBase]:
# Mark that iteration has started (used by enable_profile check)
self._iteration_started = True
try:
if self.trajs_per_batch is None and (
self.replay_write_mode != "trajectory"
or getattr(self, "_trajectory_writes_in_workers", False)
):
yield from self.iterator()
else:
yield from self._iter_by_trajectories()
except Exception:
self.shutdown()
raise
def _iter_by_trajectories(self) -> Iterator[TensorDictBase]:
"""Yield complete trajectories, either as padded batches or into a replay buffer.
**Without a replay buffer** (the default when iterating directly):
accumulates complete trajectories and yields zero-padded batches of
shape ``(trajs_per_batch, max_traj_len)`` with a
``("collector", "mask")`` boolean field marking valid timesteps.
With ``traj_format="cat"``, the trajectories are instead concatenated
along time into flat, unpadded batches (trajectories delimited by
``("next", "done")``).
**With ``replay_write_mode="trajectory"``**: each complete trajectory
is written to the buffer immediately as a **flat 1-D sequence** of
valid timesteps — no padding and no dependency on
``trajs_per_batch``. The method yields
``None`` on every write, matching the standard replay-buffer collection
convention. This flat storage is directly compatible with
:class:`~torchrl.data.SliceSampler` using
``end_key=("next", "done")``.
**Completeness guarantee**: a trajectory is considered complete when
its last step carries ``("next", "done") == True`` (which equals
``terminated | truncated``). Partial trajectories (episodes still in
flight) are held in the internal ``partial_trajs`` dict and never
written to the buffer or yielded.
**Batched environments**: when the environment has ``batch_size > 1``,
:func:`_traj_ingest` flattens the batch and groups steps by
``("collector", "traj_ids")``, so each trajectory is assembled and
written individually regardless of the environment batch shape.
**Multi-process / distributed collectors**: trajectory assembly is
delegated to each worker's inner collector, and each worker calls this
method independently to write complete trajectories to the shared
replay buffer.
"""
partial_trajs: dict[int, list] = {}
complete_trajs: list = []
# register the assembly containers on the instance so that reset()
# can flush them (the generator keeps references to the same objects)
self._traj_assembly = (partial_trajs, complete_trajs)
rb = getattr(self, "replay_buffer", None)
# _ignore_rb is a single-collector concept; multi-collectors don't have it.
# Default True so that missing attr → has_rb=False (safe for multi-collectors).
has_rb = rb is not None and not getattr(self, "_ignore_rb", True)
if has_rb:
_prev_ignore_rb = self._ignore_rb
# Prevent iterator() from writing raw frames to the replay buffer;
# we will write assembled trajectory sequences instead.
self._ignore_rb = True
try:
for batch in self.iterator():
if batch is None:
continue
if getattr(self, "_collector_progress_pending_on_ingest", False):
self._record_pending_trajectory_frames(batch.numel())
completed_frames, completed_trajectories = _traj_ingest(
batch, partial_trajs, complete_trajs
)
if completed_trajectories:
self._record_trajectory_completion(
completed_frames, completed_trajectories
)
if has_rb:
# Write each complete trajectory to the replay buffer
# immediately as a flat sequence — no padding, no
# accumulation to trajs_per_batch. This avoids the
# pad-then-unpad round-trip and works with any storage
# ndim (variable-length trajectories cannot fill a
# fixed second dimension reliably).
trajs_per_write = getattr(self, "trajs_per_write", None)
if trajs_per_write is None:
trajs_per_write = len(complete_trajs)
else:
trajs_per_write = max(int(trajs_per_write), 1)
while complete_trajs:
trajs = complete_trajs[:trajs_per_write]
del complete_trajs[:trajs_per_write]
if len(trajs) == 1:
trajs = trajs[0]
else:
trajs = torch.cat(trajs, dim=0)
trajs = _maybe_normalize_replay_buffer_tensordict_device(
trajs, rb
)
rb.extend(trajs)
self._record_replay_write(trajs.numel())
yield
else:
while len(complete_trajs) >= self.trajs_per_batch:
traj_batch = _traj_emit(
complete_trajs,
self.trajs_per_batch,
traj_format=getattr(self, "traj_format", "padded"),
)
yield traj_batch
finally:
if has_rb:
self._ignore_rb = _prev_ignore_rb
def next(self):
try:
if self._iterator is None:
self._iterator = iter(self)
out = next(self._iterator)
# if any, we don't want the device ref to be passed in distributed settings
if out is not None and (out.device != "cpu"):
out = out.copy().clear_device_()
return out
except StopIteration:
return None
def _flush_trajectory_assembly(self) -> None:
"""Drop partially-assembled and queued-but-not-yet-yielded trajectories.
Called by ``reset()`` when trajectory assembly is in use: after an
environment reset, steps queued under the pre-reset policy must not
leak into post-reset batches, and stale partial chunks must not be
merged with later episodes that reuse a rebased trajectory id.
"""
assembly = getattr(self, "_traj_assembly", None)
if assembly is not None:
assembly[0].clear()
assembly[1].clear()
self._clear_pending_trajectory_progress()
@abc.abstractmethod
def shutdown(
self,
timeout: float | None = None,
close_env: bool = True,
raise_on_error: bool = True,
) -> None:
raise NotImplementedError
@abc.abstractmethod
def iterator(self) -> Iterator[TensorDictBase]:
raise NotImplementedError
@abc.abstractmethod
def set_seed(self, seed: int, static_seed: bool = False) -> int:
raise NotImplementedError
@abc.abstractmethod
def state_dict(self) -> OrderedDict:
raise NotImplementedError
@abc.abstractmethod
def load_state_dict(self, state_dict: OrderedDict) -> None:
raise NotImplementedError
def _read_compile_kwargs(self, compile_policy, cudagraph_policy):
self.compiled_policy = compile_policy not in (False, None)
self.cudagraphed_policy = cudagraph_policy not in (False, None)
self.compiled_policy_kwargs = (
{} if not isinstance(compile_policy, typing.Mapping) else compile_policy
)
self.cudagraphed_policy_kwargs = (
{} if not isinstance(cudagraph_policy, typing.Mapping) else cudagraph_policy
)
def __repr__(self) -> str:
string = f"{self.__class__.__name__}()"
return string
def __class_getitem__(self, index):
raise NotImplementedError
def __len__(self) -> int:
if self.total_frames > 0:
return -(self.total_frames // -self.requested_frames_per_batch)
raise RuntimeError("Non-terminating collectors do not have a length")
[docs]
def init_updater(self, *args, **kwargs):
"""Initialize the weight updater with custom arguments.
This method passes the arguments to the weight updater's init method.
If no weight updater is set, this is a no-op.
Args:
*args: Positional arguments for weight updater initialization
**kwargs: Keyword arguments for weight updater initialization
"""
if self.weight_updater is not None:
self.weight_updater.init(*args, **kwargs)
