from __future__ import annotations
import _pickle
import abc
import contextlib
import sys
import time
import warnings
from collections import OrderedDict
from collections.abc import Callable, Mapping, Sequence
from typing import Any, Literal
import numpy as np
import torch
from tensordict import TensorDict, TensorDictBase
from tensordict.nn import CudaGraphModule, TensorDictModule, TensorDictModuleBase
from tensordict.utils import _zip_strict
from torch import multiprocessing as mp, nn
from torchrl import logger as torchrl_logger
from torchrl._utils import (
_check_for_faulty_process,
_get_mp_ctx,
_make_process_no_warn_cls,
_mp_sharing_strategy_for_spawn,
_set_mp_start_method_if_unset,
RL_WARNINGS,
)
from torchrl.collectors._base import _ProfilerHook, BaseCollector, ProfileConfig
from torchrl.collectors._constants import (
_Interruptor,
DEFAULT_EXPLORATION_TYPE,
ExplorationType,
INSTANTIATE_TIMEOUT,
)
from torchrl.collectors._runner import _main_async_collector
from torchrl.collectors._single import Collector
from torchrl.collectors.utils import (
_CollectorProgress,
_make_meta_policy_cm,
_TrajectoryPool,
_validate_replay_write_mode,
_validate_traj_format,
)
from torchrl.collectors.weight_update import WeightUpdaterBase
from torchrl.data import ReplayBuffer
from torchrl.data.utils import CloudpickleWrapper, DEVICE_TYPING
from torchrl.envs import EnvBase, EnvCreator
from torchrl.envs.transforms import PolicyVersion
from torchrl.weight_update import (
MultiProcessWeightSyncScheme,
SharedMemWeightSyncScheme,
WeightSyncScheme,
)
from torchrl.weight_update.utils import _resolve_model
class _MultiCollectorMeta(abc.ABCMeta):
"""Metaclass for MultiCollector that dispatches based on sync parameter.
When MultiCollector is instantiated with sync=True or sync=False, the metaclass
intercepts the call and returns the appropriate subclass instance:
- sync=True: returns MultiSyncCollector (alias: MultiSyncCollector)
- sync=False: returns MultiAsyncCollector (alias: MultiAsyncCollector)
"""
def __call__(cls, *args, sync: bool | None = None, **kwargs):
# Only dispatch if we're instantiating MultiCollector directly (not a subclass)
# and sync is explicitly provided
if cls.__name__ == "MultiCollector" and sync is not None:
if sync:
from torchrl.collectors._multi_sync import MultiSyncCollector
return MultiSyncCollector(*args, **kwargs)
else:
from torchrl.collectors._multi_async import MultiAsyncCollector
return MultiAsyncCollector(*args, **kwargs)
return super().__call__(*args, **kwargs)
[docs]
class MultiCollector(BaseCollector, metaclass=_MultiCollectorMeta):
"""Runs a given number of DataCollectors on separate processes.
.. note::
Use :class:`~torchrl.collectors.Collector` with ``num_collectors`` and
``sync`` to construct local process collectors in new code. This class
remains the concrete process-collector API.
Args:
create_env_fn (List[Callabled]): list of Callables, each returning an
instance of :class:`~torchrl.envs.EnvBase`.
policy (Callable): Policy to be executed in the environment.
Must accept :class:`tensordict.tensordict.TensorDictBase` object as input.
If ``None`` is provided (default), the policy used will be a
:class:`~torchrl.collectors.RandomPolicy` instance with the environment
``action_spec``.
Accepted policies are usually subclasses of :class:`~tensordict.nn.TensorDictModuleBase`.
This is the recommended usage of the collector.
Other callables are accepted too:
If the policy is not a ``TensorDictModuleBase`` (e.g., a regular :class:`~torch.nn.Module`
instances) it will be wrapped in a `nn.Module` first.
Then, the collector will try to assess if these
modules require wrapping in a :class:`~tensordict.nn.TensorDictModule` or not.
- If the policy forward signature matches any of ``forward(self, tensordict)``,
``forward(self, td)`` or ``forward(self, <anything>: TensorDictBase)`` (or
any typing with a single argument typed as a subclass of ``TensorDictBase``)
then the policy won't be wrapped in a :class:`~tensordict.nn.TensorDictModule`.
- In all other cases an attempt to wrap it will be undergone as such:
``TensorDictModule(policy, in_keys=env_obs_key, out_keys=env.action_keys)``.
.. note:: If the policy needs to be passed as a policy factory (e.g., in case it mustn't be serialized /
pickled directly), the ``policy_factory`` should be used instead.
.. note:: When using ``weight_sync_schemes``, both ``policy`` and ``policy_factory`` can be provided together.
In this case, the ``policy`` is used ONLY for weight extraction (via ``TensorDict.from_module()``) to
set up weight synchronization, but it is NOT sent to workers and its weights are NOT depopulated.
The ``policy_factory`` is what actually gets passed to workers to create their local policy instances.
This is useful when the policy is hard to serialize but you have a copy on the main node for
weight synchronization purposes.
Keyword Args:
sync (bool, optional): if ``True``, the collector will run in sync mode (:class:`~torchrl.collectors.MultiSyncCollector`). If
`False`, the collector will run in async mode (:class:`~torchrl.collectors.MultiAsyncCollector`).
policy_factory (Callable[[], Callable], list of Callable[[], Callable], optional): a callable
(or list of callables) that returns a policy instance.
When not using ``weight_sync_schemes``, this is mutually exclusive with the ``policy`` argument.
When using ``weight_sync_schemes``, both ``policy`` and ``policy_factory`` can be provided:
the ``policy`` is used for weight extraction only, while ``policy_factory`` creates policies on workers.
.. note:: `policy_factory` comes in handy whenever the policy cannot be serialized.
.. warning:: `policy_factory` is currently not compatible with multiprocessed data
collectors.
num_workers (int, optional): number of workers to use. If `create_env_fn` is a list, this will be ignored.
Defaults to `None` (workers determined by the `create_env_fn` length).
frames_per_batch (int, Sequence[int]): A keyword-only argument representing the
total number of elements in a batch. If a sequence is provided, represents the number of elements in a
batch per worker. Total number of elements in a batch is then the sum over the sequence.
total_frames (int, optional): A keyword-only argument representing the
total number of frames returned by the collector
during its lifespan. If the ``total_frames`` is not divisible by
``frames_per_batch``, an exception is raised.
Endless collectors can be created by passing ``total_frames=-1``.
Defaults to ``-1`` (never ending collector).
device (int, str or torch.device, optional): The generic device of the
collector. The ``device`` args fills any non-specified device: if
``device`` is not ``None`` and any of ``storing_device``, ``policy_device`` or
``env_device`` is not specified, its value will be set to ``device``.
Defaults to ``None`` (No default device).
Supports a list of devices if one wishes to indicate a different device
for each worker. The list must be as long as the number of workers.
storing_device (int, str or torch.device, optional): The device on which
the output :class:`~tensordict.TensorDict` will be stored.
If ``device`` is passed and ``storing_device`` is ``None``, it will
default to the value indicated by ``device``.
For long trajectories, it may be necessary to store the data on a different
device than the one where the policy and env are executed.
Defaults to ``None`` (the output tensordict isn't on a specific device,
leaf tensors sit on the device where they were created).
Supports a list of devices if one wishes to indicate a different device
for each worker. The list must be as long as the number of workers.
env_device (int, str or torch.device, optional): The device on which
the environment should be cast (or executed if that functionality is
supported). If not specified and the env has a non-``None`` device,
``env_device`` will default to that value. If ``device`` is passed
and ``env_device=None``, it will default to ``device``. If the value
as such specified of ``env_device`` differs from ``policy_device``
and one of them is not ``None``, the data will be cast to ``env_device``
before being passed to the env (i.e., passing different devices to
policy and env is supported). Defaults to ``None``.
Supports a list of devices if one wishes to indicate a different device
for each worker. The list must be as long as the number of workers.
policy_device (int, str or torch.device, optional): The device on which
the policy should be cast.
If ``device`` is passed and ``policy_device=None``, it will default
to ``device``. If the value as such specified of ``policy_device``
differs from ``env_device`` and one of them is not ``None``,
the data will be cast to ``policy_device`` before being passed to
the policy (i.e., passing different devices to policy and env is
supported). Defaults to ``None``.
Supports a list of devices if one wishes to indicate a different device
for each worker. The list must be as long as the number of workers.
create_env_kwargs (dict, optional): A dictionary with the
keyword arguments used to create an environment. If a list is
provided, each of its elements will be assigned to a sub-collector.
collector_class (Python class or constructor): a collector class to be remotely instantiated. Can be
:class:`~torchrl.collectors.Collector`,
:class:`~torchrl.collectors.MultiSyncCollector`,
:class:`~torchrl.collectors.MultiAsyncCollector`
or a derived class of these.
Defaults to :class:`~torchrl.collectors.Collector`.
max_frames_per_traj (int, optional): Maximum steps per trajectory.
Note that a trajectory can span across multiple batches (unless
``reset_at_each_iter`` is set to ``True``, see below).
Once a trajectory reaches ``n_steps``, the environment is reset.
If the environment wraps multiple environments together, the number
of steps is tracked for each environment independently. Negative
values are allowed, in which case this argument is ignored.
Defaults to ``None`` (i.e. no maximum number of steps).
init_random_frames (int, optional): Number of frames for which the
policy is ignored before it is called. This feature is mainly
intended to be used in offline/model-based settings, where a
batch of random trajectories can be used to initialize training.
If provided, it will be rounded up to the closest multiple of frames_per_batch.
Defaults to ``None`` (i.e. no random frames).
reset_at_each_iter (bool, optional): Whether environments should be reset
at the beginning of a batch collection.
Defaults to ``False``.
postproc (Callable, optional): A post-processing transform, such as
a :class:`~torchrl.envs.Transform` or a :class:`~torchrl.data.postprocs.MultiStep`
instance.
Defaults to ``None``.
split_trajs (bool, optional): Boolean indicating whether the resulting
TensorDict should be split according to the trajectories.
See :func:`~torchrl.collectors.utils.split_trajectories` for more
information.
Defaults to ``False``.
exploration_type (ExplorationType, optional): interaction mode to be used when
collecting data. Must be one of ``torchrl.envs.utils.ExplorationType.DETERMINISTIC``,
``torchrl.envs.utils.ExplorationType.RANDOM``, ``torchrl.envs.utils.ExplorationType.MODE``
or ``torchrl.envs.utils.ExplorationType.MEAN``.
reset_when_done (bool, optional): if ``True`` (default), an environment
that return a ``True`` value in its ``"done"`` or ``"truncated"``
entry will be reset at the corresponding indices.
update_at_each_batch (boolm optional): if ``True``, :meth:`update_policy_weights_()`
will be called before (sync) or after (async) each data collection.
Defaults to ``False``.
preemptive_threshold (:obj:`float`, optional): a value between 0.0 and 1.0 that specifies the ratio of workers
that will be allowed to finish collecting their rollout before the rest are forced to end early.
Frames that were not collected by the preempted workers are marked invalid: their
``("collector", "traj_ids")`` entry is ``-1`` and, with ``cat_results="stack"``, a
``("collector", "mask")`` entry flags the valid frames (with ``cat_results=-1`` the
invalid frames are dropped from the batch instead).
num_threads (int, optional): number of threads for this process.
Defaults to the number of workers.
num_sub_threads (int, optional): number of threads of the subprocesses.
Should be equal to one plus the number of processes launched within
each subprocess (or one if a single process is launched).
Defaults to 1 for safety: if none is indicated, launching multiple
workers may charge the cpu load too much and harm performance.
cat_results (str, int or None): (:class:`~torchrl.collectors.MultiSyncCollector` exclusively).
If ``"stack"``, the data collected from the workers will be stacked along the
first dimension. This is the preferred behavior as it is the most compatible
with the rest of the library.
If ``0``, results will be concatenated along the first dimension
of the outputs, which can be the batched dimension if the environments are
batched or the time dimension if not.
A ``cat_results`` value of ``-1`` will always concatenate results along the
time dimension. This should be preferred over the default. Intermediate values
are also accepted.
Defaults to ``"stack"``.
.. note:: From v0.5, this argument will default to ``"stack"`` for a better
interoperability with the rest of the library.
set_truncated (bool, optional): if ``True``, the truncated signals (and corresponding
``"done"`` but not ``"terminated"``) will be set to ``True`` when the last frame of
a rollout is reached. If no ``"truncated"`` key is found, an exception is raised.
Truncated keys can be set through ``env.add_truncated_keys``.
Defaults to ``False``.
With a multi-process collector writing to a shared replay buffer,
this marks the worker-batch seams that would otherwise be
invisible to a :class:`~torchrl.data.replay_buffers.SliceSampler`;
see :ref:`the trajectory-boundary documentation <ref_traj_boundaries>`
and :ref:`collectors_replay_trajs` for the trade-offs.
trajs_per_batch (int, optional): When set without ``replay_buffer``,
the multi-collector
assembles trajectories from the worker batches and yields
zero-padded batches of shape ``(trajs_per_batch, max_traj_len)``
with a ``("collector", "mask")`` boolean field, or flat unpadded
concatenations with ``traj_format="cat"``.
Defaults to ``None`` (fixed-frame batches).
replay_write_mode (``"rollout"``, ``"trajectory"``, optional): Selects
fixed-frame rollout writes or complete-trajectory writes to the
shared replay buffer. In trajectory mode, assembly is delegated
to each worker and only flat, completed trajectories are inserted.
Defaults to ``None``; the legacy combination of ``replay_buffer``
and ``trajs_per_batch`` still selects trajectory writes.
traj_format (str, optional): layout of the batches yielded when
``trajs_per_batch`` is set without a ``replay_buffer``:
``"padded"`` for zero-padded
``(trajs_per_batch, max_traj_len)`` stacks with a
``("collector", "mask")`` entry, ``"cat"`` for flat unpadded
concatenations along time (trajectories delimited by
``("next", "done")`` and ``("collector", "traj_ids")``).
Replay-buffer writes are 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.
use_buffers (bool, optional): if ``True``, a buffer will be used to stack the data.
This isn't compatible with environments with dynamic specs. Defaults to ``True``
for envs without dynamic specs, ``False`` for others.
replay_buffer (ReplayBuffer, optional): if provided, the collector will not yield tensordicts
but populate the buffer instead. Defaults to ``None``.
extend_buffer (bool, optional): if `True`, the replay buffer is extended with entire rollouts and not
with single steps. Defaults to `True` for multiprocessed data collectors.
trust_policy (bool, optional): if ``True``, a non-TensorDictModule policy will be trusted to be
assumed to be compatible with the collector. This defaults to ``True`` for CudaGraphModules
and ``False`` otherwise.
compile_policy (bool or Dict[str, Any], optional): if ``True``, the policy will be compiled
using :func:`~torch.compile` default behaviour. If a dictionary of kwargs is passed, it
will be used to compile the policy.
cudagraph_policy (bool or Dict[str, Any], optional): if ``True``, the policy will be wrapped
in :class:`~tensordict.nn.CudaGraphModule` with default kwargs.
If a dictionary of kwargs is passed, it will be used to wrap the policy.
no_cuda_sync (bool): if ``True``, explicit CUDA synchronizations calls will be bypassed.
For environments running directly on CUDA (`IsaacLab <https://github.com/isaac-sim/IsaacLab/>`_
or `ManiSkills <https://github.com/haosulab/ManiSkill/>`_) cuda synchronization may cause unexpected
crashes.
Defaults to ``False``.
auto_register_policy_transforms (bool, optional): forwarded to each
worker :class:`~torchrl.collectors.Collector`. When ``True``,
workers append :class:`~torchrl.envs.transforms.InitTracker` and
recurrent-state :class:`~torchrl.envs.transforms.TensorDictPrimer`
transforms to their envs if the env specs don't already provide
them. ``False`` disables it; ``None`` (default through v0.14)
preserves pre-v0.15 behavior and emits a
:class:`FutureWarning` when wrapping would have been needed.
Default flips to ``True`` in v0.15.
weight_updater (WeightUpdaterBase or constructor, optional): An instance of :class:`~torchrl.collectors.WeightUpdaterBase`
or its subclass, responsible for updating the policy weights on remote inference workers.
If not provided, a :class:`~torchrl.collectors.MultiProcessedWeightUpdater` will be used by default,
which handles weight synchronization across multiple processes.
Consider using a constructor if the updater needs to be serialized.
weight_sync_schemes (dict[str, WeightSyncScheme], optional): Dictionary of weight sync schemes for
SENDING weights to worker sub-collectors. Keys are model identifiers (e.g., "policy")
and values are WeightSyncScheme instances configured to send weights to child processes.
If not provided, a :class:`~torchrl.collectors.MultiProcessWeightSyncScheme` will be used by default.
This is for propagating weights DOWN the hierarchy (parent -> children).
weight_recv_schemes (dict[str, WeightSyncScheme], optional): Dictionary of weight sync schemes for
RECEIVING weights from parent collectors. Keys are model identifiers (e.g., "policy")
and values are WeightSyncScheme instances configured to receive weights.
This enables cascading in hierarchies like: RPCCollector -> MultiSyncCollector -> Collector.
Received weights are automatically propagated to sub-collectors if matching model_ids exist.
Defaults to ``None``.
track_policy_version (bool or PolicyVersion, optional): if ``True``, the collector will track the version of the policy.
A :class:`~torchrl.envs.transforms.PolicyVersion` transform is
installed on each worker's environment, tagging every collected frame with the
current version under the ``"policy_version"`` key. Each worker's transform is
bumped after the new weights have actually been applied in that worker, so
per-frame tagging tracks real weight updates rather than rollout iterations.
Note that in asynchronous mode a batch that was already in flight when
:meth:`update_policy_weights_` is called may straddle the bump (some frames
tagged with the old version, the remainder with the new). Treat the value as
the version under which each individual frame was produced, not as a batch-level
label.
For multi-process collectors, the ``"policy_version"`` entries in the
collected tensordict are produced by worker-local transforms and are the
source of truth for data provenance. The parent collector's
:attr:`policy_version` property is only available while all workers are
known to have acknowledged the same sequence of weight updates. Use
:meth:`worker_policy_versions` for an explicit per-worker snapshot.
The recommended path is ``track_policy_version=True``: let the collector own
the transform. Passing a :class:`~torchrl.envs.transforms.PolicyVersion`
instance directly is reserved for advanced use cases that wire up a
``PolicyVersion`` **without** going through a collector. With multi-process
collectors that pre-built tracker lives in the *parent* and is not propagated
into workers, so per-frame tagging will still be driven by per-worker
transforms — favor ``True``.
Defaults to ``False``.
compact_obs (bool, optional): if ``True``, each worker drops the
observation and state keys from the ``("next", ...)`` sub-tensordict
before stacking. See
:class:`~torchrl.collectors.Collector` for the full
explanation and tradeoffs (most notably:
:class:`~torchrl.envs.transforms.MultiStepTransform` cannot be used
in compact mode), plus the pairing with
:class:`~torchrl.envs.transforms.NextStateReconstructor` at
sampling time, the boundary-preserving lossy alternative
:class:`~torchrl.envs.transforms.NextObservationDelta`, and the
*Memory-efficient RL training* tutorial. Defaults to ``False``.
worker_idx (int, optional): the index of the worker.
Examples:
>>> from torchrl.collectors import MultiCollector
>>> from torchrl.envs import GymEnv
>>>
>>> def make_env():
... return GymEnv("CartPole-v1")
>>>
>>> # Synchronous collection (for on-policy algorithms like PPO)
>>> sync_collector = MultiCollector(
... create_env_fn=[make_env] * 4, # 4 parallel workers
... policy=my_policy,
... frames_per_batch=1000,
... total_frames=100000,
... sync=True, # All workers complete before batch is delivered
... )
>>>
>>> # Asynchronous collection (for off-policy algorithms like SAC)
>>> async_collector = MultiCollector(
... create_env_fn=[make_env] * 4,
... policy=my_policy,
... frames_per_batch=1000,
... total_frames=100000,
... sync=False, # First-come-first-serve delivery
... )
>>>
>>> # Iterate over collected data
>>> for data in sync_collector:
... # data is a TensorDict with collected transitions
... pass
>>> sync_collector.shutdown()
"""
def __init__(
self,
create_env_fn: Sequence[Callable[[], EnvBase]],
policy: None
| (TensorDictModule | Callable[[TensorDictBase], TensorDictBase]) = None,
*,
num_workers: int | None = None,
policy_factory: (
Callable[[], Callable] | list[Callable[[], Callable]] | None
) = None,
frames_per_batch: int | Sequence[int],
total_frames: int | None = -1,
device: DEVICE_TYPING | Sequence[DEVICE_TYPING] | None = None,
storing_device: DEVICE_TYPING | Sequence[DEVICE_TYPING] | None = None,
env_device: DEVICE_TYPING | Sequence[DEVICE_TYPING] | None = None,
policy_device: DEVICE_TYPING | Sequence[DEVICE_TYPING] | None = None,
create_env_kwargs: Sequence[dict] | None = None,
collector_class: type | Callable[[], BaseCollector] | None = None,
max_frames_per_traj: int | None = None,
init_random_frames: int | None = None,
reset_at_each_iter: bool = False,
postproc: Callable[[TensorDictBase], TensorDictBase] | None = None,
split_trajs: bool | None = None,
exploration_type: ExplorationType = DEFAULT_EXPLORATION_TYPE,
reset_when_done: bool = True,
update_at_each_batch: bool = False,
preemptive_threshold: float | None = None,
num_threads: int | None = None,
num_sub_threads: int = 1,
cat_results: str | int | None = None,
set_truncated: bool = False,
use_buffers: bool | None = None,
replay_buffer: ReplayBuffer | None = None,
extend_buffer: bool = True,
trust_policy: bool | None = None,
compile_policy: bool | dict[str, Any] | None = None,
cudagraph_policy: bool | dict[str, Any] | None = None,
no_cuda_sync: bool = False,
weight_updater: (
WeightUpdaterBase | Callable[[], WeightUpdaterBase] | None
) = None,
weight_sync_schemes: dict[str, WeightSyncScheme] | None = None,
weight_recv_schemes: dict[str, WeightSyncScheme] | None = None,
track_policy_version: bool = False,
worker_idx: int | 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"] | None = None,
init_fn: Callable[[], None] | None = None,
auto_register_policy_transforms: bool | None = None,
pre_collect_hook: Callable[[], None] | None = None,
post_collect_hook: Callable[[TensorDictBase], None] | None = None,
compact_obs: bool = False,
):
self.closed = True
self._shutdown_complete = False
self.worker_idx = worker_idx
self.trajs_per_batch = trajs_per_batch
self.trajs_per_write = trajs_per_write
self.replay_write_mode = _validate_replay_write_mode(
replay_write_mode,
has_replay_buffer=replay_buffer is not None,
trajs_per_batch=trajs_per_batch,
trajs_per_write=trajs_per_write,
)
# Preserve the raw selector for workers. In the legacy
# replay_buffer + trajs_per_batch form it must stay None so that the
# worker can resolve that combination without seeing two selectors.
self._worker_replay_write_mode = replay_write_mode
self._trajectory_writes_in_workers = self.replay_write_mode == "trajectory"
self.traj_format = _validate_traj_format(
traj_format, trajs_per_batch, has_replay_buffer=replay_buffer is not None
)
self._auto_register_policy_transforms = auto_register_policy_transforms
super().__init__(
pre_collect_hook=pre_collect_hook,
post_collect_hook=post_collect_hook,
)
self._worker_pre_collect_hook = (
CloudpickleWrapper(pre_collect_hook)
if pre_collect_hook is not None
else None
)
self._worker_post_collect_hook = (
CloudpickleWrapper(post_collect_hook)
if post_collect_hook is not None
else None
)
self._worker_trajs_per_batch = None
# Wrap init_fn with CloudpickleWrapper to support lambdas / closures
# across the spawn start method.
self._worker_init_fn = (
CloudpickleWrapper(init_fn) if init_fn is not None else None
)
# Set up workers and environment functions
create_env_fn, total_frames_per_batch = self._setup_workers_and_env_fns(
create_env_fn, num_workers, frames_per_batch
)
# Set up basic configuration
self.set_truncated = set_truncated
self.num_sub_threads = num_sub_threads
self.num_threads = num_threads
self.create_env_fn = create_env_fn
self._read_compile_kwargs(compile_policy, cudagraph_policy)
# Set up environment kwargs
self.create_env_kwargs = self._setup_env_kwargs(create_env_kwargs)
# Set up devices
storing_devices, policy_devices, env_devices = self._get_devices(
storing_device=storing_device,
env_device=env_device,
policy_device=policy_device,
device=device,
)
self.storing_device = storing_devices
self.policy_device = policy_devices
self.env_device = env_devices
self.collector_class = collector_class
del storing_device, env_device, policy_device, device
self.no_cuda_sync = no_cuda_sync
# Set up replay buffer
self._use_buffers = use_buffers
self.replay_buffer = replay_buffer
# Set up policy and weights
if trust_policy is None:
trust_policy = policy is not None and isinstance(policy, CudaGraphModule)
self.trust_policy = trust_policy
policy_factory = self._setup_policy_factory(policy_factory)
# Set up weight synchronization
if weight_sync_schemes is None and weight_updater is None:
weight_sync_schemes = {}
elif weight_sync_schemes is not None and weight_updater is not None:
raise TypeError(
"Cannot specify both weight_sync_schemes and weight_updater."
)
if (
weight_sync_schemes is not None
and not weight_sync_schemes
and weight_updater is None
and isinstance(policy, nn.Module)
):
weight_sync_schemes["policy"] = SharedMemWeightSyncScheme()
self._setup_multi_weight_sync(weight_updater, weight_sync_schemes)
# Store policy and policy_factory - temporary set to make them visible to the receiver
self.policy = policy
self.policy_factory = policy_factory
# Policy-version tracking must be configured before replay-buffer
# initialization so the shared storage reserves worker-produced fields.
self._setup_multi_policy_version_tracking(track_policy_version)
self._setup_multi_replay_buffer(replay_buffer, extend_buffer)
# Set up weight receivers if provided
if weight_recv_schemes is not None:
self.register_scheme_receiver(weight_recv_schemes)
self._setup_multi_policy_and_weights(
self.policy, self.policy_factory, weight_updater, weight_sync_schemes
)
# # Set up fallback policy for weight extraction
# self._setup_fallback_policy(policy, policy_factory, weight_sync_schemes)
# Set up total frames and other parameters
self._setup_multi_total_frames(
total_frames, total_frames_per_batch, frames_per_batch
)
self.reset_at_each_iter = reset_at_each_iter
self.postproc = postproc
self.compact_obs = bool(compact_obs)
self.max_frames_per_traj = (
int(max_frames_per_traj) if max_frames_per_traj is not None else 0
)
# Set up split trajectories
self.requested_frames_per_batch = total_frames_per_batch
self.reset_when_done = reset_when_done
self._setup_split_trajs(split_trajs, reset_when_done)
# Set up other parameters
self.init_random_frames = (
int(init_random_frames) if init_random_frames is not None else 0
)
self.update_at_each_batch = update_at_each_batch
self.exploration_type = exploration_type
self.frames_per_worker = np.inf
# Set up preemptive threshold
self._setup_preemptive_threshold(preemptive_threshold)
# Run worker processes
try:
self._run_processes()
except Exception as e:
self.shutdown(raise_on_error=False)
raise e
# Set up frame tracking and other options
self._exclude_private_keys = True
self._frames = 0
self._iter = -1
# Validate cat_results
self._validate_cat_results(cat_results)
def _setup_workers_and_env_fns(
self,
create_env_fn: Sequence[Callable] | Callable,
num_workers: int | None,
frames_per_batch: int | Sequence[int],
) -> tuple[list[Callable], int]:
"""Set up workers and environment functions."""
if isinstance(create_env_fn, Sequence):
self.num_workers = len(create_env_fn)
else:
self.num_workers = num_workers
create_env_fn = [create_env_fn] * self.num_workers
if (
isinstance(frames_per_batch, Sequence)
and len(frames_per_batch) != self.num_workers
):
raise ValueError(
"If `frames_per_batch` is provided as a sequence, it should contain exactly one value per worker."
f"Got {len(frames_per_batch)} values for {self.num_workers} workers."
)
self._frames_per_batch = frames_per_batch
total_frames_per_batch = (
sum(frames_per_batch)
if isinstance(frames_per_batch, Sequence)
else frames_per_batch
)
return create_env_fn, total_frames_per_batch
def _setup_env_kwargs(
self, create_env_kwargs: Sequence[dict] | dict | None
) -> list[dict]:
"""Set up environment kwargs for each worker."""
if isinstance(create_env_kwargs, Mapping):
create_env_kwargs = [create_env_kwargs] * self.num_workers
elif create_env_kwargs is None:
create_env_kwargs = [{}] * self.num_workers
elif isinstance(create_env_kwargs, (tuple, list)):
create_env_kwargs = list(create_env_kwargs)
if len(create_env_kwargs) != self.num_workers:
raise ValueError(
f"len(create_env_kwargs) must be equal to num_workers, got {len(create_env_kwargs)=} and {self.num_workers=}"
)
return create_env_kwargs
def _setup_multi_replay_buffer(
self,
replay_buffer: ReplayBuffer | None,
extend_buffer: bool,
) -> None:
"""Set up replay buffer for multi-process collector."""
self.local_init_rb = True
self._check_replay_buffer_init()
self.extend_buffer = extend_buffer
if (
replay_buffer is not None
and hasattr(replay_buffer, "shared")
and not replay_buffer.shared
):
torchrl_logger.info("Replay buffer is not shared. Sharing it.")
replay_buffer.share()
def _setup_policy_factory(
self, policy_factory: Callable | list[Callable] | None
) -> list[Callable | None]:
"""Set up policy factory for each worker."""
if not isinstance(policy_factory, Sequence):
policy_factory = [policy_factory] * self.num_workers
return policy_factory
def _setup_multi_policy_and_weights(
self,
policy: TensorDictModule | Callable | None,
policy_factory: list[Callable | None],
weight_updater: WeightUpdaterBase | Callable | None,
weight_sync_schemes: dict[str, WeightSyncScheme] | None,
) -> None:
"""Set up policy for multi-process collector.
With weight sync schemes: validates and stores policy without weight extraction.
With weight updater: extracts weights and creates stateful policies.
When both policy and policy_factory are provided (with weight_sync_schemes):
- The policy is used ONLY for weight extraction via get_model()
- The policy is NOT depopulated of its weights
- The policy is NOT sent to workers
- The policy_factory is used to create policies on workers
"""
if any(policy_factory) and policy is not None:
if weight_sync_schemes is None:
raise TypeError(
"policy_factory and policy are mutually exclusive when not using weight_sync_schemes. "
"When using weight_sync_schemes, policy can be provided alongside policy_factory "
"for weight extraction purposes only (the policy will not be sent to workers)."
)
# Store policy as fallback for weight extraction only
# The policy keeps its weights and is NOT sent to workers
self._fallback_policy = policy
if weight_sync_schemes is not None:
weight_sync_policy = weight_sync_schemes.get("policy")
if weight_sync_policy is None:
return
# # If we only have a policy_factory (no policy instance), the scheme must
# # be pre-initialized on the sender, since there is no policy on the
# # collector to extract weights from.
# if any(p is not None for p in policy_factory) and policy is None:
# if not weight_sync_policy.initialized_on_sender:
# raise RuntimeError(
# "the weight sync scheme must be initialized on sender ahead of time "
# "when passing a policy_factory without a policy instance on the collector. "
# f"Got {policy_factory=}"
# )
# # When a policy instance is provided alongside a policy_factory, the scheme
# # can rely on the collector context (and its policy) to extract weights.
# # Weight sync scheme initialization then happens in _run_processes where
# # pipes and workers are available.
else:
# Using legacy weight updater - extract weights and create stateful policies
self._setup_multi_policy_and_weights_legacy(
policy, policy_factory, weight_updater, weight_sync_schemes
)
def _setup_multi_policy_and_weights_legacy(
self,
policy: TensorDictModule | Callable | None,
policy_factory: list[Callable | None],
weight_updater: WeightUpdaterBase | Callable | None,
weight_sync_schemes: dict[str, WeightSyncScheme] | None,
) -> None:
"""Set up policy and extract weights for each device.
Creates stateful policies with weights extracted and placed in shared memory.
Used with weight updater for in-place weight replacement.
"""
self._policy_weights_dict = {}
self._fallback_policy = None # Policy to use for weight extraction fallback
if not any(policy_factory):
for policy_device, env_maker, env_maker_kwargs in _zip_strict(
self.policy_device, self.create_env_fn, self.create_env_kwargs
):
policy_new_device, get_weights_fn = self._get_policy_and_device(
policy=policy,
policy_device=policy_device,
env_maker=env_maker,
env_maker_kwargs=env_maker_kwargs,
)
if type(policy_new_device) is not type(policy):
policy = policy_new_device
weights = (
TensorDict.from_module(policy_new_device)
if isinstance(policy_new_device, nn.Module)
else TensorDict()
)
# For multi-process collectors, ensure weights are in shared memory
if policy_device and policy_device.type == "cpu":
weights = weights.share_memory_()
self._policy_weights_dict[policy_device] = weights
# Store the first policy instance for fallback weight extraction
if self._fallback_policy is None:
self._fallback_policy = policy_new_device
self._get_weights_fn = get_weights_fn
if weight_updater is None:
# For multiprocessed collectors, use MultiProcessWeightSyncScheme by default
if weight_sync_schemes is None:
weight_sync_schemes = {"policy": MultiProcessWeightSyncScheme()}
self._weight_sync_schemes = weight_sync_schemes
elif weight_updater is None:
warnings.warn(
"weight_updater is None, but policy_factory is provided. This means that the server will "
"not know how to send the weights to the workers. If the workers can handle their weight synchronization "
"on their own (via some specialized worker type / constructor) this may well work, but make sure "
"your weight synchronization strategy is properly set. To suppress this warning, you can use "
"RemoteModuleWeightUpdater() which enforces explicit weight passing when calling update_policy_weights_(weights). "
"This will work whenever your inference and training policies are nn.Module instances with similar structures."
)
def _setup_multi_weight_sync(
self,
weight_updater: WeightUpdaterBase | Callable | None,
weight_sync_schemes: dict[str, WeightSyncScheme] | None,
) -> None:
"""Set up weight synchronization for multi-process collector."""
if weight_sync_schemes is not None:
# Use weight sync schemes for weight distribution
self._weight_sync_schemes = weight_sync_schemes
# Senders will be created in _run_processes
self.weight_updater = None
else:
# Use weight updater for weight distribution
self.weight_updater = weight_updater
self._weight_sync_schemes = None
def _setup_multi_policy_version_tracking(
self, track_policy_version: bool | PolicyVersion
) -> None:
"""Set up policy version tracking for multi-process collector."""
self.policy_version_tracker = track_policy_version
if PolicyVersion is not None:
if isinstance(track_policy_version, bool) and track_policy_version:
self.policy_version_tracker = PolicyVersion()
elif hasattr(track_policy_version, "increment_version"):
self.policy_version_tracker = track_policy_version
else:
self.policy_version_tracker = None
else:
if track_policy_version:
raise ImportError(
"PolicyVersion is not available. Please install the LLM dependencies or set track_policy_version=False."
)
self.policy_version_tracker = None
self._acknowledged_policy_version = (
self.policy_version_tracker.version
if self.policy_version_tracker is not None
else None
)
[docs]
def update_policy_weights_(
self,
policy_or_weights: (
TensorDictBase | TensorDictModuleBase | nn.Module | dict | None
) = None,
*,
worker_ids: int | list[int] | torch.device | list[torch.device] | None = None,
**kwargs,
) -> None:
"""Update worker policy weights and track an acknowledged version.
The aggregate :attr:`policy_version` is advanced only when every worker
synchronously acknowledges the update. Partial or asynchronous updates
invalidate the aggregate because a single version can no longer describe
every worker.
"""
acknowledged_version = self._acknowledged_policy_version
if self.policy_version_tracker is not None:
# If synchronization fails after some workers have applied the update,
# retaining the previous aggregate would be misleading.
self._acknowledged_policy_version = None
super().update_policy_weights_(
policy_or_weights=policy_or_weights, worker_ids=worker_ids, **kwargs
)
if self.policy_version_tracker is None or acknowledged_version is None:
return
if worker_ids is None:
all_workers_targeted = True
elif isinstance(worker_ids, int):
all_workers_targeted = self.num_workers == 1 and worker_ids == 0
elif isinstance(worker_ids, torch.device):
all_workers_targeted = False
elif all(isinstance(worker_id, int) for worker_id in worker_ids):
all_workers_targeted = len(worker_ids) == self.num_workers and set(
worker_ids
) == set(range(self.num_workers))
else:
all_workers_targeted = False
weights = kwargs.get("weights", policy_or_weights)
per_worker_weights = (
isinstance(weights, dict)
and bool(weights)
and all(isinstance(worker_id, int) for worker_id in weights)
)
if per_worker_weights:
all_workers_targeted = all_workers_targeted and set(weights) == set(
range(self.num_workers)
)
weights_dict = kwargs.get("weights_dict")
if weights_dict is not None:
model_ids = list(weights_dict)
for model_weights in weights_dict.values():
if (
isinstance(model_weights, dict)
and model_weights
and all(isinstance(worker_id, int) for worker_id in model_weights)
):
all_workers_targeted = all_workers_targeted and set(
model_weights
) == set(range(self.num_workers))
else:
model_ids = [kwargs.get("model_id") or "policy"]
if not model_ids:
self._acknowledged_policy_version = acknowledged_version
return
schemes = self._weight_sync_schemes
synchronously_acknowledged = bool(schemes) and all(
model_id in schemes and getattr(schemes[model_id], "sync", True)
for model_id in model_ids
)
if all_workers_targeted and synchronously_acknowledged:
# BaseCollector synchronizes each model separately, and the worker
# cascade increments its local tracker once per synced model.
# Rebuild the aggregate with that same per-model count only after
# every worker has acknowledged the full synchronous update.
self.policy_version_tracker.version = acknowledged_version
for _ in model_ids:
self.increment_version()
# TODO: Remove this
def _setup_fallback_policy(
self,
policy: TensorDictModule | Callable | None,
policy_factory: list[Callable | None],
weight_sync_schemes: dict[str, WeightSyncScheme] | None,
) -> None:
"""Set up fallback policy for weight extraction when using policy_factory."""
# _fallback_policy is already set in _setup_multi_policy_and_weights if a policy was provided
# If policy_factory was used, create a policy instance to use as fallback
if policy is None and any(policy_factory) and weight_sync_schemes is not None:
if not hasattr(self, "_fallback_policy") or self._fallback_policy is None:
first_factory = (
policy_factory[0]
if isinstance(policy_factory, list)
else policy_factory
)
if first_factory is not None:
# Create a policy instance for weight extraction
# This will be a reference to a policy with the same structure
# For shared memory, modifications to any policy will be visible here
self._fallback_policy = first_factory()
def _setup_multi_total_frames(
self,
total_frames: int,
total_frames_per_batch: int,
frames_per_batch: int | Sequence[int],
) -> None:
"""Validate and set total frames for multi-process collector."""
if total_frames is None or total_frames < 0:
total_frames = float("inf")
else:
remainder = total_frames % total_frames_per_batch
if remainder != 0 and RL_WARNINGS:
warnings.warn(
f"total_frames ({total_frames}) is not exactly divisible by frames_per_batch ({total_frames_per_batch}). "
f"This means {total_frames_per_batch - remainder} additional frames will be collected. "
"To silence this message, set the environment variable RL_WARNINGS to False."
)
self.total_frames = (
int(total_frames) if total_frames != float("inf") else total_frames
)
def _setup_split_trajs(
self, split_trajs: bool | None, reset_when_done: bool
) -> None:
"""Set up split trajectories option."""
if split_trajs is None:
split_trajs = False
elif not reset_when_done and split_trajs:
raise RuntimeError(
"Cannot split trajectories when reset_when_done is False."
)
elif split_trajs:
warnings.warn(
"split_trajs=True produces a (N_traj, T_max) zero-padded "
"tensordict with a 'mask' key. For sequence training, prefer "
"the contiguous-trajectory layout: pass a replay_buffer to "
"the collector and sample with "
":class:`~torchrl.data.SliceSampler` (variable-length slices, "
"no padding, no mask). See "
":ref:`Data layout: contiguous trajectories <data-layout>` "
"in the docs. This advisory will become a "
"DeprecationWarning in a future release.",
stacklevel=3,
)
self.split_trajs = split_trajs
def _setup_preemptive_threshold(self, preemptive_threshold: float | None) -> None:
"""Set up preemptive threshold for early stopping."""
if preemptive_threshold is not None:
preemptive_threshold = float(preemptive_threshold)
if not 0.0 <= preemptive_threshold <= 1.0:
raise ValueError(
f"preemptive_threshold must be between 0.0 and 1.0, got {preemptive_threshold}."
)
self.preemptive_threshold = preemptive_threshold
# A threshold of 1.0 waits for every worker, so preemption can never
# fire: skip the interruptor and its per-step polling in the workers.
self.interruptor = _Interruptor() if preemptive_threshold < 1.0 else None
else:
self.preemptive_threshold = 1.0
self.interruptor = None
def _should_use_random_frames(self) -> bool:
"""Determine if random frames should be used instead of the policy.
When a replay buffer is provided, uses `replay_buffer.write_count` as the
global step counter to support `.start()` mode where `_frames` isn't updated
until after collection. Otherwise, uses the internal `_frames` counter.
Returns:
bool: True if random frames should be used, False otherwise.
"""
if self.init_random_frames is None or self.init_random_frames <= 0:
return False
# Use replay_buffer.write_count when available for accurate counting in .start() mode
if self.replay_buffer is not None:
return self.replay_buffer.write_count < self.init_random_frames
return self._frames < self.init_random_frames
def _validate_cat_results(self, cat_results: str | int | None) -> None:
"""Validate cat_results parameter."""
if cat_results is not None and (
not isinstance(cat_results, (int, str))
or (isinstance(cat_results, str) and cat_results != "stack")
):
raise ValueError(
"cat_results must be a string ('stack') "
f"or an integer representing the cat dimension. Got {cat_results}."
)
# Lazy import to avoid circular dependency
from torchrl.collectors._multi_sync import MultiSyncCollector
if not isinstance(self, MultiSyncCollector) and cat_results not in (
"stack",
None,
):
raise ValueError(
"cat_results can only be used with ``MultiSyncCollector``."
)
self.cat_results = cat_results
def _check_replay_buffer_init(self):
if self.replay_buffer is None:
return
# Warn when a SliceSampler is used without trajectory writes: workers
# write batches independently so adjacent frames in the buffer can
# come from different episodes without an intervening done signal.
# This hazard is specific to multi-process collectors: a single
# Collector writes batches in temporal order, so consecutive batches
# are contiguous continuations of the same trajectories and the only
# mid-trajectory edge is the live write cursor, which SliceSampler
# already resolves at read time (see the trajectory-boundary section
# of the replay-buffer docs).
from torchrl.data.replay_buffers.samplers import SliceSampler
if (
self.replay_write_mode != "trajectory"
and isinstance(getattr(self.replay_buffer, "_sampler", None), SliceSampler)
and not self.set_truncated
):
warnings.warn(
"A SliceSampler is used with a multi-process collector but "
"replay_write_mode='trajectory' is not set and "
"set_truncated is False. "
"Adjacent frames in the replay buffer may come from different "
"workers and different episodes, causing SliceSampler to "
"sample slices that cross trajectory boundaries. "
"Consider setting replay_write_mode='trajectory' to write "
"only complete trajectories, or set_truncated=True to mark batch "
"boundaries (note: this introduces artificial truncations).",
category=UserWarning,
stacklevel=2,
)
if self.replay_write_mode == "trajectory":
# Trajectory assembly will happen at the worker level: each worker's
# inner Collector uses _iter_by_trajectories() to assemble complete
# trajectories and write them to the shared replay buffer.
# Null out trajs_per_batch on the multi-collector so that __iter__
# routes to self.iterator() directly (not _iter_by_trajectories,
# which would spin forever on the None yields from the RB path).
if self.trajs_per_batch is not None:
self._worker_trajs_per_batch = self.trajs_per_batch
self.trajs_per_batch = None
is_init = hasattr(self.replay_buffer, "_storage") and getattr(
self.replay_buffer._storage, "initialized", True
)
if is_init:
self._validate_policy_version_replay_schema()
if not is_init:
storage = self.replay_buffer._storage
if self._should_init_replay_buffer_from_worker(storage):
self._enable_replay_buffer_worker_init(storage)
self.replay_buffer.share()
return
if self.local_init_rb and getattr(storage, "shared_init", False):
# New behavior: storage handles all coordination itself
# Nothing to do here - the storage will coordinate during first write
self.replay_buffer.share()
return
# Legacy behavior: fake tensordict initialization
if isinstance(self.create_env_fn[0], EnvCreator):
fake_td = self.create_env_fn[0].meta_data.tensordict
elif isinstance(self.create_env_fn[0], EnvBase):
fake_td = self.create_env_fn[0].fake_tensordict()
else:
fake_td = self.create_env_fn[0](
**self.create_env_kwargs[0]
).fake_tensordict()
fake_td = self._add_collector_outputs_to_fake_td(fake_td)
if self.replay_write_mode == "trajectory":
# In trajectory mode, workers write flat 1-D timesteps to
# the buffer. Initialise the storage as 1-D so that the
# shapes match when real trajectories are written.
fake_td = fake_td.reshape(-1)[:1]
fake_td["collector", "traj_ids"] = torch.zeros(
fake_td.shape, dtype=torch.long
)
self.replay_buffer.extend(fake_td)
else:
fake_td["collector", "traj_ids"] = torch.zeros(
fake_td.shape, dtype=torch.long
)
# Use extend to avoid time-related transforms to fail
self.replay_buffer.extend(fake_td.unsqueeze(-1))
self.replay_buffer.empty()
def _should_init_replay_buffer_from_worker(self, storage):
return (
self.local_init_rb
and self.policy is None
and any(self.policy_factory)
and hasattr(storage, "shared_init")
and hasattr(storage, "_make_init_directory")
)
@staticmethod
def _enable_replay_buffer_worker_init(storage):
if getattr(storage, "_compilable", False):
raise RuntimeError(
"Cannot initialize a compilable replay-buffer storage from "
"multi-collector workers."
)
if getattr(storage, "shared_init", False):
return
storage_device = getattr(storage, "device", None)
if (
storage_device is not None
and storage_device != "auto"
and torch.device(storage_device).type != "cpu"
):
warnings.warn(
f"Worker-initialized replay buffers store data in a CPU "
f"memory-mapped tensordict; the storage device "
f"({storage_device}) cannot be honored and will be reset to "
f"'cpu' at initialization time."
)
storage.shared_init = True
storage._init_lock = mp.Lock()
storage._init_event = mp.Event()
storage._make_init_directory()
def _validate_policy_version_replay_schema(self):
if self.policy_version_tracker is None:
return
storage = getattr(self.replay_buffer._storage, "_storage", None)
if not isinstance(storage, TensorDictBase):
return
key = ("next", "policy_version")
if key not in storage.keys(True, True):
raise RuntimeError(
"The replay-buffer storage was initialized without the required "
"('next', 'policy_version') field. Use an uninitialized lazy "
"storage or include a field matching the configured policy-version "
"type at that key before enabling track_policy_version."
)
value = storage.get(key)
version_type = getattr(self.policy_version_tracker, "version_type", int)
if version_type in (int, "int"):
if not isinstance(value, torch.Tensor) or value.dtype != torch.int64:
dtype = getattr(value, "dtype", type(value).__name__)
raise RuntimeError(
"The replay-buffer field ('next', 'policy_version') must have "
f"dtype torch.int64 for integer policy versions, got {dtype}."
)
elif version_type in (str, "uuid") and isinstance(value, torch.Tensor):
raise RuntimeError(
"The replay-buffer field ('next', 'policy_version') must contain "
"non-tensor values for UUID policy versions."
)
def _add_collector_outputs_to_fake_td(self, fake_td):
policy = getattr(self, "policy", None)
out_keys = getattr(policy, "out_keys", None)
if out_keys:
with torch.no_grad():
policy_output = policy(fake_td.copy())
policy_output_keys = policy_output.keys(True, True)
for key in out_keys:
if key in fake_td.keys(True, True) or key not in policy_output_keys:
continue
fake_td.set(key, policy_output.get(key))
if self.policy_version_tracker is not None:
next_td = self.policy_version_tracker._step(fake_td, fake_td.get("next"))
fake_td.set("next", next_td)
return fake_td
[docs]
def fake_tensordict(self) -> TensorDictBase:
"""Not implemented for multi-process collectors.
Honoring the multi-collector contract here would require either
creating an env in the main process (which defeats the purpose of
a multi-process collector — Isaac Lab / mujoco-mjx etc. can only
run in workers) or routing a request to a worker over the pipe
(which requires workers to be alive and adds protocol surface).
Neither is implemented; call :meth:`~torchrl.collectors.Collector.fake_tensordict`
on a single :class:`~torchrl.collectors.Collector` instead, or
build the template directly from the env spec.
"""
raise NotImplementedError(
f"{type(self).__name__}.fake_tensordict() is not implemented. "
"Use Collector.fake_tensordict() on a single-process collector "
"for storage / cudagraph warmup, or build the template from the "
"env spec directly."
)
@classmethod
def _total_workers_from_env(cls, env_creators):
if isinstance(env_creators, (tuple, list)):
return sum(
cls._total_workers_from_env(env_creator) for env_creator in env_creators
)
from torchrl.envs import ParallelEnv
if isinstance(env_creators, ParallelEnv):
return env_creators.num_workers
return 1
def _get_devices(
self,
*,
storing_device: torch.device,
policy_device: torch.device,
env_device: torch.device,
device: torch.device,
):
# convert all devices to lists
if not isinstance(storing_device, (list, tuple)):
storing_device = [
storing_device,
] * self.num_workers
if not isinstance(policy_device, (list, tuple)):
policy_device = [
policy_device,
] * self.num_workers
if not isinstance(env_device, (list, tuple)):
env_device = [
env_device,
] * self.num_workers
if not isinstance(device, (list, tuple)):
device = [
device,
] * self.num_workers
if not (
len(device)
== len(storing_device)
== len(policy_device)
== len(env_device)
== self.num_workers
):
raise RuntimeError(
f"The length of the devices does not match the number of workers: {self.num_workers}."
)
storing_device, policy_device, env_device = zip(
*[
Collector._get_devices(
storing_device=storing_device,
policy_device=policy_device,
env_device=env_device,
device=device,
)
for (storing_device, policy_device, env_device, device) in zip(
storing_device, policy_device, env_device, device
)
]
)
return storing_device, policy_device, env_device
def frames_per_batch_worker(self, *, worker_idx: int | None = None) -> int:
raise NotImplementedError
@property
def _queue_len(self) -> int:
raise NotImplementedError
def _recv_and_check(
self,
pipe,
*,
timeout: float | None = None,
check_interval: float = 1.0,
worker_idx: int | None = None,
):
"""Receive from a pipe while periodically checking worker health.
This method prevents the main process from hanging indefinitely if a worker
dies while we're waiting for a response. It polls the pipe with a timeout
and checks if all worker processes are still alive between polls.
The overhead is minimal: if data is already available, `poll()` returns
immediately and no health check is performed. Health checks only run
when actually waiting for a slow response.
Args:
pipe: The pipe to receive from.
timeout: Maximum total time to wait for a message (seconds).
If None (default), wait indefinitely but still check worker health
periodically.
check_interval: How often to check worker health (seconds). Default 1.0.
worker_idx: Optional worker index for error messages.
Returns:
The received message.
Raises:
RuntimeError: If a worker process dies while waiting.
TimeoutError: If no message is received within the timeout (only if
timeout is not None).
"""
# Fast path: check if data is already available (no overhead)
if pipe.poll(0):
return pipe.recv()
# Slow path: wait with periodic health checks
elapsed = 0.0
while timeout is None or elapsed < timeout:
if pipe.poll(check_interval):
return pipe.recv()
elapsed += check_interval
# Check if any worker has died
_check_for_faulty_process(self.procs)
torchrl_logger.debug(
f"MultiCollector._recv_and_check: Still waiting after {elapsed:.1f}s"
+ (f" for worker {worker_idx}" if worker_idx is not None else "")
)
# Final check before timeout
_check_for_faulty_process(self.procs)
worker_info = f" from worker {worker_idx}" if worker_idx is not None else ""
raise TimeoutError(
f"Timed out after {timeout}s waiting for message{worker_info}. "
f"All workers are still alive - this may indicate a deadlock or very slow operation."
)
def _run_processes(self) -> None:
if self.num_threads is None:
total_workers = self._total_workers_from_env(self.create_env_fn)
self.num_threads = max(
1, torch.get_num_threads() - total_workers
) # 1 more thread for this proc
# Set up for worker processes
torch.set_num_threads(self.num_threads)
ctx = _get_mp_ctx()
# Best-effort global init (only if unset) to keep other mp users consistent.
_set_mp_start_method_if_unset(ctx.get_start_method())
if sys.platform == "linux" and ctx.get_start_method() == "spawn":
# On older PyTorch versions (< 2.8), pickling Process objects for "spawn"
# can pass file descriptors for shared storages, causing spawn-time failures.
# The strategy function returns "file_system" for old PyTorch, None otherwise.
strategy = _mp_sharing_strategy_for_spawn()
if strategy is not None:
mp.set_sharing_strategy(strategy)
queue_out = ctx.Queue(self._queue_len) # sends data from proc to main
self.queue_out = queue_out
self.procs = []
self._traj_pool = _TrajectoryPool(ctx=ctx, lock=True)
# Workers and coordinator each get a single-writer row. The extra
# coordinator row accounts for trajectory assembly performed by this
# process when no replay buffer is attached.
self._collector_progress = _CollectorProgress(self.num_workers + 1, ctx=ctx)
self._collector_progress_worker_idx = self.num_workers
self._collector_progress_aggregate = True
self._collector_progress_pending_on_ingest = True
# Create all pipes upfront (needed for weight sync scheme initialization)
# Store as list of (parent, child) tuples for use in worker creation
pipe_pairs = [ctx.Pipe() for _ in range(self.num_workers)]
# Extract parent pipes for external use (e.g., polling, receiving messages)
self.pipes = [pipe_parent for pipe_parent, _ in pipe_pairs]
_ProcessNoWarnCtx = _make_process_no_warn_cls(ctx)
# Initialize all weight sync schemes now that pipes are available
# Both SharedMemWeightSyncScheme (uses queues) and MultiProcessWeightSyncScheme (uses pipes)
# can be initialized here since all required resources exist
if self._weight_sync_schemes:
for model_id, scheme in self._weight_sync_schemes.items():
if not scheme.initialized_on_sender:
torchrl_logger.debug(
f"Init weight sync scheme {type(scheme).__name__} for {model_id=}."
)
scheme.init_on_sender(model_id=model_id, context=self, ctx=ctx)
# Create a policy on the right device
policy_factory = self.policy_factory
has_policy_factory = any(policy_factory)
if has_policy_factory:
policy_factory = [
CloudpickleWrapper(_policy_factory)
for _policy_factory in policy_factory
]
for i, (env_fun, env_fun_kwargs) in enumerate(
zip(self.create_env_fn, self.create_env_kwargs)
):
pipe_parent, pipe_child = pipe_pairs[i] # use pre-created pipes
if env_fun.__class__.__name__ != "EnvCreator" and not isinstance(
env_fun, EnvBase
): # to avoid circular imports
env_fun = CloudpickleWrapper(env_fun)
policy_device = self.policy_device[i]
storing_device = self.storing_device[i]
env_device = self.env_device[i]
# Prepare policy for worker based on weight synchronization method.
# IMPORTANT: when a policy_factory is provided, the policy instance
# is used ONLY on the main process (for weight extraction etc.) and
# is NOT sent to workers.
policy = self.policy
if self._weight_sync_schemes:
# With weight sync schemes, send stateless policies.
# Schemes handle weight distribution on worker side.
if has_policy_factory:
# Factory will create policy in worker; don't send policy.
policy_to_send = None
cm = contextlib.nullcontext()
elif policy is not None:
# Send a stateless policy down to workers: schemes apply weights.
policy_to_send = policy
cm = _make_meta_policy_cm(
policy, mp_start_method=ctx.get_start_method()
)
else:
policy_to_send = None
cm = contextlib.nullcontext()
elif hasattr(self, "_policy_weights_dict"):
# LEGACY:
# With weight updater, use in-place weight replacement.
# Take the weights and locally dispatch them to the policy before sending.
# This ensures a given set of shared weights for a device are shared
# for all policies that rely on that device.
policy_weights = self._policy_weights_dict.get(policy_device)
if has_policy_factory:
# Even in legacy mode, when a policy_factory is present, do not
# send the stateful policy down to workers.
policy_to_send = None
cm = contextlib.nullcontext()
else:
policy_to_send = policy
if policy is not None and policy_weights is not None:
cm = policy_weights.to_module(
policy, preserve_module_state=False
)
else:
cm = contextlib.nullcontext()
else:
# Parameter-less policy.
cm = contextlib.nullcontext()
# When a policy_factory exists, never send the policy instance.
policy_to_send = None if has_policy_factory else policy
with cm:
kwargs = {
"policy_factory": policy_factory[i],
"pipe_child": pipe_child,
"queue_out": queue_out,
"create_env_fn": env_fun,
"create_env_kwargs": env_fun_kwargs,
"policy": policy_to_send,
"max_frames_per_traj": self.max_frames_per_traj,
"frames_per_batch": self.frames_per_batch_worker(worker_idx=i),
"reset_at_each_iter": self.reset_at_each_iter,
"policy_device": policy_device,
"storing_device": storing_device,
"env_device": env_device,
"exploration_type": self.exploration_type,
"reset_when_done": self.reset_when_done,
"idx": i,
"interruptor": self.interruptor,
"set_truncated": self.set_truncated,
"use_buffers": self._use_buffers,
"replay_buffer": self.replay_buffer,
"extend_buffer": self.extend_buffer,
"traj_pool": self._traj_pool,
"trust_policy": self.trust_policy,
"compile_policy": (
self.compiled_policy_kwargs if self.compiled_policy else False
),
"cudagraph_policy": (
self.cudagraphed_policy_kwargs
if self.cudagraphed_policy
else False
),
"no_cuda_sync": self.no_cuda_sync,
"collector_class": self.collector_class,
"postproc": (
self.postproc if self.replay_buffer is not None else None
),
"weight_sync_schemes": self._weight_sync_schemes,
"worker_idx": i, # Worker index for queue-based weight distribution
"init_random_frames": self.init_random_frames,
"trajs_per_batch": self._worker_trajs_per_batch,
"trajs_per_write": self.trajs_per_write,
"replay_write_mode": self._worker_replay_write_mode,
"init_fn": self._worker_init_fn,
"auto_register_policy_transforms": self._auto_register_policy_transforms,
"track_policy_version": self.policy_version_tracker is not None,
"pre_collect_hook": self._worker_pre_collect_hook,
"post_collect_hook": self._worker_post_collect_hook,
"compact_obs": self.compact_obs,
"collector_progress": self._collector_progress,
}
proc = _ProcessNoWarnCtx(
target=_main_async_collector,
num_threads=self.num_sub_threads,
_start_method=ctx.get_start_method(),
kwargs=kwargs,
)
# proc.daemon can't be set as daemonic processes may be launched by the process itself
try:
proc.start()
except TypeError as err:
if "cannot pickle" in str(err):
raise RuntimeError(
"A non-serializable object was passed to the collector workers."
) from err
except RuntimeError as err:
if "Cowardly refusing to serialize non-leaf tensor" in str(err):
raise RuntimeError(
"At least one of the tensors in the policy, replay buffer, environment constructor or postprocessor requires gradients. "
"This is not supported in multiprocessed data collectors.\n- For ReplayBuffer transforms, use a `transform_factory` instead with `delayed_init=True`.\n"
"- Make sure your environment constructor does not reference tensors already instantiated on the main process.\n"
"- Since no gradient can be propagated through the Collector pipes, the backward graph is never needed. Consider using detached tensors instead."
) from err
elif "_share_fd_: only available on CPU" in str(
err
) or "_share_filename_: only available on CPU" in str(err):
# This is a common failure mode on older PyTorch versions when using the
# "spawn" multiprocessing start method: the process object contains a
# CUDA/MPS tensor (or a module/buffer on a non-CPU device), which must be
# pickled when spawning workers.
#
# See: https://github.com/pytorch/pytorch/issues/87688#issuecomment-1968901877
start_method = None
try:
start_method = mp.get_start_method(allow_none=True)
except Exception:
# Best effort: some environments may disallow querying here.
start_method = None
raise RuntimeError(
"Failed to start a collector worker process because a non-CPU tensor "
"was captured in the worker process arguments and had to be serialized "
"(pickled) at process start.\n\n"
f"Detected multiprocessing start method: {start_method!r}.\n\n"
"Workarounds:\n"
"- Keep any tensors/modules referenced by your collector constructor "
"(policy, replay buffer, postprocs, env factory captures, etc.) on CPU "
"when using a spawning start method (common on macOS/Windows).\n"
"- Or set the multiprocessing start method to 'fork' *before* creating "
"the collector (Unix only). Example:\n\n"
" import torch.multiprocessing as mp\n"
" if __name__ == '__main__':\n"
" mp.set_start_method('fork', force=True)\n\n"
"Upstream context: https://github.com/pytorch/pytorch/issues/87688#issuecomment-1968901877"
) from err
else:
raise err
except ValueError as err:
if "bad value(s) in fds_to_keep" in str(err):
# This error occurs on old Python versions (e.g., 3.9) with old PyTorch (e.g., 2.3)
# when using the spawn multiprocessing start method. The spawn implementation tries to
# preserve file descriptors across exec, but some descriptors may be invalid/closed.
# This is a compatibility issue with old Python multiprocessing implementations.
python_version = (
f"{sys.version_info.major}.{sys.version_info.minor}"
)
raise RuntimeError(
f"Failed to start collector worker process due to file descriptor issues "
f"with spawn multiprocessing on Python {python_version}.\n\n"
f"This is a known compatibility issue with old Python/PyTorch stacks. "
f"Consider upgrading to Python >= 3.10 and PyTorch >= 2.5, or use the 'fork' "
f"multiprocessing start method on Unix systems.\n\n"
f"Workarounds:\n"
f"- Upgrade Python to >= 3.10 and PyTorch to >= 2.5\n"
f"- On Unix systems, force fork start method:\n"
f" import torch.multiprocessing as mp\n"
f" if __name__ == '__main__':\n"
f" mp.set_start_method('fork', force=True)\n\n"
f"Upstream Python issue: https://github.com/python/cpython/issues/87706"
) from err
except _pickle.PicklingError as err:
if "<lambda>" in str(err):
raise RuntimeError(
"""Can't open a process with doubly cloud-pickled lambda function.
This error is likely due to an attempt to use a ParallelEnv in a
multiprocessed data collector. To do this, consider wrapping your
lambda function in an `torchrl.envs.EnvCreator` wrapper as follows:
`env = ParallelEnv(N, EnvCreator(my_lambda_function))`.
This will not only ensure that your lambda function is cloud-pickled once, but
also that the state dict is synchronised across processes if needed."""
) from err
pipe_child.close()
self.procs.append(proc)
# Synchronize initial weights with workers AFTER starting processes but BEFORE waiting for "instantiated"
# This must happen after proc.start() but before workers send "instantiated" to avoid deadlock:
# Workers will call receiver.collect() during init and may block waiting for data
if self._weight_sync_schemes:
# start with policy
policy_scheme = self._weight_sync_schemes.get("policy")
if policy_scheme is not None:
policy_scheme.connect()
for key, scheme in self._weight_sync_schemes.items():
if key == "policy":
continue
scheme.connect()
# Wait for workers to be ready
for i, pipe_parent in enumerate(self.pipes):
pipe_parent.poll(timeout=INSTANTIATE_TIMEOUT)
try:
msg = pipe_parent.recv()
except EOFError as e:
raise RuntimeError(
f"Worker {i} failed to initialize and closed the connection before sending status. "
f"This typically indicates that the worker process crashed during initialization. "
f"Check the worker process logs for the actual error."
) from e
if msg != "instantiated":
# Check if it's an error dict from worker
if isinstance(msg, dict) and msg.get("error"):
# Reconstruct the exception from the worker
exc_type_name = msg["exception_type"]
exc_msg = msg["exception_msg"]
traceback_str = msg["traceback"]
# Try to get the actual exception class
exc_class = None
exc_module = msg["exception_module"]
if exc_module == "builtins":
# Get from builtins
import builtins
exc_class = getattr(builtins, exc_type_name, None)
else:
# Try to import from the module
try:
import importlib
mod = importlib.import_module(exc_module)
exc_class = getattr(mod, exc_type_name, None)
except Exception:
pass
# Re-raise with original exception type if possible
if exc_class is not None:
raise exc_class(
f"{exc_msg}\n\nWorker traceback:\n{traceback_str}"
)
else:
# Fall back to RuntimeError if we can't get the original type
raise RuntimeError(
f"Worker {i} raised {exc_type_name}: {exc_msg}\n\nWorker traceback:\n{traceback_str}"
)
else:
# Legacy string error message
raise RuntimeError(msg)
self.queue_out = queue_out
self.closed = False
_running_free = False
[docs]
def start(self):
"""Starts the collector(s) for asynchronous data collection.
The collected data is stored in the provided replay buffer. This method initiates the background collection of
data across multiple processes, allowing for decoupling of data collection and training.
Raises:
RuntimeError: If no replay buffer is defined during the collector's initialization.
Example:
>>> from torchrl.modules import RandomPolicy >>> >>> import time
>>> from functools import partial
>>>
>>> import tqdm
>>>
>>> from torchrl.collectors import MultiAsyncCollector
>>> from torchrl.data import LazyTensorStorage, ReplayBuffer
>>> from torchrl.envs import GymEnv, set_gym_backend
>>> import ale_py
>>>
>>> # Set the gym backend to gymnasium
>>> set_gym_backend("gymnasium").set()
>>>
>>> if __name__ == "__main__":
... # Create a random policy for the Pong environment
... env_fn = partial(GymEnv, "ALE/Pong-v5")
... policy = RandomPolicy(env_fn().action_spec)
...
... # Initialize a shared replay buffer
... rb = ReplayBuffer(storage=LazyTensorStorage(10000), shared=True)
...
... # Create a multi-async data collector with 16 environments
... num_envs = 16
... collector = MultiAsyncCollector(
... [env_fn] * num_envs,
... policy=policy,
... replay_buffer=rb,
... frames_per_batch=num_envs * 16,
... total_frames=-1,
... )
...
... # Progress bar to track the number of collected frames
... pbar = tqdm.tqdm(total=100_000)
...
... # Start the collector asynchronously
... collector.start()
...
... # Track the write count of the replay buffer
... prec_wc = 0
... while True:
... wc = rb.write_count
... c = wc - prec_wc
... prec_wc = wc
...
... # Update the progress bar
... pbar.update(c)
... pbar.set_description(f"Write Count: {rb.write_count}")
...
... # Check the write count every 0.5 seconds
... time.sleep(0.5)
...
... # Stop when the desired number of frames is reached
... if rb.write_count . 100_000:
... break
...
... # Shut down the collector
... collector.async_shutdown()
"""
if self.replay_buffer is None:
raise RuntimeError("Replay buffer must be defined for execution.")
self._running_free = True
torchrl_logger.debug(
f"MultiCollector.start(): Sending run_free to {len(self.pipes)} workers..."
)
for i, pipe in enumerate(self.pipes):
pipe.send((None, "run_free"))
torchrl_logger.debug(f"MultiCollector.start(): Sent run_free to worker {i}")
[docs]
@contextlib.contextmanager
def pause(self):
"""Context manager that pauses the collector if it is running free."""
if self._running_free:
for pipe in self.pipes:
pipe.send((None, "pause"))
# Make sure all workers are paused
for i in range(len(self.pipes)):
# Use timeout with health check to avoid hanging if a worker dies
timeout = 30.0
check_interval = 1.0
elapsed = 0.0
while elapsed < timeout:
try:
idx, msg = self.queue_out.get(timeout=check_interval)
break
except Exception:
elapsed += check_interval
_check_for_faulty_process(self.procs)
else:
_check_for_faulty_process(self.procs)
raise TimeoutError(
f"Timed out waiting for worker {i} to pause after {timeout}s"
)
if msg != "paused":
raise ValueError(f"Expected paused, but got {msg=}.")
torchrl_logger.debug(f"Worker {idx} is paused.")
self._running_free = False
yield None
for pipe in self.pipes:
pipe.send((None, "restart"))
self._running_free = True
else:
raise RuntimeError("Collector cannot be paused.")
def _install_profile_hooks(self, config: ProfileConfig) -> None:
"""Install per-worker :class:`_ProfilerHook` on each selected worker.
Each worker gets its own ``_ProfilerHook(config, worker_idx=idx)``
wrapped in :class:`CloudpickleWrapper` and pushed via the existing
``"setattr"`` pipe message. Workers not in ``config.workers`` are left
untouched.
"""
if not hasattr(self, "pipes") or getattr(self, "closed", True):
raise RuntimeError(
"MultiCollector workers are not running; call enable_profile() "
"before close()/shutdown()."
)
_check_for_faulty_process(self.procs)
targeted = [idx for idx in config.workers if idx < self.num_workers]
for idx in targeted:
hook = CloudpickleWrapper(_ProfilerHook(config, worker_idx=idx))
self.pipes[idx].send((("post_collect_hook", hook), "setattr"))
for idx in targeted:
result, msg = self._recv_and_check(self.pipes[idx], worker_idx=idx)
if msg != "setattr":
raise RuntimeError(f"Worker {idx}: Expected 'setattr' ack, got {msg}")
if isinstance(result, Exception):
raise result
def _uninstall_profile_hooks(self, config: ProfileConfig) -> None:
"""Stop the per-worker profiler hooks and clear ``post_collect_hook``."""
if not hasattr(self, "pipes") or getattr(self, "closed", True):
return
targeted = [idx for idx in config.workers if idx < self.num_workers]
# Stop the inner profiler — early-stop is harmless if it already auto-stopped.
for idx in targeted:
self.pipes[idx].send(
(("post_collect_hook.fn.stop", (), {}), "cascade_execute")
)
for idx in targeted:
try:
result, msg = self._recv_and_check(self.pipes[idx], worker_idx=idx)
if msg != "cascade_execute":
raise RuntimeError(
f"Worker {idx}: Expected 'cascade_execute' ack, got {msg}"
)
if isinstance(result, Exception):
raise result
except Exception:
# Swallow per-worker stop errors — we still want to clear the hook below.
pass
# Clear the hook on each targeted worker.
for idx in targeted:
self.pipes[idx].send((("post_collect_hook", None), "setattr"))
for idx in targeted:
self._recv_and_check(self.pipes[idx], worker_idx=idx)
def _set_worker_attr(self, attr_name: str, value: Any) -> None:
if not hasattr(self, "pipes") or getattr(self, "closed", True):
return
_check_for_faulty_process(self.procs)
for pipe in self.pipes:
pipe.send(((attr_name, value), "setattr"))
for idx, pipe in enumerate(self.pipes):
result, msg = self._recv_and_check(pipe, worker_idx=idx)
if msg != "setattr":
raise RuntimeError(f"Expected msg='setattr', got {msg}")
if isinstance(result, Exception):
raise result
@property
def pre_collect_hook(self) -> Callable[[], None] | None:
return self._pre_collect_hook
@pre_collect_hook.setter
def pre_collect_hook(self, hook: Callable[[], None] | None) -> None:
self._pre_collect_hook = hook
self._worker_pre_collect_hook = (
CloudpickleWrapper(hook) if hook is not None else None
)
self._set_worker_attr("pre_collect_hook", self._worker_pre_collect_hook)
@property
def post_collect_hook(self) -> Callable[[TensorDictBase], None] | None:
return self._post_collect_hook
@post_collect_hook.setter
def post_collect_hook(self, hook: Callable[[TensorDictBase], None] | None) -> None:
self._post_collect_hook = hook
self._worker_post_collect_hook = (
CloudpickleWrapper(hook) if hook is not None else None
)
self._set_worker_attr("post_collect_hook", self._worker_post_collect_hook)
def _normalize_worker_calls(
self,
list_of_args: list[tuple] | None = None,
list_of_kwargs: list[dict] | None = None,
) -> tuple[list[tuple], list[dict]]:
if list_of_args is None and list_of_kwargs is None:
list_of_args = [()] * self.num_workers
list_of_kwargs = [{}] * self.num_workers
elif list_of_args is None:
list_of_args = [()] * len(list_of_kwargs)
elif list_of_kwargs is None:
list_of_kwargs = [{}] * len(list_of_args)
if len(list_of_args) != self.num_workers:
raise ValueError(
f"Expected {self.num_workers} argument entries, got {len(list_of_args)}."
)
if len(list_of_kwargs) != self.num_workers:
raise ValueError(
f"Expected {self.num_workers} keyword-argument entries, got {len(list_of_kwargs)}."
)
return list_of_args, list_of_kwargs
[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 worker collector."""
list_of_args, list_of_kwargs = self._normalize_worker_calls(
list_of_args, list_of_kwargs
)
_check_for_faulty_process(self.procs)
for pipe, args, kwargs in zip(self.pipes, list_of_args, list_of_kwargs):
pipe.send(((method_name, args, kwargs), "cascade_execute"))
results = []
for idx, pipe in enumerate(self.pipes):
result, msg = self._recv_and_check(pipe, worker_idx=idx)
if msg != "cascade_execute":
raise RuntimeError(f"Expected msg='cascade_execute', got {msg}")
if isinstance(result, Exception):
raise result
results.append(result)
return results
[docs]
def get_distant_attr(self, attr: str) -> list[Any]:
"""Get a nested attribute from each worker collector."""
_check_for_faulty_process(self.procs)
for pipe in self.pipes:
pipe.send((attr, "get_distant_attr"))
results = []
for idx, pipe in enumerate(self.pipes):
result, msg = self._recv_and_check(pipe, worker_idx=idx)
if msg != "get_distant_attr":
raise RuntimeError(f"Expected msg='get_distant_attr', got {msg}")
if isinstance(result, Exception):
raise result
results.append(result)
return results
def __del__(self):
try:
self.shutdown()
except Exception:
# an AttributeError will typically be raised if the collector is deleted when the program ends.
# In the future, insignificant changes to the close method may change the error type.
# We excplicitely assume that any error raised during closure in
# __del__ will not affect the program.
pass
[docs]
def shutdown(
self,
timeout: float | None = None,
close_env: bool = True,
raise_on_error: bool = True,
) -> None:
"""Shuts down all processes. This operation is irreversible.
Args:
timeout (float, optional): The timeout for closing pipes between workers.
close_env (bool, optional): Whether to close the environment. Defaults to `True`.
raise_on_error (bool, optional): Whether to raise an error if the shutdown fails. Defaults to `True`.
"""
if not close_env:
raise RuntimeError(
f"Cannot shutdown {type(self).__name__} collector without environment being closed."
)
try:
self._shutdown_main(timeout)
except Exception as e:
if raise_on_error:
raise e
else:
pass
def _shutdown_main(self, timeout: float | None = None) -> None:
if timeout is None:
timeout = 10
if timeout < 0:
raise ValueError(f"timeout must be non-negative, got {timeout}.")
if getattr(self, "_shutdown_complete", False):
return
# Claim shutdown before touching any IPC resource. This makes repeated
# calls no-ops even when the first call finds an already-dead worker or
# a pipe that the worker has closed.
self.closed = True
started_at = time.monotonic()
# Reserve part of the caller's single deadline for forced cleanup:
# terminate after 80%, then retain the final 5% for a kill fallback.
graceful_timeout = timeout * 0.8
terminate_timeout = timeout * 0.95
procs = getattr(self, "procs", ())
pipes = getattr(self, "pipes", ())
try:
for proc, pipe in zip(procs, pipes):
if proc._closed or not proc.is_alive():
continue
try:
pipe.send((None, "close"))
except (EOFError, OSError, ValueError):
# The worker may already have observed the parent closing
# its end of the channel. It will be joined below.
pass
# Workers acknowledge "close" on their control pipes, but process
# exit is the shutdown contract; draining those replies would spend
# the same bounded deadline without improving teardown guarantees.
pending = {
idx
for idx, proc in enumerate(procs)
if not proc._closed and proc.is_alive()
}
while pending and time.monotonic() - started_at < graceful_timeout:
for idx in tuple(pending):
proc = procs[idx]
if proc._closed or not proc.is_alive():
pending.remove(idx)
continue
remaining = graceful_timeout - (time.monotonic() - started_at)
if remaining <= 0:
break
proc.join(timeout=min(0.05, remaining))
if proc._closed or not proc.is_alive():
pending.remove(idx)
stragglers = [
proc for proc in procs if not proc._closed and proc.is_alive()
]
for proc in stragglers:
try:
proc.terminate()
except (OSError, ValueError):
pass
for proc in stragglers:
remaining = terminate_timeout - (time.monotonic() - started_at)
proc.join(timeout=max(remaining, 0.0))
stragglers = [
proc for proc in procs if not proc._closed and proc.is_alive()
]
for proc in stragglers:
try:
proc.kill()
except (AttributeError, OSError, ValueError):
# ``kill`` is unavailable on older Python versions. The
# earlier terminate request still bounds the parent call.
pass
for proc in stragglers:
remaining = timeout - (time.monotonic() - started_at)
proc.join(timeout=max(remaining, 0.0))
finally:
import torchrl
queue_out = getattr(self, "queue_out", None)
if queue_out is not None:
try:
queue_out.close()
except (OSError, ValueError):
pass
for pipe in pipes:
try:
pipe.close()
except (OSError, ValueError):
pass
num_threads = min(
torchrl._THREAD_POOL_INIT,
torch.get_num_threads()
+ self._total_workers_from_env(self.create_env_fn),
)
torch.set_num_threads(num_threads)
progress = getattr(self, "_collector_progress", None)
if progress is not None:
progress.clear_pending()
self._shutdown_complete = True
[docs]
def async_shutdown(self, timeout: float | None = None):
return self.shutdown(timeout=timeout)
[docs]
def set_seed(self, seed: int, static_seed: bool = False) -> int:
"""Sets the seeds of the environments stored in the DataCollector.
Args:
seed: integer representing the seed to be used for the environment.
static_seed (bool, optional): if ``True``, the seed is not incremented.
Defaults to False
Returns:
Output seed. This is useful when more than one environment is
contained in the DataCollector, as the seed will be incremented for
each of these. The resulting seed is the seed of the last
environment.
Examples:
>>> from torchrl.envs import ParallelEnv
>>> from torchrl.envs.libs.gym import GymEnv
>>> from tensordict.nn import TensorDictModule
>>> from torch import nn
>>> env_fn = lambda: GymEnv("Pendulum-v1")
>>> env_fn_parallel = lambda: ParallelEnv(6, env_fn)
>>> policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
>>> collector = Collector(env_fn_parallel, policy, frames_per_batch=100, total_frames=300)
>>> out_seed = collector.set_seed(1) # out_seed = 6
"""
_check_for_faulty_process(self.procs)
for idx in range(self.num_workers):
self.pipes[idx].send(((seed, static_seed), "seed"))
new_seed, msg = self._recv_and_check(self.pipes[idx], worker_idx=idx)
if msg != "seeded":
raise RuntimeError(f"Expected msg='seeded', got {msg}")
seed = new_seed
self.reset()
return seed
[docs]
def reset(self, reset_idx: Sequence[bool] | None = None) -> None:
"""Resets the environments to a new initial state.
Args:
reset_idx: Optional. Sequence indicating which environments have
to be reset. If None, all environments are reset.
"""
_check_for_faulty_process(self.procs)
# drop parent-level trajectory-assembly state (trajs_per_batch
# without a replay buffer assembles trajectories on this process)
self._flush_trajectory_assembly()
if reset_idx is None:
reset_idx = [True for _ in range(self.num_workers)]
for idx in range(self.num_workers):
if reset_idx[idx]:
self.pipes[idx].send((None, "reset"))
for idx in range(self.num_workers):
if reset_idx[idx]:
j, msg = self._recv_and_check(self.pipes[idx], worker_idx=idx)
if msg != "reset":
raise RuntimeError(f"Expected msg='reset', got {msg}")
[docs]
def state_dict(self) -> OrderedDict:
"""Returns the state_dict of the data collector.
Each field represents a worker containing its own state_dict.
"""
for idx in range(self.num_workers):
self.pipes[idx].send((None, "state_dict"))
state_dict = OrderedDict()
traj_pool_state = None
for idx in range(self.num_workers):
_state_dict, msg = self._recv_and_check(self.pipes[idx], worker_idx=idx)
if msg != "state_dict":
raise RuntimeError(f"Expected msg='state_dict', got {msg}")
worker_traj_pool_state = _state_dict.pop("traj_pool", None)
if traj_pool_state is None and worker_traj_pool_state is not None:
traj_pool_state = worker_traj_pool_state
state_dict[f"worker{idx}"] = _state_dict
state_dict.update(
{
"frames": self._frames,
"iter": self._iter,
"collector_progress_coordinator": self._progress_state_dict(),
}
)
if traj_pool_state is not None:
state_dict["traj_pool"] = traj_pool_state
if self.policy_version_tracker is not None:
state_dict["policy_version"] = self.policy_version
return state_dict
[docs]
def load_state_dict(self, state_dict: OrderedDict) -> None:
"""Loads the state_dict on the workers.
Args:
state_dict (OrderedDict): state_dict of the form
``{"worker0": state_dict0, "worker1": state_dict1}``.
"""
traj_pool_state = state_dict.get("traj_pool")
if traj_pool_state is not None:
self._traj_pool.load_state_dict(traj_pool_state)
for idx in range(self.num_workers):
self.pipes[idx].send((state_dict[f"worker{idx}"], "load_state_dict"))
for idx in range(self.num_workers):
_, msg = self._recv_and_check(self.pipes[idx], worker_idx=idx)
if msg != "loaded":
raise RuntimeError(f"Expected msg='loaded', got {msg}")
self._frames = state_dict["frames"]
self._iter = state_dict["iter"]
self._load_progress_state_dict(state_dict.get("collector_progress_coordinator"))
if "policy_version" in state_dict and self.policy_version_tracker is not None:
policy_version = state_dict["policy_version"]
self._acknowledged_policy_version = policy_version
if policy_version is not None:
self.policy_version_tracker.version = policy_version
[docs]
def increment_version(self):
"""Increment the policy version."""
if self.policy_version_tracker is not None:
if not hasattr(self.policy_version_tracker, "increment_version"):
raise RuntimeError(
"Policy version tracker is not a PolicyVersion instance. Please pass a PolicyVersion instance to the collector."
)
self.policy_version_tracker.increment_version()
self._acknowledged_policy_version = self.policy_version_tracker.version
@property
def policy_version(self) -> str | int | None:
"""The policy version acknowledged by every worker.
For multi-process collectors, worker-local
:class:`~torchrl.envs.transforms.PolicyVersion`
transforms write the per-frame ``"policy_version"`` values in returned
batches and remain the source of truth for collected data. This property
returns ``None`` after a partial-worker or asynchronous update because
no single scalar is then known to describe every worker. Worker-local
versions can be queried explicitly with :meth:`worker_policy_versions`.
"""
return self._acknowledged_policy_version
[docs]
def get_policy_version(self) -> str | int | None:
"""Get the policy version acknowledged by every worker.
This method exists to support remote calls in Ray actors, since properties
cannot be accessed directly through Ray's RPC mechanism.
Returns:
The aggregate version number (int) or UUID (str), or ``None`` if
version tracking is disabled or workers may differ.
"""
return self.policy_version
[docs]
def worker_policy_versions(self) -> dict[int, str | int | None]:
"""Query the policy version currently reported by each worker.
Unlike :attr:`policy_version`, this method performs worker RPCs and can
expose divergent versions after partial updates. It shares the worker
control channels with other coordinator commands and therefore should
not race with weight updates issued from another thread.
Returns:
A mapping from worker index to its local policy version. Values are
``None`` when policy-version tracking is disabled.
"""
versions = dict(enumerate(self.map_fn("get_policy_version")))
if getattr(self, "running", False):
# map_fn has consumed every worker reply, so the control pipes are
# ready for the command that resumes asynchronous collection.
msg = "continue_random" if self._should_use_random_frames() else "continue"
for idx, pipe in enumerate(self.pipes):
pipe.send((idx, msg))
return versions
[docs]
def getattr_policy(self, attr):
"""Get an attribute from the policy of the first worker.
Args:
attr (str): The attribute name to retrieve from the policy.
Returns:
The attribute value from the policy of the first worker.
Raises:
AttributeError: If the attribute doesn't exist on the policy.
"""
_check_for_faulty_process(self.procs)
# Send command to first worker (index 0)
self.pipes[0].send((attr, "getattr_policy"))
result, msg = self._recv_and_check(self.pipes[0], worker_idx=0)
if msg != "getattr_policy":
raise RuntimeError(f"Expected msg='getattr_policy', got {msg}")
# If the worker returned an AttributeError, re-raise it
if isinstance(result, AttributeError):
raise result
return result
[docs]
def getattr_env(self, attr):
"""Get an attribute from the environment of the first worker.
Args:
attr (str): The attribute name to retrieve from the environment.
Returns:
The attribute value from the environment of the first worker.
Raises:
AttributeError: If the attribute doesn't exist on the environment.
"""
_check_for_faulty_process(self.procs)
# Send command to first worker (index 0)
self.pipes[0].send((attr, "getattr_env"))
result, msg = self._recv_and_check(self.pipes[0], worker_idx=0)
if msg != "getattr_env":
raise RuntimeError(f"Expected msg='getattr_env', got {msg}")
# If the worker returned an AttributeError, re-raise it
if isinstance(result, AttributeError):
raise result
return result
[docs]
def getattr_rb(self, attr):
"""Get an attribute from the replay buffer."""
return getattr(self.replay_buffer, attr)
[docs]
def get_model(self, model_id: str):
"""Get model instance by ID (for weight sync schemes).
Args:
model_id: Model identifier (e.g., "policy", "value_net")
Returns:
The model instance
Raises:
ValueError: If model_id is not recognized
"""
if model_id == "policy":
# Return the fallback policy instance
if (fallback_policy := getattr(self, "_fallback_policy", None)) is not None:
return fallback_policy
elif hasattr(self, "policy") and self.policy is not None:
return self.policy
else:
raise ValueError(f"No policy found for model_id '{model_id}'")
else:
# Try to resolve via attribute access
return _resolve_model(self, model_id)
[docs]
def get_cached_weights(self, model_id: str):
"""Get cached shared memory weights if available (for weight sync schemes).
Args:
model_id: Model identifier
Returns:
Cached TensorDict weights or None if not available
"""
if model_id == "policy" and hasattr(self, "_policy_weights_dict"):
# Get the policy device (first device if list)
policy_device = self.policy_device
if isinstance(policy_device, (list, tuple)):
policy_device = policy_device[0] if len(policy_device) > 0 else None
# Return cached weights for this device
return self._policy_weights_dict.get(policy_device)
return None
def _weight_update_impl(
self,
policy_or_weights: TensorDictBase | nn.Module | 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:
"""Update weights on workers.
Weight sync schemes now use background threads on the receiver side.
The scheme's send() method:
1. Puts weights in the queue (or updates shared memory)
2. Sends a "receive" instruction to the worker's background thread
3. Waits for acknowledgment (if sync=True)
No pipe signaling is needed - the scheme handles everything internally.
"""
# Call parent implementation which calls scheme.send()
# The scheme handles instruction delivery and acknowledgments
super()._weight_update_impl(
policy_or_weights=policy_or_weights,
worker_ids=worker_ids,
model_id=model_id,
weights_dict=weights_dict,
**kwargs,
)
# for RPC
[docs]
def receive_weights(self, policy_or_weights: TensorDictBase | None = None):
return super().receive_weights(policy_or_weights)
# for RPC
def _receive_weights_scheme(self, model_version: int | None = None):
return super()._receive_weights_scheme(model_version=model_version)
# Backward-compatible alias (deprecated, use MultiCollector instead)
MultiCollector = MultiCollector
