torchrl 0.11.0__cp314-cp314t-win_amd64.whl

torchrl/collectors/distributed/rpc.py

@@ -0,0 +1,1006 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+r"""Generic distributed data-collector using torch.distributed.rpc backend."""
+from __future__ import annotations
+
+import collections
+import os
+import socket
+import time
+import warnings
+from collections import OrderedDict
+from collections.abc import Callable, Sequence
+from copy import copy, deepcopy
+from typing import Any
+
+import torch.cuda
+
+from tensordict import TensorDict, TensorDictBase
+from torch import nn
+
+from torch.distributed import rpc
+from torchrl._utils import _ProcessNoWarn, logger as torchrl_logger, VERBOSE
+from torchrl.collectors._base import _LegacyCollectorMeta, BaseCollector
+
+from torchrl.collectors._constants import DEFAULT_EXPLORATION_TYPE
+from torchrl.collectors._multi_async import MultiAsyncCollector
+from torchrl.collectors._multi_sync import MultiSyncCollector
+from torchrl.collectors._single import Collector
+from torchrl.collectors.distributed.default_configs import (
+    DEFAULT_SLURM_CONF,
+    DEFAULT_TENSORPIPE_OPTIONS,
+    IDLE_TIMEOUT,
+    TCP_PORT,
+)
+from torchrl.collectors.utils import _NON_NN_POLICY_WEIGHTS, split_trajectories
+from torchrl.collectors.weight_update import WeightUpdaterBase
+from torchrl.data.utils import CloudpickleWrapper
+from torchrl.envs.common import EnvBase
+from torchrl.envs.env_creator import EnvCreator
+from torchrl.weight_update.weight_sync_schemes import WeightSyncScheme
+
+SUBMITIT_ERR = None
+try:
+    import submitit
+
+    _has_submitit = True
+except ModuleNotFoundError as err:
+    _has_submitit = False
+    SUBMITIT_ERR = err
+
+
+def _rpc_init_collection_node(
+    rank,
+    rank0_ip,
+    tcp_port,
+    world_size,
+    visible_device,
+    tensorpipe_options,
+    backend="gloo",
+    verbose=VERBOSE,
+):
+    os.environ["MASTER_ADDR"] = str(rank0_ip)
+    os.environ["MASTER_PORT"] = str(tcp_port)
+
+    # Initialize torch.distributed process group for efficient weight transfer
+    if verbose:
+        torchrl_logger.debug(
+            f"init distributed with rank={rank}, world_size={world_size}, backend={backend}"
+        )
+    torch.distributed.init_process_group(
+        backend=backend,
+        rank=rank,
+        world_size=world_size,
+    )
+
+    if isinstance(visible_device, list):
+        pass
+    elif isinstance(visible_device, (str, int, torch.device)):
+        visible_device = [visible_device]
+    elif visible_device is None:
+        pass
+    else:
+        raise RuntimeError(f"unrecognised dtype {type(visible_device)}")
+
+    options = rpc.TensorPipeRpcBackendOptions(
+        devices=visible_device,
+        **tensorpipe_options,
+    )
+    if verbose:
+        torchrl_logger.debug(
+            f"init rpc with master addr: {os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}"
+        )
+    rpc.init_rpc(
+        f"COLLECTOR_NODE_{rank}",
+        rank=rank,
+        backend=rpc.BackendType.TENSORPIPE,
+        rpc_backend_options=options,
+        world_size=world_size,
+    )
+    rpc.shutdown()
+    torch.distributed.destroy_process_group()
+
+
+class RPCCollector(BaseCollector):
+    """An RPC-based distributed data collector.
+
+    Supports sync and async data collection.
+
+    Args:
+        create_env_fn (Callable or List[Callabled]): list of Callables, each returning an
+            instance of :class:`~torchrl.envs.EnvBase`.
+        policy (Callable): Policy to be executed in the environment.
+            Must accept :class:`tensordict.tensordict.TensorDictBase` object as input.
+            If ``None`` is provided, the policy used will be a
+            :class:`~torchrl.collectors.RandomPolicy` instance with the environment
+            ``action_spec``.
+            Accepted policies are usually subclasses of :class:`~tensordict.nn.TensorDictModuleBase`.
+            This is the recommended usage of the collector.
+            Other callables are accepted too:
+            If the policy is not a ``TensorDictModuleBase`` (e.g., a regular :class:`~torch.nn.Module`
+            instances) it will be wrapped in a `nn.Module` first.
+            Then, the collector will try to assess if these
+            modules require wrapping in a :class:`~tensordict.nn.TensorDictModule` or not.
+
+            - If the policy forward signature matches any of ``forward(self, tensordict)``,
+              ``forward(self, td)`` or ``forward(self, <anything>: TensorDictBase)`` (or
+              any typing with a single argument typed as a subclass of ``TensorDictBase``)
+              then the policy won't be wrapped in a :class:`~tensordict.nn.TensorDictModule`.
+
+            - In all other cases an attempt to wrap it will be undergone as such: ``TensorDictModule(policy, in_keys=env_obs_key, out_keys=env.action_keys)``.
+
+            .. note:: If the policy needs to be passed as a policy factory (e.g., in case it mustn't be serialized /
+                pickled directly), the ``policy_factory`` should be used instead.
+
+    Keyword Args:
+        policy_factory (Callable[[], Callable], list of Callable[[], Callable], optional): a callable
+            (or list of callables) that returns a policy instance. This is exclusive with the `policy` argument.
+
+            .. note:: `policy_factory` comes in handy whenever the policy cannot be serialized.
+
+        frames_per_batch (int): A keyword-only argument representing the total
+            number of elements in a batch.
+        total_frames (int): A keyword-only argument representing the total
+            number of frames returned by the collector
+            during its lifespan. If the ``total_frames`` is not divisible by
+            ``frames_per_batch``, an exception is raised.
+            Endless collectors can be created by passing ``total_frames=-1``.
+            Defaults to ``-1`` (endless collector).
+        device (int, str or torch.device, optional): The generic device of the
+            collector. The ``device`` args fills any non-specified device: if
+            ``device`` is not ``None`` and any of ``storing_device``, ``policy_device`` or
+            ``env_device`` is not specified, its value will be set to ``device``.
+            Defaults to ``None`` (No default device).
+            Lists of devices are supported.
+        storing_device (int, str or torch.device, optional): The *remote* device on which
+            the output :class:`~tensordict.TensorDict` will be stored.
+            If ``device`` is passed and ``storing_device`` is ``None``, it will
+            default to the value indicated by ``device``.
+            For long trajectories, it may be necessary to store the data on a different
+            device than the one where the policy and env are executed.
+            Defaults to ``None`` (the output tensordict isn't on a specific device,
+            leaf tensors sit on the device where they were created).
+            Lists of devices are supported.
+        env_device (int, str or torch.device, optional): The *remote* device on which
+            the environment should be cast (or executed if that functionality is
+            supported). If not specified and the env has a non-``None`` device,
+            ``env_device`` will default to that value. If ``device`` is passed
+            and ``env_device=None``, it will default to ``device``. If the value
+            as such specified of ``env_device`` differs from ``policy_device``
+            and one of them is not ``None``, the data will be cast to ``env_device``
+            before being passed to the env (i.e., passing different devices to
+            policy and env is supported). Defaults to ``None``.
+            Lists of devices are supported.
+        policy_device (int, str or torch.device, optional): The *remote* device on which
+            the policy should be cast.
+            If ``device`` is passed and ``policy_device=None``, it will default
+            to ``device``. If the value as such specified of ``policy_device``
+            differs from ``env_device`` and one of them is not ``None``,
+            the data will be cast to ``policy_device`` before being passed to
+            the policy (i.e., passing different devices to policy and env is
+            supported). Defaults to ``None``.
+            Lists of devices are supported.
+        max_frames_per_traj (int, optional): Maximum steps per trajectory.
+            Note that a trajectory can span across multiple batches (unless
+            ``reset_at_each_iter`` is set to ``True``, see below).
+            Once a trajectory reaches ``n_steps``, the environment is reset.
+            If the environment wraps multiple environments together, the number
+            of steps is tracked for each environment independently. Negative
+            values are allowed, in which case this argument is ignored.
+            Defaults to ``None`` (i.e., no maximum number of steps).
+        init_random_frames (int, optional): Number of frames for which the
+            policy is ignored before it is called. This feature is mainly
+            intended to be used in offline/model-based settings, where a
+            batch of random trajectories can be used to initialize training.
+            If provided, it will be rounded up to the closest multiple of frames_per_batch.
+            Defaults to ``None`` (i.e. no random frames).
+        reset_at_each_iter (bool, optional): Whether environments should be reset
+            at the beginning of a batch collection.
+            Defaults to ``False``.
+        postproc (Callable, optional): A post-processing transform, such as
+            a :class:`~torchrl.envs.Transform` or a :class:`~torchrl.data.postprocs.MultiStep`
+            instance.
+            Defaults to ``None``.
+        split_trajs (bool, optional): Boolean indicating whether the resulting
+            TensorDict should be split according to the trajectories.
+            See :func:`~torchrl.collectors.utils.split_trajectories` for more
+            information.
+            Defaults to ``False``.
+        exploration_type (ExplorationType, optional): interaction mode to be used when
+            collecting data. Must be one of ``torchrl.envs.utils.ExplorationType.DETERMINISTIC``,
+            ``torchrl.envs.utils.ExplorationType.RANDOM``, ``torchrl.envs.utils.ExplorationType.MODE``
+            or ``torchrl.envs.utils.ExplorationType.MEAN``.
+            Defaults to ``torchrl.envs.utils.ExplorationType.RANDOM``.
+        collector_class (Type or str, optional): a collector class for the remote node. Can be
+            :class:`~torchrl.collectors.Collector`,
+            :class:`~torchrl.collectors.MultiSyncCollector`,
+            :class:`~torchrl.collectors.MultiAsyncCollector`
+            or a derived class of these. The strings "single", "sync" and
+            "async" correspond to respective class.
+            Defaults to :class:`~torchrl.collectors.Collector`.
+
+            .. note::
+
+              Support for :class:`MultiSyncCollector` and :class:`MultiAsyncCollector`
+              is experimental, and :class:`~torchrl.collectors.Collector`
+              should always be preferred. If multiple simultaneous environment
+              need to be executed on a single node, consider using a
+              :class:`~torchrl.envs.ParallelEnv` instance.
+        collector_kwargs (dict or list, optional): a dictionary of parameters to be passed to the
+            remote data-collector. If a list is provided, each element will
+            correspond to an individual set of keyword arguments for the
+            dedicated collector.
+        num_workers_per_collector (int, optional): the number of copies of the
+            env constructor that is to be used on the remote nodes.
+            Defaults to 1 (a single env per collector).
+            On a single worker node all the sub-workers will be
+            executing the same environment. If different environments need to
+            be executed, they should be dispatched across worker nodes, not
+            subnodes.
+        sync (bool, optional): if ``True``, the resulting tensordict is a stack of all the
+            tensordicts collected on each node. If ``False`` (default), each
+            tensordict results from a separate node in a "first-ready,
+            first-served" fashion.
+        slurm_kwargs (dict): a dictionary of parameters to be passed to the
+            submitit executor.
+        update_after_each_batch (bool, optional): if ``True``, the weights will
+            be updated after each collection. For ``sync=True``, this means that
+            all workers will see their weights updated. For ``sync=False``,
+            only the worker from which the data has been gathered will be
+            updated.
+            Defaults to ``False``, ie. updates have to be executed manually
+            through
+            :meth:`~torchrl.collectors.distributed.DistributedDataCollector.update_policy_weights_`.
+        max_weight_update_interval (int, optional): the maximum number of
+            batches that can be collected before the policy weights of a worker
+            is updated.
+            For sync collections, this parameter is overwritten by ``update_after_each_batch``.
+            For async collections, it may be that one worker has not seen its
+            parameters being updated for a certain time even if ``update_after_each_batch``
+            is turned on.
+            Defaults to -1 (no forced update).
+        launcher (str, optional): how jobs should be launched.
+            Can be one of "submitit" or "mp" for multiprocessing. The former
+            can launch jobs across multiple nodes, whilst the latter will only
+            launch jobs on a single machine. "submitit" requires the homonymous
+            library to be installed.
+            To find more about submitit, visit
+            https://github.com/facebookincubator/submitit
+            Defaults to "submitit".
+        tcp_port (int, optional): the TCP port to be used. Defaults to 10003.
+        backend (str, optional): the torch.distributed backend to use for weight synchronization.
+            Must be one of ``"gloo"``, ``"mpi"``, ``"nccl"`` or ``"ucc"``. See the torch.distributed
+            documentation for more information. Defaults to ``"gloo"``.
+        visible_devices (list of Union[int, torch.device, str], optional): a
+            list of the same length as the number of nodes containing the
+            device used to pass data to main.
+        tensorpipe_options (dict, optional): a dictionary of keyword argument
+            to pass to :class:`torch.distributed.rpc.TensorPipeRpcBackendOption`.
+        weight_updater (WeightUpdaterBase or constructor, optional): An instance of :class:`~torchrl.collectors.WeightUpdaterBase`
+            or its subclass, responsible for updating the policy weights on remote inference workers using RPC.
+            If not provided, an :class:`~torchrl.collectors.distributed.RPCWeightUpdater` will be used by default, which
+            handles weight synchronization via RPC.
+            Consider using a constructor if the updater needs to be serialized.
+        weight_sync_schemes (dict[str, WeightSyncScheme], optional): Dictionary of weight sync schemes for
+            SENDING weights to remote collector workers. Keys are model identifiers (e.g., "policy")
+            and values are WeightSyncScheme instances configured to send weights via RPC.
+            If not provided, an :class:`~torchrl.weight_update.RPCWeightSyncScheme` will be used by default.
+            This is for propagating weights from the main process to remote collectors.
+        weight_recv_schemes (dict[str, WeightSyncScheme], optional): Dictionary of weight sync schemes for
+            RECEIVING weights from a parent process or training loop. Keys are model identifiers (e.g., "policy")
+            and values are WeightSyncScheme instances configured to receive weights.
+            This is typically used when RPCDataCollector is itself a worker in a larger distributed setup.
+            Defaults to ``None``.
+
+    """
+
+    _VERBOSE = VERBOSE  # for debugging
+
+    def __init__(
+        self,
+        create_env_fn,
+        policy: Callable[[TensorDictBase], TensorDictBase] | None = None,
+        *,
+        policy_factory: Callable[[], Callable]
+        | list[Callable[[]], Callable]
+        | None = None,
+        frames_per_batch: int,
+        total_frames: int = -1,
+        device: torch.device | list[torch.device] = None,
+        storing_device: torch.device | list[torch.device] = None,
+        env_device: torch.device | list[torch.device] = None,
+        policy_device: torch.device | list[torch.device] = None,
+        max_frames_per_traj: int = -1,
+        init_random_frames: int = -1,
+        reset_at_each_iter: bool = False,
+        postproc: Callable | None = None,
+        split_trajs: bool = False,
+        exploration_type: ExporationType = DEFAULT_EXPLORATION_TYPE,  # noqa
+        collector_class: type = Collector,
+        collector_kwargs: dict[str, Any] | None = None,
+        num_workers_per_collector: int = 1,
+        sync: bool = False,
+        slurm_kwargs: dict[str, Any] | None = None,
+        update_after_each_batch: bool = False,
+        max_weight_update_interval: int = -1,
+        launcher: str = "submitit",
+        tcp_port: str | None = None,
+        backend: str = "gloo",
+        visible_devices: list[torch.device] | None = None,
+        tensorpipe_options: dict[str, Any] | None = None,
+        weight_updater: WeightUpdaterBase
+        | Callable[[], WeightUpdaterBase]
+        | None = None,
+        weight_sync_schemes: dict[str, WeightSyncScheme] | None = None,
+        weight_recv_schemes: dict[str, WeightSyncScheme] | None = None,
+    ):
+
+        if self._VERBOSE:
+            torchrl_logger.setLevel("DEBUG")
+
+        if collector_class == "async":
+            collector_class = MultiAsyncCollector
+        elif collector_class == "sync":
+            collector_class = MultiSyncCollector
+        elif collector_class == "single":
+            collector_class = Collector
+        self.collector_class = collector_class
+        self.env_constructors = create_env_fn
+        self.policy = policy
+        if isinstance(policy, nn.Module):
+            policy_weights = TensorDict.from_module(policy)
+            policy_weights = policy_weights.data.lock_()
+        else:
+            if weight_updater is None and (
+                policy_factory is None
+                or (isinstance(policy_factory, Sequence) and not any(policy_factory))
+            ):
+                warnings.warn(_NON_NN_POLICY_WEIGHTS)
+            policy_weights = TensorDict(lock=True)
+
+        if not isinstance(policy_factory, Sequence):
+            policy_factory = [policy_factory] * len(create_env_fn)
+        self.policy_factory = policy_factory
+        self.policy_weights = policy_weights
+        self.num_workers = len(create_env_fn)
+        self.frames_per_batch = frames_per_batch
+        self.requested_frames_per_batch = frames_per_batch
+
+        self.device = device
+        self.storing_device = storing_device
+        self.env_device = env_device
+        self.policy_device = policy_device
+
+        self.storing_device = storing_device
+        # make private to avoid changes from users during collection
+        self._sync = sync
+        self.update_after_each_batch = update_after_each_batch
+        self.max_weight_update_interval = max_weight_update_interval
+        if self.update_after_each_batch and self.max_weight_update_interval > -1:
+            raise RuntimeError(
+                "Got conflicting update instructions: `update_after_each_batch` "
+                "`max_weight_update_interval` are incompatible."
+            )
+        self.launcher = launcher
+        self._batches_since_weight_update = [0 for _ in range(self.num_workers)]
+        if tcp_port is None:
+            self.tcp_port = os.environ.get("TCP_PORT", TCP_PORT)
+        else:
+            self.tcp_port = str(tcp_port)
+        self.visible_devices = visible_devices
+        if self._sync:
+            if self.frames_per_batch % self.num_workers != 0:
+                raise RuntimeError(
+                    f"Cannot dispatch {self.frames_per_batch} frames across {self.num_workers}. "
+                    f"Consider using a number of frames per batch that is divisible by the number of workers."
+                )
+            self._frames_per_batch_corrected = self.frames_per_batch // self.num_workers
+        else:
+            self._frames_per_batch_corrected = self.frames_per_batch
+
+        self.num_workers_per_collector = num_workers_per_collector
+        self.total_frames = total_frames
+        self.slurm_kwargs = copy(DEFAULT_SLURM_CONF)
+        if slurm_kwargs is not None:
+            self.slurm_kwargs.update(slurm_kwargs)
+
+        collector_kwargs = collector_kwargs if collector_kwargs is not None else {}
+        self.collector_kwargs = (
+            deepcopy(collector_kwargs)
+            if isinstance(collector_kwargs, (list, tuple))
+            else [copy(collector_kwargs) for _ in range(self.num_workers)]
+        )
+
+        # update collector kwargs
+        for i, collector_kwarg in enumerate(self.collector_kwargs):
+            collector_kwarg["max_frames_per_traj"] = max_frames_per_traj
+            collector_kwarg["init_random_frames"] = (
+                init_random_frames // self.num_workers
+            )
+            if not self._sync and init_random_frames > 0:
+                warnings.warn(
+                    "async distributed data collection with init_random_frames > 0 "
+                    "may have unforeseen consequences as we do not control that once "
+                    "non-random data is being collected all nodes are returning non-random data. "
+                    "If this is a feature that you feel should be fixed, please raise an issue on "
+                    "torchrl's repo."
+                )
+            collector_kwarg["reset_at_each_iter"] = reset_at_each_iter
+            collector_kwarg["exploration_type"] = exploration_type
+            collector_kwarg["device"] = self.device[i]
+            collector_kwarg["storing_device"] = self.storing_device[i]
+            collector_kwarg["env_device"] = self.env_device[i]
+            collector_kwarg["policy_device"] = self.policy_device[i]
+
+        self.postproc = postproc
+        self.split_trajs = split_trajs
+        self.backend = backend
+
+        if tensorpipe_options is None:
+            self.tensorpipe_options = copy(DEFAULT_TENSORPIPE_OPTIONS)
+        else:
+            self.tensorpipe_options = copy(DEFAULT_TENSORPIPE_OPTIONS).update(
+                tensorpipe_options
+            )
+
+        # 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 import RPCWeightSyncScheme
+
+            weight_sync_schemes = {"policy": RPCWeightSyncScheme()}
+
+        if weight_sync_schemes is not None:
+            # Use new weight synchronization system
+            self._weight_sync_schemes = weight_sync_schemes
+            self.weight_updater = None
+        else:
+            # Fall back to legacy weight updater system
+            if weight_updater is None:
+                weight_updater = RPCWeightUpdater(
+                    collector_infos=self.collector_infos,
+                    collector_class=self.collector_class,
+                    collector_rrefs=self.collector_rrefs,
+                    policy_weights=self.policy_weights,
+                    num_workers=self.num_workers,
+                )
+            self.weight_updater = weight_updater
+            self._weight_sync_schemes = None
+
+        self._init()
+
+        if weight_sync_schemes is not None:
+            # Set up weight senders now that remote collectors exist
+            for model_id, scheme in self._weight_sync_schemes.items():
+                scheme.init_on_sender(
+                    model_id=model_id,
+                    num_workers=self.num_workers,
+                    context=self,
+                )
+                scheme.connect()
+
+        # Set up weight receivers if provided
+        if weight_recv_schemes is not None:
+            self.register_scheme_receiver(weight_recv_schemes)
+
+    @property
+    def device(self) -> list[torch.device]:
+        return self._device
+
+    @property
+    def storing_device(self) -> list[torch.device]:
+        return self._storing_device
+
+    @property
+    def env_device(self) -> list[torch.device]:
+        return self._env_device
+
+    @property
+    def policy_device(self) -> list[torch.device]:
+        return self._policy_device
+
+    @device.setter
+    def device(self, value):
+        if isinstance(value, (tuple, list)):
+            if len(value) != self.num_workers:
+                raise RuntimeError(
+                    "The number of devices passed to the collector must match the number of workers."
+                )
+            self._device = value
+        else:
+            self._device = [value] * self.num_workers
+
+    @storing_device.setter
+    def storing_device(self, value):
+        if isinstance(value, (tuple, list)):
+            if len(value) != self.num_workers:
+                raise RuntimeError(
+                    "The number of devices passed to the collector must match the number of workers."
+                )
+            self._storing_device = value
+        else:
+            self._storing_device = [value] * self.num_workers
+
+    @env_device.setter
+    def env_device(self, value):
+        if isinstance(value, (tuple, list)):
+            if len(value) != self.num_workers:
+                raise RuntimeError(
+                    "The number of devices passed to the collector must match the number of workers."
+                )
+            self._env_device = value
+        else:
+            self._env_device = [value] * self.num_workers
+
+    @policy_device.setter
+    def policy_device(self, value):
+        if isinstance(value, (tuple, list)):
+            if len(value) != self.num_workers:
+                raise RuntimeError(
+                    "The number of devices passed to the collector must match the number of workers."
+                )
+            self._policy_device = value
+        else:
+            self._policy_device = [value] * self.num_workers
+
+    def _init_master_rpc(
+        self,
+        world_size,
+    ):
+        """Init torch.distributed and RPC on main node."""
+        # Initialize torch.distributed process group for efficient weight transfer
+        torchrl_logger.debug(
+            f"init distributed with rank=0, world_size={world_size}, backend={self.backend}"
+        )
+        torch.distributed.init_process_group(
+            backend=self.backend,
+            rank=0,
+            world_size=world_size,
+        )
+
+        # Initialize RPC for control/signaling
+        options = rpc.TensorPipeRpcBackendOptions(**self.tensorpipe_options)
+        if torch.cuda.is_available():
+            if self.visible_devices:
+                for i in range(self.num_workers):
+                    rank = i + 1
+                    options.set_device_map(
+                        f"COLLECTOR_NODE_{rank}", {0: self.visible_devices[i]}
+                    )
+        torchrl_logger.debug("init rpc")
+        rpc.init_rpc(
+            "TRAINER_NODE",
+            rank=0,
+            backend=rpc.BackendType.TENSORPIPE,
+            rpc_backend_options=options,
+            world_size=world_size,
+        )
+
+    def _start_workers(
+        self,
+        world_size,
+        env_constructors,
+        collector_class,
+        num_workers_per_collector,
+        policy,
+        policy_factory,
+        frames_per_batch,
+        total_frames,
+        collector_kwargs,
+    ):
+        """Instantiate remote collectors."""
+        num_workers = world_size - 1
+        time_interval = 1.0
+        collector_infos = []
+        for i in range(num_workers):
+            counter = 0
+            while True:
+                counter += 1
+                time.sleep(time_interval)
+                try:
+                    torchrl_logger.debug(f"trying to connect to collector node {i + 1}")
+                    collector_info = rpc.get_worker_info(f"COLLECTOR_NODE_{i + 1}")
+                    break
+                except RuntimeError as err:
+                    if counter * time_interval > self.tensorpipe_options["rpc_timeout"]:
+                        raise RuntimeError("Could not connect to remote node") from err
+                    continue
+            collector_infos.append(collector_info)
+
+        collector_rrefs = []
+        for i in range(num_workers):
+            env_make = env_constructors[i]
+            if not isinstance(env_make, (EnvBase, EnvCreator)):
+                env_make = CloudpickleWrapper(env_make)
+            torchrl_logger.debug("Making collector in remote node")
+            # When using weight sync schemes together with a policy_factory, the
+            # main-node `policy` should be used only as a weight source on the
+            # trainer, and NOT sent to remote collectors (which will build their
+            # own policies from the factory). This mirrors the behaviour of
+            # `DistributedDataCollector` with multi-process collectors.
+            policy_to_send = (
+                None
+                if (
+                    policy is not None
+                    and policy_factory[i] is not None
+                    and getattr(self, "_weight_sync_schemes", None) is not None
+                )
+                else policy
+            )
+
+            collector_rref = rpc.remote(
+                collector_infos[i],
+                collector_class,
+                args=(
+                    [env_make] * num_workers_per_collector
+                    if collector_class is not Collector
+                    else env_make,
+                    policy_to_send,
+                ),
+                kwargs={
+                    "policy_factory": policy_factory[i],
+                    "frames_per_batch": frames_per_batch,
+                    "total_frames": -1,
+                    "split_trajs": False,
+                    "weight_recv_schemes": self._weight_sync_schemes,
+                    "worker_idx": i,
+                    **collector_kwargs[i],
+                },
+            )
+            collector_rrefs.append(collector_rref)
+
+        # Set up receiver schemes on remote collectors (if using new weight sync system)
+        # This enables cascading: RPC -> MultiSync -> Sync
+        if getattr(self, "_weight_sync_schemes", None) is not None:
+            for i in range(num_workers):
+                torchrl_logger.debug(
+                    f"Setting up receiver schemes on remote collector {i}"
+                )
+                # Call register_scheme_receiver on the remote collector using rref.rpc_sync()
+                # This properly dereferences the rref and calls the instance method
+                collector_rrefs[i].rpc_sync().register_scheme_receiver(
+                    self._weight_sync_schemes
+                )
+
+        futures = collections.deque(maxlen=self.num_workers)
+
+        if not self._sync:
+            for i in range(num_workers):
+                torchrl_logger.debug("Asking for the first batch")
+                # Use rref.rpc_async() to properly call instance method
+                future = collector_rrefs[i].rpc_async().next()
+                futures.append((future, i))
+        self.futures = futures
+        self.collector_rrefs = collector_rrefs
+        self.collector_infos = collector_infos
+
+    def _init_worker_rpc(self, executor, i):
+        """Init RPC node if necessary."""
+        visible_device = (
+            self.visible_devices[i] if self.visible_devices is not None else None
+        )
+        if self.launcher == "submitit":
+            if not _has_submitit:
+                raise ImportError("submitit not found.") from SUBMITIT_ERR
+            job = executor.submit(
+                _rpc_init_collection_node,
+                i + 1,
+                self.IPAddr,
+                self.tcp_port,
+                self.num_workers + 1,
+                visible_device,
+                self.tensorpipe_options,
+                self.backend,
+                self._VERBOSE,
+            )
+            torchrl_logger.debug(f"job id {job.job_id}")  # ID of your job
+            return job
+        elif self.launcher == "mp":
+            job = _ProcessNoWarn(
+                target=_rpc_init_collection_node,
+                args=(
+                    i + 1,
+                    self.IPAddr,
+                    self.tcp_port,
+                    self.num_workers + 1,
+                    visible_device,
+                    self.tensorpipe_options,
+                    self.backend,
+                    self._VERBOSE,
+                ),
+            )
+            job.start()
+            return job
+        elif self.launcher == "submitit_delayed":
+            # job is already launched
+            return None
+        else:
+            raise NotImplementedError(f"Unknown launcher {self.launcher}")
+
+    def _init(self):
+        self._shutdown = False
+        if self.launcher == "submitit":
+            executor = submitit.AutoExecutor(folder="log_test")
+            executor.update_parameters(**self.slurm_kwargs)
+        else:
+            executor = None
+
+        hostname = socket.gethostname()
+        if self.launcher != "mp":
+            IPAddr = socket.gethostbyname(hostname)
+        else:
+            IPAddr = "localhost"
+        self.IPAddr = IPAddr
+
+        os.environ["MASTER_ADDR"] = str(self.IPAddr)
+        os.environ["MASTER_PORT"] = str(self.tcp_port)
+
+        self.jobs = []
+        for i in range(self.num_workers):
+            torchrl_logger.debug(f"Submitting job {i}")
+            job = self._init_worker_rpc(
+                executor,
+                i,
+            )
+            self.jobs.append(job)
+
+        self._init_master_rpc(
+            self.num_workers + 1,
+        )
+        self._start_workers(
+            world_size=self.num_workers + 1,
+            env_constructors=self.env_constructors,
+            collector_class=self.collector_class,
+            num_workers_per_collector=self.num_workers_per_collector,
+            policy=self.policy,
+            policy_factory=self.policy_factory,
+            frames_per_batch=self._frames_per_batch_corrected,
+            total_frames=self.total_frames,
+            collector_kwargs=self.collector_kwargs,
+        )
+
+    def iterator(self):
+        self._collected_frames = 0
+        while self._collected_frames < self.total_frames:
+            if self._sync:
+                data = self._next_sync_rpc()
+            else:
+                data = self._next_async_rpc()
+
+            if self.split_trajs:
+                data = split_trajectories(data)
+            if self.postproc is not None:
+                data = self.postproc(data)
+            yield data
+
+            if self.max_weight_update_interval > -1 and not self._sync:
+                for j in range(self.num_workers):
+                    if (
+                        self._batches_since_weight_update[j]
+                        > self.max_weight_update_interval
+                    ):
+                        torchrl_logger.debug(
+                            f"Updating policy of worker {j} with wait=False"
+                        )
+                        self.update_policy_weights_(worker_ids=[j], wait=False)
+            elif self.max_weight_update_interval > -1:
+                ranks = [
+                    1
+                    for j in range(self.num_workers)
+                    if self._batches_since_weight_update[j]
+                    > self.max_weight_update_interval
+                ]
+                torchrl_logger.debug(
+                    f"Updating policy of workers {ranks} with wait=True"
+                )
+                self.update_policy_weights_(worker_ids=ranks, wait=True)
+
+    def _next_async_rpc(self):
+        torchrl_logger.debug("next async")
+        if not len(self.futures):
+            raise StopIteration(
+                f"The queue is empty, the collector has ran out of data after {self._collected_frames} collected frames."
+            )
+        while True:
+            future, i = self.futures.popleft()
+            if future.done():
+                if self.update_after_each_batch:
+                    self.update_policy_weights_(worker_ids=(i,), wait=False)
+                torchrl_logger.debug(f"future {i} is done")
+                data = future.value()
+                self._collected_frames += data.numel()
+                if self._collected_frames < self.total_frames:
+                    # Use rref.rpc_async() to properly call instance method
+                    future = self.collector_rrefs[i].rpc_async().next()
+                    self.futures.append((future, i))
+                return data
+            self.futures.append((future, i))
+
+    def _next_sync_rpc(self):
+        torchrl_logger.debug("next sync: futures")
+        if self.update_after_each_batch:
+            self.update_policy_weights_()
+        for i in range(self.num_workers):
+            # Use rref.rpc_async() to properly call instance method
+            future = self.collector_rrefs[i].rpc_async().next()
+            self.futures.append((future, i))
+        data = []
+        while len(self.futures):
+            future, i = self.futures.popleft()
+            # the order is NOT guaranteed: should we change that?
+            if future.done():
+                data += [future.value()]
+                torchrl_logger.debug(
+                    f"got data from {i} // data has len {len(data)} / {self.num_workers}"
+                )
+            else:
+                self.futures.append((future, i))
+        data = torch.cat(data)
+        traj_ids = data.get(("collector", "traj_ids"), None)
+        if traj_ids is not None:
+            for i in range(1, self.num_workers):
+                traj_ids[i] += traj_ids[i - 1].max()
+            data.set_(("collector", "traj_ids"), traj_ids)
+        self._collected_frames += data.numel()
+        return data
+
+    def set_seed(self, seed: int, static_seed: bool = False) -> int:
+        for worker in self.collector_infos:
+            seed = rpc.rpc_sync(worker, self.collector_class.set_seed, args=(seed,))
+
+    def state_dict(self) -> OrderedDict:
+        raise NotImplementedError
+
+    def load_state_dict(self, state_dict: OrderedDict) -> None:
+        raise NotImplementedError
+
+    def shutdown(self, timeout: float | None = None) -> None:
+        if not hasattr(self, "_shutdown"):
+            warnings.warn("shutdown has no effect has `_init` has not been called yet.")
+            return
+        if self._shutdown:
+            return
+
+        torchrl_logger.debug("shutting down")
+        for future, i in self.futures:
+            # clear the futures
+            while future is not None and not future.done():
+                torchrl_logger.debug(f"waiting for proc {i} to clear")
+                future.wait()
+        for i in range(self.num_workers):
+            torchrl_logger.debug(f"shutting down {i}")
+            # Use rref.rpc_sync() to properly call instance method
+            self.collector_rrefs[i].rpc_sync(timeout=int(IDLE_TIMEOUT)).shutdown()
+        torchrl_logger.debug("rpc shutdown")
+        rpc.shutdown(timeout=int(IDLE_TIMEOUT))
+
+        if self.launcher == "mp":
+            for job in self.jobs:
+                job.join(int(IDLE_TIMEOUT))
+        elif self.launcher == "submitit":
+            for job in self.jobs:
+                _ = job.result()
+        elif self.launcher == "submitit_delayed":
+            pass
+        else:
+            raise NotImplementedError(f"Unknown launcher {self.launcher}")
+
+        # Clean up weight sync schemes AFTER workers have exited
+        if getattr(self, "_weight_sync_schemes", None) is not None:
+            torchrl_logger.debug("shutting down weight sync schemes")
+            for scheme in self._weight_sync_schemes.values():
+                try:
+                    scheme.shutdown()
+                except Exception as e:
+                    torchrl_logger.warning(
+                        f"Error shutting down weight sync scheme: {e}"
+                    )
+            self._weight_sync_schemes = None
+
+        # Destroy torch.distributed process group
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+
+        self._shutdown = True
+
+
+class RPCWeightUpdater(WeightUpdaterBase):
+    """A remote weight updater for synchronizing policy weights across remote workers using RPC.
+
+    The `RPCWeightUpdater` class provides a mechanism for updating the weights of a policy
+    across remote inference workers using RPC. It is designed to work with the :class:`~torchrl.collectors.distributed.RPCDataCollector`
+    to ensure that each worker receives the latest policy weights.
+    This class is typically used in distributed data collection scenarios where remote workers
+    are managed via RPC and need to be kept in sync with the central policy weights.
+
+    Args:
+        collector_infos: Information about the collectors, used for RPC communication.
+        collector_class: The class of the collectors being used.
+        collector_rrefs: Remote references to the collectors.
+        policy_weights (TensorDictBase): The current weights of the policy that need to be distributed
+            to the workers.
+        num_workers (int): The number of remote workers that will receive the updated policy weights.
+
+    Methods:
+        update_weights: Updates the weights on specified or all remote workers using RPC.
+        all_worker_ids: Returns a list of all worker identifiers (not implemented in this class).
+        _sync_weights_with_worker: Synchronizes the server weights with a specific worker (not implemented).
+        _get_server_weights: Retrieves the latest weights from the server (not implemented).
+        _maybe_map_weights: Optionally maps server weights before distribution (not implemented).
+
+    .. note::
+        This class assumes that the server weights can be directly applied to the remote workers
+        without any additional processing. If your use case requires more complex weight mapping or
+        synchronization logic, consider extending `WeightUpdaterBase` with a custom implementation.
+
+    .. seealso:: :class:`~torchrl.collectors.WeightUpdaterBase` and
+        :class:`~torchrl.collectors.distributed.RPCDataCollector`.
+
+    """
+
+    _VERBOSE = VERBOSE  # for debugging
+
+    def __init__(
+        self,
+        collector_infos,
+        collector_class,
+        collector_rrefs,
+        policy_weights: TensorDictBase,
+        num_workers: int,
+    ):
+        super().__init__()
+        self.collector_infos = collector_infos
+        self.collector_class = collector_class
+        self.collector_rrefs = collector_rrefs
+        self.policy_weights = policy_weights
+        self.num_workers = num_workers
+
+    def _sync_weights_with_worker(
+        self, worker_id: int | torch.device, server_weights: TensorDictBase
+    ) -> TensorDictBase:
+        raise NotImplementedError
+
+    def _get_server_weights(self) -> TensorDictBase:
+        raise NotImplementedError
+
+    def _maybe_map_weights(self, server_weights: TensorDictBase) -> TensorDictBase:
+        raise NotImplementedError
+
+    def all_worker_ids(self) -> list[int] | list[torch.device]:
+        raise NotImplementedError
+
+    def push_weights(
+        self,
+        weights: TensorDictBase | None = None,
+        worker_ids: torch.device | int | list[int] | list[torch.device] | None = None,
+        **kwargs,
+    ):
+        workers = worker_ids
+        if isinstance(workers, int):
+            workers = [workers]
+        if workers is None:
+            workers = list(range(self.num_workers))
+        else:
+            workers = list(workers)
+        futures = []
+        weights = self.policy_weights if weights is None else weights
+        for i in workers:
+            torchrl_logger.debug(f"calling update on worker {i}")
+            # Use rref.rpc_async() to properly call instance method
+            futures.append(
+                self.collector_rrefs[i].rpc_async().update_policy_weights_(weights)
+            )
+        if kwargs.get("wait", True):
+            for i in workers:
+                torchrl_logger.debug(f"waiting for worker {i}")
+                futures[i].wait()
+                torchrl_logger.debug("got it!")
+
+
+class RPCDataCollector(RPCCollector, metaclass=_LegacyCollectorMeta):
+    """Deprecated version of :class:`~torchrl.collectors.distributed.RPCCollector`."""
+
+    ...