pytorch · vmoens · Jun 6, 2025 · Jun 6, 2025
diff --git a/docs/source/reference/collectors.rst b/docs/source/reference/collectors.rst
@@ -159,8 +159,15 @@ transformed, and applied, ensuring seamless integration with their existing infr
     VanillaWeightUpdater
     MultiProcessedWeightUpdater
     RayWeightUpdater
-    DistributedWeightUpdater
+
+.. currentmodule:: torchrl.collectors.distributed
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template.rst
+
     RPCWeightUpdater
+    DistributedWeightUpdater
 
 Collectors and replay buffers interoperability
 ----------------------------------------------

diff --git a/torchrl/collectors/__init__.py b/torchrl/collectors/__init__.py
@@ -13,7 +13,7 @@
     SyncDataCollector,
 )
 from .weight_update import (
-    MultiProcessedWeightUpdate,
+    MultiProcessedWeightUpdater,
     RayWeightUpdater,
     VanillaWeightUpdater,
     WeightUpdaterBase,
@@ -24,7 +24,7 @@
     "WeightUpdaterBase",
     "VanillaWeightUpdater",
     "RayWeightUpdater",
-    "MultiProcessedWeightUpdate",
+    "MultiProcessedWeightUpdater",
     "aSyncDataCollector",
     "DataCollectorBase",
     "MultiaSyncDataCollector",

diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -51,7 +51,7 @@
 )
 from torchrl.collectors.utils import split_trajectories
 from torchrl.collectors.weight_update import (
-    MultiProcessedWeightUpdate,
+    MultiProcessedWeightUpdater,
     VanillaWeightUpdater,
     WeightUpdaterBase,
 )
@@ -2010,7 +2010,7 @@ def __init__(
                 self._policy_weights_dict[policy_device] = weights
             self._get_weights_fn = get_weights_fn
             if weight_updater is None:
-                weight_updater = MultiProcessedWeightUpdate(
+                weight_updater = MultiProcessedWeightUpdater(
                     get_server_weights=self._get_weights_fn,
                     policy_weights=self._policy_weights_dict,
                 )

diff --git a/torchrl/collectors/distributed/__init__.py b/torchrl/collectors/distributed/__init__.py
@@ -9,7 +9,7 @@
     DistributedWeightUpdater,
 )
 from .ray import RayCollector
-from .rpc import RPCDataCollector
+from .rpc import RPCDataCollector, RPCWeightUpdater
 from .sync import DistributedSyncDataCollector
 from .utils import submitit_delayed_launcher
 
@@ -19,7 +19,7 @@
     "DistributedWeightUpdater",
     "DistributedSyncDataCollector",
     "RPCDataCollector",
-    "RPCDataCollector",
+    "RPCWeightUpdater",
     "RayCollector",
     "submitit_delayed_launcher",
 ]
diff --git a/torchrl/collectors/distributed/rpc.py b/torchrl/collectors/distributed/rpc.py
@@ -412,7 +412,7 @@ def __init__(
             )
         self._init()
         if weight_updater is None:
-            weight_updater = RPCWeightUpdaterBase(
+            weight_updater = RPCWeightUpdater(
                 collector_infos=self.collector_infos,
                 collector_class=self.collector_class,
                 collector_rrefs=self.collector_rrefs,
@@ -810,7 +810,7 @@ def shutdown(self, timeout: float | None = None) -> None:
         self._shutdown = True
 
 
-class RPCWeightUpdaterBase(WeightUpdaterBase):
+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

diff --git a/torchrl/collectors/llm/weight_update/vllm.py b/torchrl/collectors/llm/weight_update/vllm.py
@@ -4,6 +4,8 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
+import importlib.util
+
 import torch
 import torch.cuda
 import torch.distributed
@@ -13,7 +15,16 @@
 
 from torchrl.collectors import WeightUpdaterBase
 from torchrl.modules.llm.backends.vllm import stateless_init_process_group
-from vllm.utils import get_open_port
+
+_has_vllm = importlib.util.find_spec("vllm") is not None
+if _has_vllm:
+    from vllm.utils import get_open_port
+else:
+
+    def get_open_port():  # noqa: D103
+        raise ImportError(
+            "vllm is not installed. Please install it with `pip install vllm`."
+        )
 
 
 class vLLMUpdater(WeightUpdaterBase):

diff --git a/torchrl/collectors/weight_update.py b/torchrl/collectors/weight_update.py
@@ -206,7 +206,7 @@ def _sync_weights_with_worker(
         self.policy_weights.update_(server_weights)
 
 
-class MultiProcessedWeightUpdate(WeightUpdaterBase):
+class MultiProcessedWeightUpdater(WeightUpdaterBase):
     """A remote weight updater for synchronizing policy weights across multiple processes or devices.
 
     The `MultiProcessedWeightUpdater` class provides a mechanism for updating the weights

diff --git a/torchrl/envs/llm/chat.py b/torchrl/envs/llm/chat.py
@@ -7,7 +7,6 @@
 from typing import Any, Callable, Literal
 
 import torch
-import transformers
 from tensordict import lazy_stack, TensorDict, TensorDictBase
 from torch.utils.data import DataLoader
 from torchrl.data import Composite, NonTensor
@@ -116,7 +115,7 @@ def __init__(
         batch_size: tuple | torch.Size | None = None,
         system_prompt: str | None = None,
         apply_template: bool | None = None,
-        tokenizer: transformers.AutoTokenizer | None = None,
+        tokenizer: transformers.AutoTokenizer | None = None,  # noqa: F821
         template_kwargs: dict[str, Any] | None = None,
         system_role: str = "system",
         user_role: str = "user",
@@ -309,7 +308,7 @@ def __init__(
         batch_size_dl: int = 1,
         seed: int | None = None,
         group_repeats: bool = False,
-        tokenizer: transformers.AutoTokenizer | None = None,
+        tokenizer: transformers.AutoTokenizer | None = None,  # noqa: F821
         device: torch.device | None = None,
         template_kwargs: dict[str, Any] | None = None,
         apply_template: bool | None = None,