MultiProcessWeightSyncScheme

class torchrl.weight_update.MultiProcessWeightSyncScheme(strategy: Literal['state_dict', 'tensordict'] = 'state_dict')

Weight synchronization for multiprocess operations using pipes.

This scheme creates transports that communicate via multiprocessing pipes.

create_receiver() → WeightReceiver

Create a receiver for this scheme (legacy).

Returns:

WeightReceiver instance configured for this scheme.

create_sender() → WeightSender

Create a sender for this scheme (legacy).

Returns:

WeightSender instance configured for this scheme.

create_transport(pipe: Any) → TransportBackend

Create an MPTransport using the provided pipe (legacy).

get_receiver() → WeightReceiver

Get the receiver instance.

Returns:

Receiver instance for receiving weights in this worker

Raises:

RuntimeError – If init_on_worker() hasn’t been called yet

get_sender() → WeightSender

Get the sender instance.

Returns:

Sender instance for sending weights to workers

Raises:

RuntimeError – If init_on_sender() hasn’t been called yet

init_on_sender(model_id: str, context: Any = None, **kwargs) → None

Initialize on the main process (sender side).

Parameters:

• model_id – Identifier for the model being synchronized

• context – Optional context object providing pipes and num_workers

• **kwargs – Alternative to context (pipes, num_workers, etc.)

init_on_worker(model_id: str, context: Any = None, **kwargs) → None

Initialize on worker process (receiver side).

Parameters:

• model_id – Identifier for the model being synchronized

• context – Optional context object providing pipe and model

• **kwargs – Alternative to context (pipe, model, etc.)

prepare_weights(weights: Any, model_id: str, strategy: WeightStrategy, context: Any = None) → Any

Prepare weights for sending.

This method handles weight extraction, conversion, and any scheme-specific preparation (e.g., cache lookups for SharedMemWeightSyncScheme).

Parameters:

• weights – Raw weights input (can be None, nn.Module, TensorDict, dict, str reference, etc.)

• model_id – The model identifier (e.g., “policy”)

• strategy – WeightStrategy for extracting/converting weights

• context – Optional context (e.g., collector) for model resolution

Returns:

Prepared weights ready to send via transport