BaseCollector

class torchrl.collectors.BaseCollector

Base class for data collectors.

async_shutdown(timeout: float | None = None, close_env: bool = True) → None

Shuts down the collector when started asynchronously with the start method.

Parameters:

• timeout (float, optional) – The maximum time to wait for the collector to shutdown.

• close_env (bool, optional) – If True, the collector will close the contained environment. Defaults to True.

See also

start()

cascade_execute(attr_path: str, *args, **kwargs) → Any

Execute a method on a nested attribute of this collector.

This method allows remote callers to invoke methods on nested attributes of the collector without needing to know the full structure. It’s particularly useful for calling methods on weight sync schemes from the sender side.

Parameters:

• attr_path – Full path to the callable, e.g., “_receiver_schemes[‘model_id’]._set_dist_connection_info”

• *args – Positional arguments to pass to the method.

• **kwargs – Keyword arguments to pass to the method.

Returns:

The return value of the method call.

Examples

>>> collector.cascade_execute(
...     "_receiver_schemes['policy']._set_dist_connection_info",
...     connection_info_ref,
...     worker_idx=0
... )

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

pause()

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

receive_weights() → None

receive_weights(policy_or_weights: TensorDictBase | TensorDictModuleBase | nn.Module | dict, /) → None

receive_weights(*, weights: TensorDictBase | dict) → None

receive_weights(*, policy: TensorDictModuleBase | nn.Module) → None

Receive and apply weights to the collector’s policy.

This method applies weights to the local policy. When receiver schemes are registered, it delegates to those schemes. Otherwise, it directly applies the provided weights.

The method accepts weights in multiple forms for convenience:

Examples

>>> # Receive from registered schemes (distributed collectors)
>>> collector.receive_weights()
>>>
>>> # Apply weights from a policy module (positional)
>>> collector.receive_weights(trained_policy)
>>>
>>> # Apply weights from a TensorDict (positional)
>>> collector.receive_weights(weights_tensordict)
>>>
>>> # Use keyword arguments for clarity
>>> collector.receive_weights(weights=weights_td)
>>> collector.receive_weights(policy=trained_policy)

Parameters:

policy_or_weights –

The weights to apply. Can be:

• nn.Module: A policy module whose weights will be extracted and applied

• TensorDictModuleBase: A TensorDict module whose weights will be extracted

• TensorDictBase: A TensorDict containing weights

• dict: A regular dict containing weights

• None: Receive from registered schemes or mirror from original policy

Keyword Arguments:

• weights – Alternative to positional argument. A TensorDict or dict containing weights to apply. Cannot be used together with policy_or_weights or policy.

• policy – Alternative to positional argument. An nn.Module or TensorDictModuleBase whose weights will be extracted. Cannot be used together with policy_or_weights or weights.

Raises:

ValueError – If conflicting parameters are provided or if arguments are passed when receiver schemes are registered.

register_scheme_receiver(weight_recv_schemes: dict[str, torchrl.weight_update.weight_sync_schemes.WeightSyncScheme], *, synchronize_weights: bool = True)

Set up receiver schemes for this collector to receive weights from parent collectors.

This method initializes receiver schemes and stores them in _receiver_schemes for later use by _receive_weights_scheme() and receive_weights().

Receiver schemes enable cascading weight updates across collector hierarchies: - Parent collector sends weights via its weight_sync_schemes (senders) - Child collector receives weights via its weight_recv_schemes (receivers) - If child is also a parent (intermediate node), it can propagate to its own children

Parameters:

weight_recv_schemes (dict[str, WeightSyncScheme]) – Dictionary of {model_id: WeightSyncScheme} to set up as receivers. These schemes will receive weights from parent collectors.

Keyword Arguments:

synchronize_weights (bool, optional) – If True, synchronize weights immediately after registering the schemes. Defaults to True.

start()

Starts the collector for asynchronous data collection.

This method initiates the background collection of data, allowing for decoupling of data collection and training.

The collected data is typically stored in a replay buffer passed during the collector’s initialization.

Note

After calling this method, it’s essential to shut down the collector using async_shutdown() when you’re done with it to free up resources.

Warning

Asynchronous data collection can significantly impact training performance due to its decoupled nature. Ensure you understand the implications for your specific algorithm before using this mode.

Raises:

NotImplementedError – If not implemented by a subclass.

update_policy_weights_(policy_or_weights: TensorDictBase | TensorDictModuleBase | nn.Module | dict, /) → None

update_policy_weights_(policy_or_weights: TensorDictBase | TensorDictModuleBase | nn.Module | dict, /, *, worker_ids: int | list[int] | torch.device | list[torch.device] | None = None, model_id: str | None = None) → None

update_policy_weights_(*, weights: TensorDictBase | dict, model_id: str | None = None, worker_ids: int | list[int] | torch.device | list[torch.device] | None = None) → None

update_policy_weights_(*, policy: TensorDictModuleBase | nn.Module, model_id: str | None = None, worker_ids: int | list[int] | torch.device | list[torch.device] | None = None) → None

update_policy_weights_(*, weights_dict: dict[str, TensorDictBase | TensorDictModuleBase | nn.Module | dict], worker_ids: int | list[int] | torch.device | list[torch.device] | None = None) → None

Update policy weights for the data collector.

This method synchronizes the policy weights used by the collector with the latest trained weights. It supports both local and remote weight updates, depending on the collector configuration.

The method accepts weights in multiple forms for convenience:

Examples

>>> # Pass policy module as positional argument
>>> collector.update_policy_weights_(policy_module)
>>>
>>> # Pass TensorDict weights as positional argument
>>> collector.update_policy_weights_(weights_tensordict)
>>>
>>> # Use keyword arguments for clarity
>>> collector.update_policy_weights_(weights=weights_td, model_id="actor")
>>> collector.update_policy_weights_(policy=actor_module, model_id="actor")
>>>
>>> # Update multiple models atomically
>>> collector.update_policy_weights_(weights_dict={
...     "actor": actor_weights,
...     "critic": critic_weights,
... })

Parameters:

policy_or_weights –

The weights to update with. Can be:

• nn.Module: A policy module whose weights will be extracted

• TensorDictModuleBase: A TensorDict module whose weights will be extracted

• TensorDictBase: A TensorDict containing weights

• dict: A regular dict containing weights

• None: Will try to get weights from server using _get_server_weights()

Keyword Arguments:

• weights – Alternative to positional argument. A TensorDict or dict containing weights to update. Cannot be used together with policy_or_weights or policy.

• policy – Alternative to positional argument. An nn.Module or TensorDictModuleBase whose weights will be extracted. Cannot be used together with policy_or_weights or weights.

• worker_ids – Identifiers for the workers to update. Relevant when the collector has multiple workers. Can be int, list of ints, device, or list of devices.

• model_id – The model identifier to update (default: "policy"). Cannot be used together with weights_dict.

• weights_dict – Dictionary mapping model_id to weights for updating multiple models atomically. Keys should match model_ids registered in weight_sync_schemes. Cannot be used together with model_id, policy_or_weights, weights, or policy.

Raises:

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

• ValueError – If conflicting parameters are provided.

Note

Users should extend the WeightUpdaterBase classes to customize the weight update logic for specific use cases.

See also

LocalWeightsUpdaterBase and RemoteWeightsUpdaterBase().

property worker_idx: int | None

Get the worker index for this collector.

Returns:

The worker index (0-indexed).

Raises:

RuntimeError – If worker_idx has not been set.