SharedMemTransport

class torchrl.weight_update.SharedMemTransport(policy_weights: dict[str, tensordict.base.TensorDictBase] | None = None, auto_register: bool = True)

Shared memory transport for in-place weight updates.

This transport updates shared memory tensors directly without message passing. Workers automatically see weight updates without explicit communication.

The transport supports lazy registration with pipe-based buffer distribution: - On first weight send for a model, creates shared memory and sends buffer via pipes - Workers receive the buffer reference and update their local references - Subsequent updates are pure in-place shared memory (zero-copy)

This hybrid approach solves the chicken-and-egg problem: workers can start before weights are available, and they’ll receive the shared buffer references when ready.

Parameters:

• policy_weights – Dictionary mapping model_id to shared TensorDict weights. Can be empty if using lazy registration.

• auto_register – Whether to automatically register models on first weight send. Default is True. Set to False to require explicit registration via register_weights().

check_ack(message: str = 'updated') → None

No-op for shared memory - no acknowledgment needed.

check_connection() → bool

Shared memory is always ‘connected’.

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

No-op for shared memory - weights are already visible.

register_pipe(pipe: Any) → None

Register a pipe for sending buffer references on first weight send.

Parameters:

pipe – Pipe connection to a worker process.

register_weights(model_id: str, weights: TensorDictBase) → None

Register a shared memory weights TensorDict for a model.

This method allows explicit registration of shared weights. It’s optional when auto_register=True (the default), but required when auto_register=False.

If pipes are registered and this model hasn’t been sent to workers yet, this will trigger sending the buffer reference to all workers. If pipes aren’t registered yet, weights are stored and will be sent when pipes become available (during init_on_sender).

send_ack(message: str = 'updated') → None

No-op for shared memory - no acknowledgment needed.

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

Update weights in-place in shared memory.

If the model is not registered and auto_register=True, it will be automatically registered by creating a shared memory copy of the provided weights. The shared buffer reference is sent to all workers via pipes on first registration, then subsequent updates are pure in-place shared memory.

Parameters:

• model_id – Identifier for the model whose weights to update.

• weights – New weights to send. Can be a TensorDictBase or dict.

Raises:

• KeyError – If model is not registered and auto_register=False.

• ValueError – If weights type is unsupported for auto-registration.