DistributedTransport

class torchrl.weight_update.DistributedTransport(store=None, rank=None, sync=True)

torch.distributed transport for communicating with a single distributed worker.

This transport handles weight updates for ONE specific distributed worker via torch.distributed send/recv. Multiple transports are created for multiple workers, following the same pattern as multiprocess collectors.

check_connection() → bool

Check if torch.distributed is initialized.

receive_weights(timeout: float = 1.0) → tuple[str, Any] | None

Distributed workers receive weights through torch.distributed primitives.

send_weights(model_id: str, weights: Any) → None

Send weights to the distributed worker.

Note: We don’t pass model_id to the remote collector because remote collectors don’t have weight senders - they apply weights directly to their local policy.

send_weights_async(model_id: str, weights: Any) → None

Send weights to distributed worker without waiting for acknowledgment.

Use wait_ack() to wait for acknowledgment after sending to all workers.

wait_ack() → None

Wait for acknowledgment from distributed worker.