Update

[ghstack-poisoned]

Author: vmoens

benchmarks/bench_collectors.py

@@ -12,7 +12,8 @@
   2. Collector (ParallelEnv x N)        -- single-process, N envs in sub-procs
   3. MultiCollector (sync, x N)         -- N sub-processes, sync delivery
   4. MultiCollector (async, x N)        -- N sub-processes, async delivery
-  5. AsyncBatchedCollector (threading)   -- AsyncEnvPool + InferenceServer
+  5. AsyncBatched (env=thread, pol=thread)  -- threading pool + threading transport
+  6. AsyncBatched (env=mp, pol=thread)      -- multiprocessing pool + threading transport
 """
 from __future__ import annotations
 
@@ -368,33 +369,33 @@ def policy_factory():
         )
     )
 
-    # 5. AsyncBatchedCollector (threading backend)
+    # 5. AsyncBatchedCollector (env=threading, policy=threading)
     results.append(
         bench(
-            f"AsyncBatchedCollector threading (x{num_envs})",
+            f"AsyncBatched env=thread pol=thread (x{num_envs})",
             lambda: AsyncBatchedCollector(
                 create_env_fn=[make_env_fn] * num_envs,
                 policy=policy_factory(),
                 frames_per_batch=frames_per_batch,
                 total_frames=-1,
                 max_batch_size=num_envs,
-                backend="threading",
+                env_backend="threading",
             ),
             target_frames=total_frames,
         )
     )
 
-    # 6. AsyncBatchedCollector (multiprocessing backend)
+    # 6. AsyncBatchedCollector (env=multiprocessing, policy=threading)
     results.append(
         bench(
-            f"AsyncBatchedCollector mp (x{num_envs})",
+            f"AsyncBatched env=mp pol=thread (x{num_envs})",
             lambda: AsyncBatchedCollector(
                 create_env_fn=[make_env_fn] * num_envs,
                 policy=policy_factory(),
                 frames_per_batch=frames_per_batch,
                 total_frames=-1,
                 max_batch_size=num_envs,
-                backend="multiprocessing",
+                env_backend="multiprocessing",
             ),
             target_frames=total_frames,
         )

examples/collectors/async_batched_collector.py

@@ -6,14 +6,21 @@
 
 Architecture:
   - An :class:`~torchrl.envs.AsyncEnvPool` runs environments in parallel
-    using the chosen backend (``"threading"`` or ``"multiprocessing"``).
+    using the chosen ``env_backend`` (``"threading"`` or ``"multiprocessing"``).
   - One lightweight coordinator thread per environment owns a slot in the pool
     and an inference client.
   - An :class:`~torchrl.modules.InferenceServer` batches incoming observations
-    and runs a single forward pass.
+    and runs a single forward pass.  The communication layer (transport) is
+    controlled by ``policy_backend`` (``"threading"``, ``"multiprocessing"``,
+    ``"ray"``, or ``"monarch"``).
   - There is no global synchronisation barrier -- fast envs keep stepping
     while slow ones wait for inference.
 
+Backend parameters:
+  - ``backend``        -- global default for both env pool and policy transport.
+  - ``env_backend``    -- override for the env pool (falls back to ``backend``).
+  - ``policy_backend`` -- override for the transport (falls back to ``backend``).
+
 The user only supplies:
   - A list of environment factories
   - A policy (or policy factory)

test/test_inference_server.py

@@ -733,7 +733,7 @@ def test_basic_collection(self):
             frames_per_batch=frames_per_batch,
             total_frames=total_frames,
             max_batch_size=num_envs,
-            backend="threading",
+            env_backend="threading",
         )
         total_collected = 0
         for batch in collector:
@@ -751,7 +751,7 @@ def test_policy_factory(self):
             frames_per_batch=10,
             total_frames=20,
             max_batch_size=num_envs,
-            backend="threading",
+            env_backend="threading",
         )
         total_collected = 0
         for batch in collector:
@@ -791,7 +791,7 @@ def test_yield_completed_trajectories(self):
             total_frames=30,
             yield_completed_trajectories=True,
             max_batch_size=num_envs,
-            backend="threading",
+            env_backend="threading",
         )
         count = 0
         for batch in collector:
@@ -809,7 +809,7 @@ def test_shutdown_idempotent(self):
             policy=policy,
             frames_per_batch=10,
             total_frames=10,
-            backend="threading",
+            env_backend="threading",
         )
         # Consume one batch to start
         for _batch in collector:
@@ -825,7 +825,7 @@ def test_endless_collector(self):
             policy=policy,
             frames_per_batch=10,
             total_frames=-1,
-            backend="threading",
+            env_backend="threading",
         )
         collected = 0
         for batch in collector:
@@ -862,7 +862,7 @@ def postproc(td):
             frames_per_batch=10,
             total_frames=20,
             postproc=postproc,
-            backend="threading",
+            env_backend="threading",
         )
         for _ in collector:
             pass

torchrl/collectors/_async_batched.py

@@ -21,6 +21,45 @@
 
 _ENV_IDX_KEY = "env_index"
 
+_POLICY_BACKENDS = ("threading", "multiprocessing", "ray", "monarch")
+_ENV_BACKENDS = ("threading", "multiprocessing")
+
+
+def _make_transport(
+    policy_backend: str, num_slots: int | None = None
+) -> InferenceTransport:
+    """Create an :class:`InferenceTransport` from a backend name.
+
+    Args:
+        policy_backend: one of ``"threading"``, ``"multiprocessing"``,
+            ``"ray"``, or ``"monarch"``.
+        num_slots: when set and ``policy_backend="threading"``, a
+            :class:`~torchrl.modules.SlotTransport` is created instead of
+            the generic :class:`~torchrl.modules.ThreadingTransport`.
+    """
+    if policy_backend == "threading":
+        if num_slots is not None:
+            from torchrl.modules.inference_server._slot import SlotTransport
+
+            return SlotTransport(num_slots)
+        return ThreadingTransport()
+    if policy_backend == "multiprocessing":
+        from torchrl.modules.inference_server._mp import MPTransport
+
+        return MPTransport()
+    if policy_backend == "ray":
+        from torchrl.modules.inference_server._ray import RayTransport
+
+        return RayTransport()
+    if policy_backend == "monarch":
+        from torchrl.modules.inference_server._monarch import MonarchTransport
+
+        return MonarchTransport()
+    raise ValueError(
+        f"Unknown policy_backend {policy_backend!r}. "
+        f"Expected one of {_POLICY_BACKENDS}."
+    )
+
 
 def _env_loop(
     pool: AsyncEnvPool,
@@ -47,9 +86,7 @@ def _env_loop(
 
         while not shutdown_event.is_set():
             pool.async_step_and_maybe_reset_send(action_td, env_index=env_id)
-            cur_td, next_obs = pool.async_step_and_maybe_reset_recv(
-                env_index=env_id
-            )
+            cur_td, next_obs = pool.async_step_and_maybe_reset_recv(env_index=env_id)
             cur_td.set(_ENV_IDX_KEY, env_id)
             result_queue.put(cur_td)
             if shutdown_event.is_set():
@@ -104,22 +141,35 @@ class AsyncBatchedCollector(BaseCollector):
         max_batch_size (int, optional): upper bound on the number of
             requests the inference server processes in a single forward pass.
             Defaults to ``64``.
+        min_batch_size (int, optional): minimum number of requests the
+            inference server accumulates before dispatching a batch.  After
+            the first request arrives the server keeps draining for up to
+            ``server_timeout`` seconds until this many items are collected.
+            ``1`` (default) dispatches immediately.
         server_timeout (float, optional): seconds the server waits for work
             before dispatching a partial batch.  Defaults to ``0.01``.
         transport (InferenceTransport, optional): a pre-built transport
-            backend.  When ``None`` (default) a
-            :class:`~torchrl.modules.ThreadingTransport` is created
-            automatically (since worker threads always live in the main
-            process).  Pass a :class:`~torchrl.modules.RayTransport` or
-            :class:`~torchrl.modules.MonarchTransport` for distributed
-            setups where the inference server is remote.
+            object.  When provided, it takes precedence over
+            ``policy_backend``.  When ``None`` (default) a transport is
+            created automatically from the resolved ``policy_backend``.
         device (torch.device or str, optional): device for policy inference.
             Passed to the inference server.  Defaults to ``None``.
-        backend (str, optional): backend for the
+        backend (str, optional): global default backend for both
+            environments and policy inference.  Specific overrides
+            ``env_backend`` and ``policy_backend`` take precedence when set.
+            One of ``"threading"``, ``"multiprocessing"``, ``"ray"``, or
+            ``"monarch"``.  Defaults to ``"threading"``.
+        env_backend (str, optional): backend for the
             :class:`~torchrl.envs.AsyncEnvPool` that runs environments.  One
-            of ``"threading"`` or ``"multiprocessing"``.  The coordinator
-            threads are always Python threads regardless of this setting.
-            Defaults to ``"threading"``.
+            of ``"threading"`` or ``"multiprocessing"``.  Falls back to
+            ``backend`` when ``None``.  The coordinator threads are always
+            Python threads regardless of this setting.  Defaults to ``None``.
+        policy_backend (str, optional): backend for the inference transport
+            used to communicate with the
+            :class:`~torchrl.modules.InferenceServer`.  One of
+            ``"threading"``, ``"multiprocessing"``, ``"ray"``, or
+            ``"monarch"``.  Falls back to ``backend`` when ``None``.
+            Defaults to ``None``.
         reset_at_each_iter (bool, optional): whether to reset all envs at the
             start of every collection batch.  Defaults to ``False``.
         postproc (Callable, optional): post-processing transform applied to
@@ -169,10 +219,16 @@ def __init__(
         frames_per_batch: int,
         total_frames: int = -1,
         max_batch_size: int = 64,
+        min_batch_size: int = 1,
         server_timeout: float = 0.01,
         transport: InferenceTransport | None = None,
         device: torch.device | str | None = None,
-        backend: Literal["threading", "multiprocessing"] = "threading",
+        backend: Literal[
+            "threading", "multiprocessing", "ray", "monarch"
+        ] = "threading",
+        env_backend: Literal["threading", "multiprocessing"] | None = None,
+        policy_backend: Literal["threading", "multiprocessing", "ray", "monarch"]
+        | None = None,
         reset_at_each_iter: bool = False,
         postproc: Callable[[TensorDictBase], TensorDictBase] | None = None,
         yield_completed_trajectories: bool = False,
@@ -196,19 +252,34 @@ def __init__(
             raise TypeError("create_env_fn must be a list of env factories.")
         self._create_env_fn = list(create_env_fn)
         self._num_envs = len(create_env_fn)
-        self._backend = backend
         self._create_env_kwargs = create_env_kwargs
 
+        # ---- resolve backends -------------------------------------------------
+        effective_env_backend = env_backend if env_backend is not None else backend
+        effective_policy_backend = (
+            policy_backend if policy_backend is not None else backend
+        )
+        if effective_env_backend not in _ENV_BACKENDS:
+            raise ValueError(
+                f"env_backend={effective_env_backend!r} is not supported. "
+                f"Expected one of {_ENV_BACKENDS}."
+            )
+        self._env_backend = effective_env_backend
+        self._policy_backend = effective_policy_backend
+
         # ---- build transport --------------------------------------------------
         if transport is None:
-            transport = ThreadingTransport()
+            transport = _make_transport(
+                effective_policy_backend, num_slots=self._num_envs
+            )
         self._transport = transport
 
         # ---- build inference server -------------------------------------------
         self._server = InferenceServer(
             model=policy,
             transport=transport,
             max_batch_size=max_batch_size,
+            min_batch_size=min_batch_size,
             timeout=server_timeout,
             device=device,
             weight_sync=weight_sync,
@@ -252,7 +323,7 @@ def _ensure_started(self) -> None:
             kwargs["create_env_kwargs"] = self._create_env_kwargs
         self._env_pool = AsyncEnvPool(
             self._create_env_fn,
-            backend=self._backend,
+            backend=self._env_backend,
             **kwargs,
         )
 
@@ -303,9 +374,18 @@ def _rollout_frames(self) -> TensorDictBase:
         transitions: list[TensorDictBase] = []
 
         while collected < self.frames_per_batch:
+            # Block for at least one transition
             td = rq.get()
             transitions.append(td)
             collected += td.numel()
+            # Batch-drain any additional items already in the queue
+            while collected < self.frames_per_batch:
+                try:
+                    td = rq.get_nowait()
+                except queue.Empty:
+                    break
+                transitions.append(td)
+                collected += td.numel()
             if self.verbose:
                 torchrl_logger.debug(
                     f"AsyncBatchedCollector: {collected}/{self.frames_per_batch} frames"

torchrl/envs/async_envs.py

@@ -673,9 +673,7 @@ def async_step_and_maybe_reset_send(
         for _env_idx, local_td in _zip_strict(env_idx, local_tds):
             if not _per_env:
                 self._current_step_reset = self._current_step_reset + 1
-            self.input_queue[_env_idx].put(
-                ("step_and_maybe_reset", local_td, _per_env)
-            )
+            self.input_queue[_env_idx].put(("step_and_maybe_reset", local_td, _per_env))
 
     def async_step_and_maybe_reset_recv(
         self, min_get: int = 1, env_index: int | None = None
@@ -807,29 +805,29 @@ def _env_exec(
             elif msg == "batch_size":
                 output_queue.put(env.batch_size)
             elif msg == "reset":
-                data = env.reset(data.copy())
+                # No .copy() needed: data was deserialized from the queue
+                # and is not referenced after this call.
+                data = env.reset(data)
                 data.set(cls._env_idx_key, NonTensorData(i))
                 target = per_env_reset_queue if per_env else reset_queue
                 target.put(data)
             elif msg == "_reset":
-                data = env._reset(data.copy())
+                data = env._reset(data)
                 data.set(cls._env_idx_key, NonTensorData(i))
                 reset_queue.put(data)
             elif msg == "step_and_maybe_reset":
-                data, data_ = env.step_and_maybe_reset(data.copy())
+                data, data_ = env.step_and_maybe_reset(data)
                 data.set(cls._env_idx_key, NonTensorData(i))
                 data_.set(cls._env_idx_key, NonTensorData(i))
-                target = (
-                    per_env_step_reset_queue if per_env else step_reset_queue
-                )
+                target = per_env_step_reset_queue if per_env else step_reset_queue
                 target.put((data, data_))
             elif msg == "step":
-                data = env.step(data.copy())
+                data = env.step(data)
                 data.set(cls._env_idx_key, NonTensorData(i))
                 target = per_env_step_queue if per_env else step_queue
                 target.put(data)
             elif msg == "_step":
-                data = env._step(data.copy())
+                data = env._step(data)
                 data.set(cls._env_idx_key, NonTensorData(i))
                 step_queue.put(data)
             elif msg == "shutdown":

torchrl/modules/inference_server/__init__.py

@@ -7,6 +7,7 @@
 from torchrl.modules.inference_server._mp import MPTransport
 from torchrl.modules.inference_server._ray import RayTransport
 from torchrl.modules.inference_server._server import InferenceClient, InferenceServer
+from torchrl.modules.inference_server._slot import SlotTransport
 from torchrl.modules.inference_server._threading import ThreadingTransport
 from torchrl.modules.inference_server._transport import InferenceTransport
 
@@ -17,5 +18,6 @@
     "MonarchTransport",
     "MPTransport",
     "RayTransport",
+    "SlotTransport",
     "ThreadingTransport",
 ]