RemoteWeightUpdaterBase

class torchrl.collectors.RemoteWeightUpdaterBase

A base class for updating remote policy weights on inference workers.

This class provides an interface for uploading and synchronizing the weights of a policy across remote inference workers. The update process is centralized, meaning the server is responsible for distributing the weights to the inference nodes.

To extend this class, implement the following abstract methods:

• _sync_weights_with_worker: Define how to synchronize weights with a specific worker.

• _get_server_weights: Define how to retrieve the weights from the server.

• _maybe_map_weights: Optionally transform the server weights before distribution.

• all_worker_ids: Provide a list of all worker identifiers.

Variables:

policy (Policy, optional) – The policy whose weights are to be updated.

update_weights()

Updates the weights on specified or all remote workers.

register_collector()

Registers a collector. This should be called automatically by the collector upon registration of the updater.

See also

LocalWeightsUpdaterBase and update_policy_weights_().

register_collector(collector: DataCollectorBase)

Register a collector in the updater.

Once registered, the updater will not accept another collector.

Parameters:

collector (DataCollectorBase) – The collector to register.