RayLLMCollector¶

class torchrl.collectors.llm.RayLLMCollector(env: torchrl.envs.common.EnvBase | collections.abc.Callable[[], torchrl.envs.common.EnvBase], *, policy: collections.abc.Callable[[tensordict.base.TensorDictBase], tensordict.base.TensorDictBase] | None = None, policy_factory: collections.abc.Callable[[], collections.abc.Callable[[tensordict.base.TensorDictBase], tensordict.base.TensorDictBase]] | None = None, dialog_turns_per_batch: int, total_dialog_turns: int = - 1, yield_only_last_steps: bool | None = None, yield_completed_trajectories: bool | None = None, postproc: collections.abc.Callable[[tensordict.base.TensorDictBase], tensordict.base.TensorDictBase] | None = None, async_envs: bool | None = None, replay_buffer: torchrl.data.replay_buffers.replay_buffers.ReplayBuffer | None = None, reset_at_each_iter: bool = False, flatten_data: bool | None = None, weight_updater: torchrl.collectors.weight_update.WeightUpdaterBase | collections.abc.Callable[[], torchrl.collectors.weight_update.WeightUpdaterBase] | None = None, ray_init_config: dict[str, Any] | None = None, remote_config: dict[str, Any] | None = None, track_policy_version: bool | torchrl.envs.llm.transforms.policy_version.PolicyVersion = False, sync_iter: bool = True, verbose: bool = False, num_cpus: int | None = None, num_gpus: int | None = None)[source]¶

A lightweight Ray implementation of the LLM Collector that can be extended and sampled remotely.

Parameters:

env (EnvBase or EnvBase constructor) – the environment to be used for data collection.

Keyword Arguments:

policy (Callable[[TensorDictBase], TensorDictBase]) – the policy to be used for data collection.
policy_factory (Callable[[], Callable], optional) – a callable that returns a policy instance. This is exclusive with the policy argument.
dialog_turns_per_batch (int) – A keyword-only argument representing the total number of elements in a batch.
total_dialog_turns (int) – A keyword-only argument representing the total number of dialog turns returned by the collector during its lifespan.
yield_only_last_steps (bool, optional) – whether to yield every step of a trajectory, or only the last (done) steps.
yield_completed_trajectories (bool, optional) – whether to yield batches of rollouts with a given number of steps or single, completed trajectories.
postproc (Callable, optional) – A post-processing transform.
async_envs (bool, optional) – if True, the environment will be run asynchronously.
replay_buffer (ReplayBuffer, optional) – if provided, the collector will not yield tensordicts but populate the buffer instead.
reset_at_each_iter (bool, optional) – if True, the environment will be reset at each iteration.
flatten_data (bool, optional) – if True, the collector will flatten the collected data before returning it.
weight_updater (WeightUpdaterBase or constructor, optional) – An instance of WeightUpdaterBase or its subclass, responsible for updating the policy weights on remote inference workers.
ray_init_config (dict[str, Any], optional) – keyword arguments to pass to ray.init().
remote_config (dict[str, Any], optional) – keyword arguments to pass to cls.as_remote().
num_cpus (int, optional) – Number of CPUs of the actor. Defaults to None (taken from remote_config).
num_gpus (int, optional) – Number of GPUs of the actor. Defaults to None (taken from remote_config).
sync_iter (bool, optional) –
if True, items yeilded by the collector will be synced to the local process. If False, the collector will collect the next batch of data in between yielding. This has no effect when data is collected through the start() method. For example:
```
>>> collector = RayLLMCollector(..., sync_iter=True)
>>> for data in collector:  # blocking
...     # expensive operation - collector is idle
>>> collector = RayLLMCollector(..., sync_iter=False)
>>> for data in collector:  # non-blocking
...     # expensive operation - collector is collecting data
```
This is somehwat equivalent to using MultiSyncCollector (sync_iter=True) or MultiAsyncDataCollector (sync_iter=False). Defaults to True.
verbose (bool, optional) – if True, the collector will print progress information. Defaults to False.

classmethod as_remote(remote_config: dict[str, Any] | None = None)¶

Creates an instance of a remote ray class.

Parameters:

cls (Python Class) – class to be remotely instantiated.
remote_config (dict) – the quantity of CPU cores to reserve for this class.

Returns:

A function that creates ray remote class instances.

async_shutdown(timeout=None)[source]¶: Shuts down the collector asynchronously.

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

property dialog_turns_per_batch: int¶: Number of dialog turns per batch.

get_model(model_id: str)¶

Get model instance by ID (for weight sync schemes).

Parameters:: model_id – Model identifier (e.g., “policy”, “value_net”)
Returns:: The model instance
Raises:: ValueError – If model_id is not recognized

get_policy_model()¶

Get the policy model.

This method is used by RayLLMCollector to get the remote LLM instance for weight updates.

Returns:: The policy model instance

get_policy_version() → str | int | None¶

Get the current policy version.

This method exists to support remote calls in Ray actors, since properties cannot be accessed directly through Ray’s RPC mechanism.

Returns:: The current version number (int) or UUID (str), or None if version tracking is disabled.

getattr_env(attr)¶: Get an attribute from the environment.

getattr_policy(attr)¶: Get an attribute from the policy.

getattr_rb(attr)¶: Get an attribute from the replay buffer.

increment_version()[source]¶: Increment the policy version.

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

Initialize the weight updater with custom arguments.

This method calls init_updater on the remote collector.

Parameters:

*args – Positional arguments for weight updater initialization
**kwargs – Keyword arguments for weight updater initialization

is_initialized() → bool¶

Check if the collector is initialized and ready.

Returns:: True if the collector is initialized and ready to collect data.
Return type:: bool

iterator() → Iterator[TensorDictBase]¶

Iterates through the DataCollector.

Yields: TensorDictBase objects containing (chunks of) trajectories

load_state_dict(state_dict: OrderedDict, **kwargs) → None¶

Loads a state_dict on the environment and policy.

Parameters:: state_dict (OrderedDict) – ordered dictionary containing the fields “policy_state_dict” and "env_state_dict".

next() → None[source]¶

Get the next batch of data from the collector.

Returns:: None as the data is written directly to the replay buffer.

pause()¶: Context manager that pauses the collector if it is running free.

property policy_version: str | int | None¶

The current version of the policy.

Returns:: The current version number (int) or UUID (str), or None if version tracking is disabled.

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.

reset(index=None, **kwargs) → None¶: Resets the environments to a new initial state.

property rollout: Callable[[], TensorDictBase]¶: Returns the rollout function.

set_seed(seed: int, static_seed: bool = False) → int¶

Sets the seeds of the environments stored in the DataCollector.

Parameters:

seed (int) – 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 = ParallelEnv(6, env_fn)
>>> policy = TensorDictModule(nn.Linear(3, 1), in_keys=["observation"], out_keys=["action"])
>>> collector = Collector(env_fn_parallel, policy, total_frames=300, frames_per_batch=100)
>>> out_seed = collector.set_seed(1)  # out_seed = 6

shutdown()[source]¶: Shuts down the collector.

start()[source]¶: Starts the collector in a background thread.

state_dict() → OrderedDict¶

Returns the local state_dict of the data collector (environment and policy).

Returns:: an ordered dictionary with fields "policy_state_dict" and “env_state_dict”.

property total_dialog_turns¶: Total number of dialog turns to collect.

Updates the policy weights on remote workers.

Parameters:

policy_or_weights – 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 – The workers to update. If None, updates all workers.

property weight_updater: WeightUpdaterBase¶

The weight updater instance.

We can pass the weight updater because it’s stateless, hence serializable.

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.

RayLLMCollector¶

Docs

Tutorials

Resources