RayCollector

class torchrl.collectors.distributed.RayCollector(create_env_fn: collections.abc.Callable | torchrl.envs.common.EnvBase | list[collections.abc.Callable] | list[torchrl.envs.common.EnvBase], policy: collections.abc.Callable[[tensordict.base.TensorDictBase], tensordict.base.TensorDictBase] | None = None, *, policy_factory: collections.abc.Callable[[], collections.abc.Callable] | list[collections.abc.Callable[[], collections.abc.Callable]] | None = None, trust_policy: bool | None = None, frames_per_batch: int, total_frames: int = -1, device: torch.device | list[torch.device] | None = None, storing_device: torch.device | list[torch.device] | None = None, env_device: torch.device | list[torch.device] | None = None, policy_device: torch.device | list[torch.device] | None = None, max_frames_per_traj=-1, init_random_frames=-1, reset_at_each_iter=False, postproc=None, split_trajs=False, exploration_type=InteractionType.RANDOM, collector_class: ~collections.abc.Callable[[~tensordict._td.TensorDict], ~tensordict._td.TensorDict] = <class 'torchrl.collectors.collectors.SyncDataCollector'>, collector_kwargs: dict[str, typing.Any] | list[dict] | None = None, num_workers_per_collector: int = 1, sync: bool = False, ray_init_config: dict[str, typing.Any] | None = None, remote_configs: dict[str, typing.Any] | list[dict[str, typing.Any]] | None = None, num_collectors: int | None = None, update_after_each_batch: bool = False, max_weight_update_interval: int = -1, replay_buffer: torchrl.data.replay_buffers.replay_buffers.ReplayBuffer | None = None, weight_updater: torchrl.collectors.weight_update.WeightUpdaterBase | collections.abc.Callable[[], torchrl.collectors.weight_update.WeightUpdaterBase] | None = None, weight_sync_schemes: dict[str, torchrl.weight_update.weight_sync_schemes.WeightSyncScheme] | None = None, use_env_creator: bool = False, no_cuda_sync: bool | None = None)

Distributed data collector with Ray backend.

This Python class serves as a ray-based solution to instantiate and coordinate multiple data collectors in a distributed cluster. Like TorchRL non-distributed collectors, this collector is an iterable that yields TensorDicts until a target number of collected frames is reached, but handles distributed data collection under the hood.

The class dictionary input parameter “ray_init_config” can be used to provide the kwargs to call Ray initialization method ray.init(). If “ray_init_config” is not provided, the default behavior is to autodetect an existing Ray cluster or start a new Ray instance locally if no existing cluster is found. Refer to Ray documentation for advanced initialization kwargs.

Similarly, dictionary input parameter “remote_configs” can be used to specify the kwargs for ray.remote() when called to create each remote collector actor, including collector compute resources.The sum of all collector resources should be available in the cluster. Refer to Ray documentation for advanced configuration of the ray.remote() method. Default kwargs are:

>>> kwargs = {
...     "num_cpus": 1,
...     "num_gpus": 0.2,
...     "memory": 2 * 1024 ** 3,
... }

The coordination between collector instances can be specified as “synchronous” or “asynchronous”. In synchronous coordination, this class waits for all remote collectors to collect a rollout, concatenates all rollouts into a single TensorDict instance and finally yields the concatenated data. On the other hand, if the coordination is to be carried out asynchronously, this class provides the rollouts as they become available from individual remote collectors.

Parameters:

• create_env_fn (Callable or List[Callabled]) – list of Callables, each returning an instance of EnvBase.

• policy (Callable, optional) –

  Policy to be executed in the environment. Must accept tensordict.tensordict.TensorDictBase object as input. If None is provided, the policy used will be a RandomPolicy instance with the environment action_spec. Accepted policies are usually subclasses of 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 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 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 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.

Keyword Arguments:

• policy_factory (Callable[[], Callable], list of Callable[[], Callable], optional) –

  a callable (or list of callables) that returns a policy instance. This is exclusive with the policy argument.

  Note

  policy_factory comes in handy whenever the policy cannot be serialized.

• 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.

• frames_per_batch (int) – A keyword-only argument representing the total number of elements in a batch.

• total_frames (int, Optional) – lower bound of the total number of frames returned by the collector. The iterator will stop once the total number of frames equates or exceeds the total number of frames passed to the collector. Default value is -1, which mean no target total number of frames (i.e. the collector will run indefinitely).

• 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). Lists of devices are supported.

• storing_device (int, str or torch.device, optional) – The remote device on which the output 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). Lists of devices are supported.

• env_device (int, str or torch.device, optional) – The remote 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. Lists of devices are supported.

• policy_device (int, str or torch.device, optional) – The remote 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. Lists of devices are supported.

• create_env_kwargs (dict, optional) – Dictionary of kwargs for create_env_fn.

• 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 Transform or a MultiStep instance. Defaults to None.

• split_trajs (bool, optional) – Boolean indicating whether the resulting TensorDict should be split according to the trajectories. See 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.

• collector_class (Python class or constructor) – a collector class to be remotely instantiated. Can be SyncDataCollector, MultiSyncDataCollector, MultiaSyncDataCollector or a derived class of these. Defaults to SyncDataCollector.

• collector_kwargs (dict or list, optional) – a dictionary of parameters to be passed to the remote data-collector. If a list is provided, each element will correspond to an individual set of keyword arguments for the dedicated collector.

• num_workers_per_collector (int) – the number of copies of the env constructor that is to be used on the remote nodes. Defaults to 1 (a single env per collector). On a single worker node all the sub-workers will be executing the same environment. If different environments need to be executed, they should be dispatched across worker nodes, not subnodes.

• ray_init_config (dict, Optional) – kwargs used to call ray.init().

• remote_configs (list of dicts, Optional) – ray resource specs for each remote collector. A single dict can be provided as well, and will be used in all collectors.

• num_collectors (int, Optional) – total number of collectors to be instantiated.

• sync (bool) – if True, the resulting tensordict is a stack of all the tensordicts collected on each node. If False (default), each tensordict results from a separate node in a “first-ready, first-served” fashion.

• update_after_each_batch (bool, optional) – if True, the weights will be updated after each collection. For sync=True, this means that all workers will see their weights updated. For sync=False, only the worker from which the data has been gathered will be updated. This is equivalent to max_weight_update_interval=0. Defaults to False, i.e. updates have to be executed manually through torchrl.collectors.DataCollector.update_policy_weights_()

• max_weight_update_interval (int, optional) – the maximum number of batches that can be collected before the policy weights of a worker is updated. For sync collections, this parameter is overwritten by update_after_each_batch. For async collections, it may be that one worker has not seen its parameters being updated for a certain time even if update_after_each_batch is turned on. Defaults to -1 (no forced update).

• replay_buffer (RayReplayBuffer, optional) –

  if provided, the collector will not yield tensordicts but populate the buffer instead. Defaults to None.

  Note

  although it is not enfoced (to allow users to implement their own replay buffer class), a RayReplayBuffer instance should be used here.

• weight_updater (WeightUpdaterBase or constructor, optional) – (Deprecated) An instance of WeightUpdaterBase or its subclass, responsible for updating the policy weights on remote inference workers managed by Ray. If not provided, a RayWeightUpdater will be used by default, leveraging Ray’s distributed capabilities. Consider using a constructor if the updater needs to be serialized.

• weight_sync_schemes (dict[str, WeightSyncScheme], optional) – Dictionary mapping model identifiers to WeightSyncScheme instances. This is the recommended way to configure weight synchronization. If not provided, defaults to {"policy": RayWeightSyncScheme()}.

• use_env_creator (bool, optional) – if True, the environment constructor functions will be wrapped in EnvCreator. This is useful for multiprocessed settings where shared memory needs to be managed, but Ray has its own object storage mechanism, so this is typically not needed. Defaults to False.

Examples

>>> from torch import nn
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.envs.libs.gym import GymEnv
>>> from torchrl.collectors import SyncDataCollector
>>> from torchrl.collectors.distributed import RayCollector
>>> env_maker = lambda: GymEnv("Pendulum-v1", device="cpu")
>>> policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
>>> distributed_collector = RayCollector(
...     create_env_fn=[env_maker],
...     policy=policy,
...     collector_class=SyncDataCollector,
...     max_frames_per_traj=50,
...     init_random_frames=-1,
...     reset_at_each_iter=-False,
...     collector_kwargs={
...         "device": "cpu",
...         "storing_device": "cpu",
...     },
...     num_collectors=1,
...     total_frames=10000,
...     frames_per_batch=200,
... )
>>> for i, data in enumerate(collector):
...     if i == 2:
...         print(data)
...         break

add_collectors(create_env_fn, num_envs, policy, collector_kwargs, remote_configs)

Creates and adds a number of remote collectors to the set.

async async_shutdown(shutdown_ray: bool = False)

Finishes processes started by the collector during async execution.

Parameters:

shutdown_ray (bool) – If True, also shutdown the Ray cluster. Defaults to False. Note: Setting this to True will kill all Ray actors in the cluster, including any replay buffers or other services. Only set to True if you’re sure you want to shut down the entire Ray cluster.

init_updater(*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.

Parameters:

• *args – Positional arguments for weight updater initialization

• **kwargs – Keyword arguments for weight updater initialization

load_state_dict(state_dict: collections.OrderedDict | list[collections.OrderedDict]) → None

Calls parent method for each remote collector.

local_policy()

Returns local collector.

pause()

Context manager that pauses the collector if it is running free.

property remote_collectors

Returns list of remote collectors.

set_seed(seed: int, static_seed: bool = False) → list[int]

Calls parent method for each remote collector iteratively and returns final seed.

shutdown(timeout: float | None = None, shutdown_ray: bool = False) → None

Finishes processes started by the collector.

Parameters:

• timeout (float, optional) – Timeout for stopping the collection thread.

• shutdown_ray (bool) – If True, also shutdown the Ray cluster. Defaults to False. Note: Setting this to True will kill all Ray actors in the cluster, including any replay buffers or other services. Only set to True if you’re sure you want to shut down the entire Ray cluster.

start()

Starts the RayCollector in a background thread.

state_dict() → list[collections.OrderedDict]

Calls parent method for each remote collector and returns a list of results.

stop_remote_collectors()

Stops all remote collectors.

update_policy_weights_(policy_or_weights: tensordict.base.TensorDictBase | tensordict.nn.common.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

Updates the policy weights for the data collector, accommodating both local and remote execution contexts.

This method ensures that the policy weights used by the data collector are synchronized with the latest trained weights. It supports both local and remote weight updates, depending on the configuration of the data collector. The local (download) update is performed before the remote (upload) update, such that weights can be transferred to the children workers from a server.

Parameters:

• policy_or_weights (TensorDictBase | TensorDictModuleBase | dict | None) – The weights to update with. Can be: - TensorDictModuleBase: A policy module whose weights will be extracted - TensorDictBase: A TensorDict containing weights - dict: A regular dict containing weights - None: Will try to get weights from server using _get_server_weights()

• worker_ids (int | List[int] | torch.device | List[torch.device] | None, optional) – Identifiers for the workers that need to be updated. This is relevant when the collector has more than one worker associated with it.

• model_id (str | None, optional) – The model identifier to update. If provided, only updates this specific model. Cannot be used together with weights_dict.

• weights_dict (dict[str, Any] | None, optional) – Dictionary mapping model_id to weights for updating multiple models atomically. Keys should match the model_ids registered in weight_sync_schemes. Cannot be used together with model_id or policy_or_weights.

Raises:

• TypeError – If worker_ids is provided but no weight_updater is configured.

• ValueError – If conflicting parameters are provided (e.g., both model_id and weights_dict).

Note

Users should extend the WeightUpdaterBase classes to customize the weight update logic for specific use cases. This method should not be overwritten.

See also

LocalWeightsUpdaterBase and RemoteWeightsUpdaterBase().