fix the fixes:

Author: vmoens

torchrl/collectors/collectors.py

@@ -2758,13 +2758,16 @@ def _setup_multi_policy_and_weights(
 
         if weight_sync_schemes is not None:
             # Weight sync schemes handle all weight distribution
-            # No need to extract weights or create stateful policies here
+            # Extract weights so schemes can access them, but don't do in-place replacement
             self._policy_weights_dict = {}
-            self._fallback_policy = policy
-            self._get_weights_fn = None
+            self._fallback_policy = None
+
+            if not any(policy_factory) and policy is not None:
+                # Extract weights for the first device so schemes can access them
+                # Use first device as representative
+                first_device = self.policy_device[0] if self.policy_device else None
 
-            # Validate device types for SharedMemWeightSyncScheme
-            if not any(policy_factory):
+                # Validate device types for SharedMemWeightSyncScheme
                 for scheme in weight_sync_schemes.values():
                     if isinstance(scheme, SharedMemWeightSyncScheme):
                         for policy_device in self.policy_device:
@@ -2776,6 +2779,29 @@ def _setup_multi_policy_and_weights(
                                     f"Device type '{policy_device.type}' not supported for SharedMemWeightSyncScheme. "
                                     f"Only 'cpu' and 'cuda' are supported."
                                 )
+
+                # Extract weights from policy
+                weights = (
+                    TensorDict.from_module(policy)
+                    if isinstance(policy, nn.Module)
+                    else TensorDict()
+                )
+
+                # For SharedMemWeightSyncScheme, share the weights
+                if any(
+                    isinstance(scheme, SharedMemWeightSyncScheme)
+                    for scheme in weight_sync_schemes.values()
+                ):
+                    if first_device and first_device.type == "cpu":
+                        weights = weights.share_memory_()
+                    elif first_device and first_device.type == "cuda":
+                        # CUDA tensors maintain shared references through mp.Queue
+                        weights = weights.to(first_device).share_memory_()
+
+                self._policy_weights_dict[first_device] = weights
+                self._fallback_policy = policy
+
+            self._get_weights_fn = None
         else:
             # Using legacy weight updater - extract weights and create stateful policies
             self._setup_multi_policy_and_weights_legacy(
@@ -3067,21 +3093,24 @@ def _run_processes(self) -> None:
                 1, torch.get_num_threads() - total_workers
             )  # 1 more thread for this proc
 
-        # Initialize weight sync schemes to create queues before workers start
+        # Set up for worker processes
         torch.set_num_threads(self.num_threads)
         queue_out = mp.Queue(self._queue_len)  # sends data from proc to main
         self.procs = []
         self.pipes = []
         self._traj_pool = _TrajectoryPool(lock=True)
 
-        # Initialize all weight sync schemes early
-        # Schemes own their queues and handle distribution internally
+        # Initialize weight sync schemes early for SharedMemWeightSyncScheme
+        # (queue created in __init__ will be pickled with scheme to workers)
+        # For MultiProcessWeightSyncScheme, we'll initialize after pipes are available
         if self._weight_sync_schemes:
             for model_id, scheme in self._weight_sync_schemes.items():
-                # Check if scheme has new API
-                if hasattr(scheme, "init_on_sender"):
+                # Only initialize SharedMemWeightSyncScheme now (needs queue before workers)
+                # MultiProcessWeightSyncScheme will be initialized after workers are created
+                if isinstance(scheme, SharedMemWeightSyncScheme) and hasattr(
+                    scheme, "init_on_sender"
+                ):
                     scheme.init_on_sender(model_id=model_id, context=self)
-                    # Get the initialized sender
                     self._weight_senders[model_id] = scheme.get_sender()
 
         # Create a policy on the right device
@@ -3257,6 +3286,18 @@ def _run_processes(self) -> None:
                     # Legacy string error message
                     raise RuntimeError(msg)
 
+        # Initialize MultiProcessWeightSyncScheme now that workers are ready and pipes are available
+        # (SharedMemWeightSyncScheme was already initialized before workers)
+        if self._weight_sync_schemes:
+            for model_id, scheme in self._weight_sync_schemes.items():
+                # Only initialize non-SharedMem schemes here (need pipes)
+                if not isinstance(scheme, SharedMemWeightSyncScheme) and hasattr(
+                    scheme, "init_on_sender"
+                ):
+                    scheme.init_on_sender(model_id=model_id, context=self)
+                    # Get the initialized sender
+                    self._weight_senders[model_id] = scheme.get_sender()
+
         self.queue_out = queue_out
         self.closed = False
 

torchrl/weight_update/weight_sync_schemes.py

@@ -5,14 +5,19 @@
 from __future__ import annotations
 
 import abc
+import time
 
 import weakref
 from collections.abc import Iterator
+from queue import Empty
 from typing import Any, Literal, Protocol
 
+import torch
+import torch.distributed
+
 from tensordict import TensorDict, TensorDictBase
 
-from torch import nn
+from torch import multiprocessing as mp, nn
 
 __all__ = [
     "TransportBackend",
@@ -195,8 +200,6 @@ def _infer_device(self, td: TensorDictBase):
         Returns:
             torch.device or None if no tensors found or all on different devices.
         """
-        import torch
-
         for value in td.values(True, True):
             if isinstance(value, torch.Tensor):
                 return value.device
@@ -688,8 +691,6 @@ def send_ack(self, message: str = "updated") -> None:
 
     def check_connection(self) -> bool:
         """Check if torch.distributed is initialized."""
-        import torch.distributed
-
         return torch.distributed.is_initialized()
 
 
@@ -1602,7 +1603,8 @@ def __init__(
         self._shared_transport = SharedMemTransport(
             self.policy_weights, auto_register=auto_register
         )
-        self._weight_init_queue = None  # Created during init_on_sender
+        # Create queue immediately so it's available when scheme is pickled to workers
+        self._weight_init_queue = mp.Queue()
 
     def register_shared_weights(self, model_id: str, weights: TensorDictBase) -> None:
         """Register shared memory weights for a model.
@@ -1659,13 +1661,8 @@ def init_on_sender(
                 "device_to_workers mapping must be provided via context or kwargs"
             )
 
-        # Create queue once for this scheme instance (owned by scheme, not collector)
-        if self._weight_init_queue is None:
-            from torch import multiprocessing as mp
-
-            self._weight_init_queue = mp.Queue()
-
         # Set worker info in transport
+        # Queue was already created in __init__ so it's available to workers
         self._shared_transport.set_worker_info(device_to_workers)
         self._shared_transport._weight_queue = self._weight_init_queue
 
@@ -1723,8 +1720,6 @@ def init_on_worker(
         # Receive weights from the scheme's queue if available
         if self._weight_init_queue is not None and worker_idx is not None:
             # Read from queue until we find our worker_idx and model_id
-            from queue import Empty
-
             timeout = kwargs.get("timeout", 10.0)
             try:
                 while True:
@@ -1746,8 +1741,6 @@ def init_on_worker(
                             (msg_worker_idx, msg_model_id, shared_weights)
                         )
                         # Small sleep to avoid immediately picking up the same message
-                        import time
-
                         time.sleep(0.001)
             except Empty:
                 # No weights pre-registered for this model (will use auto-register or policy_factory)