BaseCollector¶

class torchrl.collectors.BaseCollector[source]¶

Base class for data collectors.

async_shutdown(timeout: float | None = None, close_env: bool = True) → None[source]¶

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[source]¶

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

enable_profile(*, workers: list[int] | None = None, num_rollouts: int = 3, warmup_rollouts: int = 1, save_path: str | Path | None = None, activities: list[str] | None = None, record_shapes: bool = True, profile_memory: bool = False, with_stack: bool = True, with_flops: bool = False, on_trace_ready: Callable | None = None) → None[source]¶

Enable profiling for collector worker rollouts.

This method configures the collector to profile rollouts using PyTorch’s profiler. For multi-process collectors, profiling happens in the worker processes. For single-process collectors (Collector), profiling happens in the main process.

Parameters:

workers – List of worker indices to profile. Defaults to [0]. For single-process collectors, this is ignored.
num_rollouts – Total number of rollouts to run the profiler for (including warmup). Profiling stops after this many rollouts. Defaults to 3.
warmup_rollouts – Number of rollouts to skip before starting actual profiling. Useful for JIT/compile warmup. The profiler runs but discards data during warmup. Defaults to 1.
save_path – Path to save the profiling trace. Supports {worker_idx} placeholder for worker-specific files. If None, traces are saved to “./collector_profile_{worker_idx}.json”.
activities – List of profiler activities (“cpu”, “cuda”). Defaults to [“cpu”, “cuda”].
record_shapes – Whether to record tensor shapes. Defaults to True.
profile_memory – Whether to profile memory usage. Defaults to False.
with_stack – Whether to record Python stack traces. Defaults to True.
with_flops – Whether to compute FLOPS. Defaults to False.
on_trace_ready – Optional callback when trace is ready. If None, traces are exported to Chrome trace format at save_path.

Raises:

RuntimeError – If called after iteration has started.
ValueError – If num_rollouts <= warmup_rollouts.

Example

>>> from torchrl.collectors import MultiSyncCollector
>>> collector = MultiSyncCollector(
...     create_env_fn=[make_env] * 4,
...     policy=policy,
...     frames_per_batch=1000,
...     total_frames=100000,
... )
>>> collector.enable_profile(
...     workers=[0],
...     num_rollouts=5,
...     warmup_rollouts=2,
...     save_path="./traces/worker_{worker_idx}.json",
... )
>>> # Worker 0 will be profiled for rollouts 2, 3, 4
>>> for data in collector:
...     train(data)
>>> collector.shutdown()

Note

Profiling adds overhead, so only profile specific workers
The trace file can be viewed in Chrome’s trace viewer (chrome://tracing) or with PyTorch’s TensorBoard plugin
For multi-process collectors, this must be called BEFORE iteration starts as it needs to configure workers

init_updater(*args, **kwargs)[source]¶

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()[source]¶: Context manager that pauses the collector if it is running free.

property profile_config: ProfileConfig | None¶

Get the profiling configuration.

Returns:: ProfileConfig if profiling is enabled, None otherwise.

receive_weights() → None[source]¶

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

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

receive_weights(*, policy: TensorDictModuleBase | 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, WeightSyncScheme], *, synchronize_weights: bool = True)[source]¶

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()[source]¶

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 | Module | dict, /) → None[source]¶

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.

BaseCollector¶

Docs

Tutorials

Resources