final?

Author: vmoens

docs/source/reference/collectors_weightsync.rst

@@ -198,6 +198,9 @@ Weight Senders
     :template: rl_template.rst
 
     WeightSender
+    MPWeightSender
+    RPCWeightSender
+    DistributedWeightSender
     RayModuleTransformSender
 
 Weight Receivers
@@ -208,6 +211,9 @@ Weight Receivers
     :template: rl_template.rst
 
     WeightReceiver
+    MPWeightReceiver
+    RPCWeightReceiver
+    DistributedWeightReceiver
     RayModuleTransformReceiver
 
 Transports

examples/collectors/multi_weight_updates.py

@@ -25,7 +25,7 @@
 from torchrl.data import LazyTensorStorage, ReplayBuffer
 from torchrl.envs.libs.gym import GymEnv
 from torchrl.envs.transforms.module import ModuleTransform
-from torchrl.weight_update.weight_sync_schemes import MultiProcessWeightSyncScheme
+from torchrl.weight_update import MultiProcessWeightSyncScheme
 
 
 def make_module():

test/test_collector.py

@@ -1558,8 +1558,6 @@ def create_env():
     )  # MultiSync has known indexing issues with SharedMem
     def test_update_weights_shared_mem(self, use_async):
         """Test shared memory weight synchronization scheme."""
-        from tensordict import TensorDict
-        from torchrl.weight_update.weight_sync_schemes import SharedMemWeightSyncScheme
 
         def create_env():
             return ContinuousActionVecMockEnv()
@@ -4117,16 +4115,17 @@ def test_start_update_policy(self, total_frames, cls, weight_sync_scheme):
             frames_per_batch=16,
             **kwargs,
         )
-        if not isinstance(collector, SyncDataCollector):
-            if weight_sync_scheme is not None:
-                assert isinstance(
-                    collector._weight_sync_schemes["policy"], weight_sync_scheme
-                )
-            else:
-                assert isinstance(
-                    collector._weight_sync_schemes["policy"], SharedMemWeightSyncScheme
-                )
         try:
+            if not isinstance(collector, SyncDataCollector):
+                if weight_sync_scheme is not None:
+                    assert isinstance(
+                        collector._weight_sync_schemes["policy"], weight_sync_scheme
+                    )
+                else:
+                    assert isinstance(
+                        collector._weight_sync_schemes["policy"],
+                        SharedMemWeightSyncScheme,
+                    )
             collector.start()
             for _ in range(10):
                 time.sleep(0.1)

test/test_weightsync.py

@@ -17,8 +17,7 @@
 from tensordict.nn import TensorDictModule
 from torch import multiprocessing as mp
 from torchrl.collectors import MultiSyncDataCollector, SyncDataCollector
-from torchrl.weight_update.weight_sync_schemes import (
-    _resolve_model,
+from torchrl.weight_update import (
     DistributedWeightSyncScheme,
     MPTransport,
     MultiProcessWeightSyncScheme,
@@ -27,6 +26,9 @@
     RayWeightSyncScheme,
     RPCWeightSyncScheme,
     SharedMemTransport,
+)
+from torchrl.weight_update.utils import _resolve_model
+from torchrl.weight_update.weight_sync_schemes import (
     SharedMemWeightSyncScheme,
     WeightStrategy,
 )

torchrl/collectors/_multi_base.py

@@ -334,7 +334,8 @@ def __init__(
         policy_factory = self._setup_policy_factory(policy_factory)
 
         # Set up weight synchronization
-        weight_sync_schemes = {}
+        if weight_sync_schemes is None:
+            weight_sync_schemes = {}
         if (
             not any(policy_factory)
             and not weight_sync_schemes
@@ -516,13 +517,13 @@ def _setup_multi_policy_and_weights(
             weight_sync_policy = weight_sync_schemes.get("policy")
             if weight_sync_policy is None:
                 return
-            if weight_sync_policy._initialized_on_sender:
-                return
             if any(p is not None for p in policy_factory):
-                raise RuntimeError(
-                    f"the weight sync scheme must be initialized on sender ahead of time when passing a policy factory. Got {policy_factory=}"
-                )
-            weight_sync_policy.init_on_sender(model=policy, devices=self.policy_device)
+                if not weight_sync_policy._initialized_on_sender:
+                    raise RuntimeError(
+                        f"the weight sync scheme must be initialized on sender ahead of time when passing a policy factory. Got {policy_factory=}"
+                    )
+            # Weight sync scheme initialization happens in _run_processes
+            # where pipes and workers are available
         else:
             # Using legacy weight updater - extract weights and create stateful policies
             self._setup_multi_policy_and_weights_legacy(
@@ -821,19 +822,20 @@ def _run_processes(self) -> None:
         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 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
+        # Create all pipes upfront (needed for weight sync scheme initialization)
+        # Store as list of (parent, child) tuples for use in worker creation
+        pipe_pairs = [mp.Pipe() for _ in range(self.num_workers)]
+        # Extract parent pipes for external use (e.g., polling, receiving messages)
+        self.pipes = [pipe_parent for pipe_parent, _ in pipe_pairs]
+
+        # Initialize all weight sync schemes now that pipes are available
+        # Both SharedMemWeightSyncScheme (uses queues) and MultiProcessWeightSyncScheme (uses pipes)
+        # can be initialized here since all required resources exist
         if self._weight_sync_schemes:
             for model_id, scheme in self._weight_sync_schemes.items():
-                # 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"
-                ):
+                if hasattr(scheme, "init_on_sender"):
                     scheme.init_on_sender(model_id=model_id, context=self)
                     self._weight_senders[model_id] = scheme.get_sender()
 
@@ -848,7 +850,7 @@ def _run_processes(self) -> None:
         for i, (env_fun, env_fun_kwargs) in enumerate(
             zip(self.create_env_fn, self.create_env_kwargs)
         ):
-            pipe_parent, pipe_child = mp.Pipe()  # send messages to procs
+            pipe_parent, pipe_child = pipe_pairs[i]  # use pre-created pipes
             if env_fun.__class__.__name__ != "EnvCreator" and not isinstance(
                 env_fun, EnvBase
             ):  # to avoid circular imports
@@ -966,7 +968,6 @@ def _run_processes(self) -> None:
                         ) from err
                 pipe_child.close()
                 self.procs.append(proc)
-                self.pipes.append(pipe_parent)
 
         # Synchronize initial weights with workers AFTER starting processes but BEFORE waiting for "instantiated"
         # This must happen after proc.start() but before workers send "instantiated" to avoid deadlock:
@@ -1027,18 +1028,6 @@ 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/collectors/_runner.py

@@ -30,7 +30,7 @@
 
 
 def _make_policy_factory(
-    *, policy: Callable, policy_factory, weight_sync_scheme, worker_idx
+    *, policy: Callable, policy_factory, weight_sync_scheme, worker_idx, pipe=None
 ):
     if policy is not None and policy_factory is not None:
         raise ValueError("policy cannot be used with policy_factory")
@@ -40,7 +40,7 @@ def _make_policy_factory(
     if weight_sync_scheme is not None:
         # Initialize the receiver on the worker side
         weight_sync_scheme.init_on_worker(
-            model=policy, model_id="policy", worker_idx=worker_idx
+            model=policy, model_id="policy", worker_idx=worker_idx, pipe=pipe
         )
         # Get the receiver and synchronize initial weights
         receiver = weight_sync_scheme.get_receiver()
@@ -92,8 +92,11 @@ def _main_async_collector(
         _make_policy_factory,
         policy=policy,
         policy_factory=policy_factory,
-        weight_sync_scheme=weight_sync_schemes.get("policy"),
+        weight_sync_scheme=weight_sync_schemes.get("policy")
+        if weight_sync_schemes
+        else None,
         worker_idx=worker_idx,
+        pipe=pipe_child,
     )
     policy = None
     try:

torchrl/collectors/distributed/generic.py

@@ -570,9 +570,7 @@ def __init__(
         # Set up weight synchronization - prefer new schemes over legacy updater
         if weight_updater is None and weight_sync_schemes is None:
             # Default to Distributed weight sync scheme for distributed collectors
-            from torchrl.weight_update.weight_sync_schemes import (
-                DistributedWeightSyncScheme,
-            )
+            from torchrl.weight_update import DistributedWeightSyncScheme
 
             weight_sync_schemes = {
                 "policy": DistributedWeightSyncScheme(backend=backend, sync=self._sync)

torchrl/collectors/distributed/ray.py

@@ -539,7 +539,7 @@ def check_list_length_consistency(*lists):
         # Set up weight synchronization - prefer new schemes over legacy updater
         if weight_updater is None and weight_sync_schemes is None:
             # Default to Ray weight sync scheme for Ray collectors
-            from torchrl.weight_update.weight_sync_schemes import RayWeightSyncScheme
+            from torchrl.weight_update import RayWeightSyncScheme
 
             weight_sync_schemes = {"policy": RayWeightSyncScheme()}
 

torchrl/collectors/distributed/rpc.py

@@ -417,7 +417,7 @@ def __init__(
         # Set up weight synchronization - prefer new schemes over legacy updater
         if weight_updater is None and weight_sync_schemes is None:
             # Default to RPC weight sync scheme for RPC collectors
-            from torchrl.weight_update.weight_sync_schemes import RPCWeightSyncScheme
+            from torchrl.weight_update import RPCWeightSyncScheme
 
             weight_sync_schemes = {"policy": RPCWeightSyncScheme()}
 

torchrl/weight_update/__init__.py

@@ -3,22 +3,30 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .weight_sync_schemes import (
+from ._distributed import (
     DistributedTransport,
+    DistributedWeightReceiver,
+    DistributedWeightSender,
     DistributedWeightSyncScheme,
+)
+from ._mp import (
     MPTransport,
+    MPWeightReceiver,
+    MPWeightSender,
     MultiProcessWeightSyncScheme,
-    NoWeightSyncScheme,
+)
+from ._noupdate import NoWeightSyncScheme
+from ._ray import (
     RayActorTransport,
     RayModuleTransformReceiver,
     RayModuleTransformScheme,
     RayModuleTransformSender,
     RayTransport,
     RayWeightSyncScheme,
-    RPCTransport,
-    RPCWeightSyncScheme,
-    SharedMemTransport,
-    SharedMemWeightSyncScheme,
+)
+from ._rpc import RPCTransport, RPCWeightReceiver, RPCWeightSender, RPCWeightSyncScheme
+from ._shared import SharedMemTransport, SharedMemWeightSyncScheme
+from .weight_sync_schemes import (
     TransportBackend,
     WeightReceiver,
     WeightSender,
@@ -27,19 +35,30 @@
 )
 
 __all__ = [
+    # Base classes
     "TransportBackend",
+    "WeightStrategy",
+    "WeightSender",
+    "WeightReceiver",
+    "WeightSyncScheme",
+    # Transports
     "MPTransport",
     "SharedMemTransport",
     "RayTransport",
     "RayActorTransport",
     "RPCTransport",
     "DistributedTransport",
-    "WeightStrategy",
-    "WeightSender",
-    "WeightReceiver",
+    # Senders
+    "MPWeightSender",
+    "RPCWeightSender",
+    "DistributedWeightSender",
     "RayModuleTransformSender",
+    # Receivers
+    "MPWeightReceiver",
+    "RPCWeightReceiver",
+    "DistributedWeightReceiver",
     "RayModuleTransformReceiver",
-    "WeightSyncScheme",
+    # Schemes
     "MultiProcessWeightSyncScheme",
     "SharedMemWeightSyncScheme",
     "NoWeightSyncScheme",