[BugFix] Bound direct off-policy replay collection

test/rb/test_rb_distributed.py

@@ -10,6 +10,7 @@
 import sys
 import time
 from functools import partial
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -18,8 +19,11 @@
 from _rb_common import _has_ray
 from tensordict import TensorDict
 from torchrl._utils import logger as torchrl_logger
-from torchrl.data import RayReplayBuffer, ReplayBuffer
-from torchrl.data.replay_buffers import RemoteTensorDictReplayBuffer
+from torchrl.data import RayReplayBuffer, ReplayBuffer, TensorDictReplayBuffer
+from torchrl.data.replay_buffers import (
+    ray_buffer as ray_buffer_module,
+    RemoteTensorDictReplayBuffer,
+)
 from torchrl.data.replay_buffers.samplers import (
     RandomSampler,
     SamplerWithoutReplacement,
@@ -44,6 +48,19 @@ def __init__(self, capacity: int, scratch_dir=None):
         )
 
 
+class _RemoteResult:
+    def __init__(self, result):
+        self.result = result
+
+    def remote(self, *args, **kwargs):
+        return self.result
+
+
+class _RejectingRemoteResult:
+    def remote(self, *args, **kwargs):
+        raise AssertionError("A pre-bound transport used payload bootstrap.")
+
+
 def construct_buffer_test(rank, name, world_size):
     if name == "TRAINER":
         buffer = _construct_buffer("BUFFER")
@@ -294,6 +311,74 @@ def test_ray_replay_with_gloo_transport(self):
         finally:
             rb.shutdown()
 
+    def test_prebound_distributed_replay_skips_payload_bootstrap(self, monkeypatch):
+        extend_data = TensorDict({"value": torch.arange(8).reshape(8, 1)}, [8])
+        sample = TensorDict(
+            {
+                "value": torch.zeros(4, 1, dtype=torch.int64),
+                "index": torch.zeros(4, dtype=torch.int64),
+            },
+            [4],
+        )
+
+        def extend_endpoint(data, *, timeout=None):
+            assert data is extend_data
+            assert timeout == 30.0
+            return TensorDict({"result": torch.arange(8)}, [])
+
+        def sample_endpoint(request, *, timeout=None):
+            assert request["batch_size"].item() == 4
+            assert timeout == 30.0
+            return sample
+
+        actor = SimpleNamespace(
+            _distributed_control_client=_RemoteResult(object()),
+            _distributed_bound_extend_client=_RemoteResult(extend_endpoint),
+            _distributed_bound_sample_client=_RemoteResult((sample_endpoint, 4)),
+            _bootstrap_distributed_extend=_RejectingRemoteResult(),
+            _bootstrap_distributed_sample=_RejectingRemoteResult(),
+        )
+        monkeypatch.setattr(ray_buffer_module.ray, "get", lambda result: result)
+        monkeypatch.setattr(
+            ray_buffer_module,
+            "_set_ray_client_liveness",
+            lambda client, actor: None,
+        )
+        client = ray_buffer_module._LazyDistributedReplayClient(actor, batch_size=4)
+        assert client.extend(extend_data, timeout=30.0).shape == (8,)
+        assert client.sample(timeout=30.0) is sample
+
+    def test_ray_replay_with_prebound_gloo_transport(self):
+        extend_data = TensorDict({"value": torch.arange(8).reshape(8, 1)}, [8])
+        sample_spec = TensorDict(
+            {
+                "value": torch.zeros(4, 1, dtype=torch.int64),
+                "index": torch.zeros(4, dtype=torch.int64),
+            },
+            [4],
+        )
+        replay = TensorDictReplayBuffer(
+            storage=partial(LazyTensorStorage, 64),
+            batch_size=4,
+            service_backend="ray",
+            service_backend_options={"remote_config": {"num_cpus": 0}},
+            transport="distributed",
+            transport_options={
+                "backend": "gloo",
+                "timeout": 30.0,
+                "extend_spec": extend_data.clone().zero_(),
+                "sample_spec": sample_spec,
+            },
+        )
+        try:
+            writer = replay.client()
+            reader = replay.client()
+            assert writer.extend(extend_data).shape == (8,)
+            assert writer.write_count == 8
+            assert reader.sample().shape == (4,)
+        finally:
+            replay.shutdown()
+
     @pytest.mark.gpu
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
     def test_ray_replay_with_nccl_transport(self):

test/test_distributed.py

@@ -97,6 +97,12 @@ def forward(self, tensordict):
         return tensordict
 
 
+class BatchedCountingPolicy(CountingPolicy):
+    def forward(self, tensordict):
+        tensordict.set("action", self.weight.expand(*tensordict.shape, 1).clone())
+        return tensordict
+
+
 class MismatchedCountingPolicy(CountingPolicy):
     def __init__(self):
         super().__init__()
@@ -137,6 +143,10 @@ def load_state_dict(self, state_dict):
         self.calls = state_dict["calls"]
 
 
+def flatten_batch(tensordict):
+    return tensordict.reshape(-1)
+
+
 class MultiOptimizerRayLoss(LossModule):
     def __init__(self):
         super().__init__()
@@ -1414,6 +1424,110 @@ def test_ray_owned_inference_and_replay(self):
             inference.shutdown()
             replay.shutdown()
 
+    @pytest.mark.parametrize("flatten", [False, True])
+    def test_async_trainer_tracks_frames_independently_of_replay_shape(self, flatten):
+        num_envs = 2
+        frames_per_batch = 8
+        total_frames = 16
+        replay = TensorDictReplayBuffer(
+            storage=partial(LazyTensorStorage, 100),
+            batch_size=2,
+            service_backend="ray",
+            service_backend_options={"remote_config": {"num_cpus": 0}},
+            transport="auto",
+        )
+        policy = BatchedCountingPolicy()
+        collector = RayCollector(
+            create_env_fn=[partial(CountingEnv, batch_size=[num_envs])],
+            policy=policy,
+            replay_buffer=replay,
+            collector_class=Collector,
+            collector_kwargs={"postproc": flatten_batch} if flatten else None,
+            frames_per_batch=frames_per_batch,
+            total_frames=total_frames,
+            remote_configs={"num_cpus": 1, "num_gpus": 0},
+            sync=True,
+        )
+        loss = ScalarRayLoss()
+        with pytest.warns(UserWarning, match="experimental"):
+            trainer = DQNTrainer(
+                collector=collector,
+                total_frames=total_frames,
+                frame_skip=1,
+                optim_steps_per_batch=1,
+                loss_module=loss,
+                optimizer=torch.optim.SGD(loss.parameters(), lr=0.1),
+                replay_buffer=replay,
+                async_collection=True,
+                progress_bar=False,
+                batch_size=2,
+                target_net_updater=CountingTargetUpdater(loss),
+                learner_backend="ray",
+                learner_backend_options={
+                    "world_size": 2,
+                    "resources_per_rank": {"num_cpus": 1, "num_gpus": 0},
+                    "backend": "gloo",
+                    "setup_timeout": 60.0,
+                    "command_timeout": 60.0,
+                },
+                enable_logging=False,
+            )
+        try:
+            trainer.train()
+            assert trainer.collected_frames == total_frames
+            assert collector.collected_frames == total_frames
+            assert trainer._optim_count > 0
+            assert trainer._published_model_version == trainer._optim_count
+            if flatten:
+                assert replay.write_count == total_frames
+                assert replay.sample().shape == (2,)
+            else:
+                assert replay.write_count == total_frames // (
+                    frames_per_batch // num_envs
+                )
+                assert replay[0].shape == (frames_per_batch // num_envs,)
+                assert replay.sample().shape == (
+                    2,
+                    frames_per_batch // num_envs,
+                )
+        finally:
+            collector.shutdown()
+            replay.shutdown()
+
+    def test_async_ray_collector_does_not_write_unreserved_batches(self):
+        num_envs = 2
+        frames_per_batch = 4
+        total_frames = 16
+        replay = TensorDictReplayBuffer(
+            storage=partial(LazyTensorStorage, 100),
+            batch_size=2,
+            service_backend="ray",
+            service_backend_options={"remote_config": {"num_cpus": 0}},
+            transport="auto",
+        )
+        policy = BatchedCountingPolicy()
+        env_fn = partial(CountingEnv, batch_size=[num_envs])
+        collector = RayCollector(
+            create_env_fn=[env_fn, env_fn],
+            policy=policy,
+            replay_buffer=replay,
+            collector_class=Collector,
+            collector_kwargs={"postproc": flatten_batch},
+            frames_per_batch=frames_per_batch,
+            total_frames=total_frames,
+            remote_configs={"num_cpus": 1, "num_gpus": 0},
+            sync=False,
+        )
+        try:
+            collector.update_policy_weights_(policy)
+            assert all(batch is None for batch in collector)
+            assert collector.collected_frames == total_frames
+            assert replay.write_count == total_frames
+            assert replay.sample().shape == (2,)
+        finally:
+            collector.shutdown()
+            replay.shutdown()
+
     def test_ray_collector_pause_drains_and_resumes(self):
         replay = TensorDictReplayBuffer(
             storage=partial(LazyTensorStorage, 1000),

torchrl/collectors/distributed/ray.py

@@ -1273,15 +1273,24 @@ async def async_shutdown(self, shutdown_ray: bool = False):
     def _async_iterator(self) -> Iterator[TensorDictBase]:
         """Collects a data batch from a single remote collector in each iteration."""
         pending_tasks = {}
+
+        def can_schedule() -> bool:
+            return self.total_frames < 0 or (
+                self.collected_frames + len(pending_tasks) * self.frames_per_batch
+                < self.total_frames
+            )
+
         for index, collector in enumerate(self.remote_collectors):
+            if not can_schedule():
+                break
             future = collector.next.remote()
             pending_tasks[future] = index
 
         while (
-            self.collected_frames < self.total_frames and not self._stop_event.is_set()
-        ):
-            if not len(list(pending_tasks.keys())) == len(self.remote_collectors):
-                raise RuntimeError("Missing pending tasks, something went wrong")
+            self.total_frames < 0 or self.collected_frames < self.total_frames
+        ) and not self._stop_event.is_set():
+            if not pending_tasks:
+                raise RuntimeError("No pending Ray collector tasks remain.")
 
             # Wait for first worker to finish
             wait_results = ray.wait(list(pending_tasks.keys()))
@@ -1305,13 +1314,17 @@ def _async_iterator(self) -> Iterator[TensorDictBase]:
                 torchrl_logger.debug(f"Updating weights on worker {collector_index}")
                 self.update_policy_weights_(worker_ids=collector_index)
 
-            # Schedule a new collection task
-            future = collector.next.remote()
-            pending_tasks[future] = collector_index
+            # Reserve in-flight batches against total_frames. Without this guard,
+            # every worker writes one extra direct-replay batch while the iterator
+            # drains pending actor calls after reaching the collection target.
+            if can_schedule():
+                future = collector.next.remote()
+                pending_tasks[future] = collector_index
 
         # Wait for the in-process collections tasks to finish.
         refs = list(pending_tasks.keys())
-        ray.wait(refs, num_returns=len(refs))
+        if refs:
+            ray.wait(refs, num_returns=len(refs))
 
         # Cancel the in-process collections tasks
         # for ref in refs:

torchrl/data/replay_buffers/ray_buffer.py

@@ -181,13 +181,19 @@ def _stats(self, timeout: float | None = None) -> TensorDictBase:
         return self._control_client(request, timeout=timeout)
 
     def extend(self, data: TensorDictBase, *, timeout: float | None = None):
+        handled = False
+        result = None
         if self._extend_client is None:
-            self._extend_client, result, handled = ray.get(
-                self._actor._bootstrap_distributed_extend.remote(data)
+            self._extend_client = ray.get(
+                self._actor._distributed_bound_extend_client.remote()
             )
+            if self._extend_client is None:
+                self._extend_client, result, handled = ray.get(
+                    self._actor._bootstrap_distributed_extend.remote(data)
+                )
             _set_ray_client_liveness(self._extend_client, self._actor)
-            if handled:
-                return result
+        if handled:
+            return result
         response = self._extend_client(data, timeout=timeout)
         return response.get("result", None)
 
@@ -202,26 +208,29 @@ def sample(self, batch_size: int | None = None, *, timeout: float | None = None)
             batch_size = self.batch_size
         if batch_size is None:
             raise RuntimeError("A sample batch size must be provided.")
+        handled = False
+        result = None
         if self._sample_client is None:
-            (
-                self._sample_client,
-                result,
-                handled,
-                self._sample_batch_size,
-            ) = ray.get(self._actor._bootstrap_distributed_sample.remote(batch_size))
-            _set_ray_client_liveness(self._sample_client, self._actor)
-            if batch_size != self._sample_batch_size:
-                raise ValueError(
-                    "The distributed replay schema is bound to sample batch size "
-                    f"{self._sample_batch_size}, got {batch_size}."
+            self._sample_client, self._sample_batch_size = ray.get(
+                self._actor._distributed_bound_sample_client.remote()
+            )
+            if self._sample_client is None:
+                (
+                    self._sample_client,
+                    result,
+                    handled,
+                    self._sample_batch_size,
+                ) = ray.get(
+                    self._actor._bootstrap_distributed_sample.remote(batch_size)
                 )
-            if handled:
-                return result
-        elif batch_size != self._sample_batch_size:
+            _set_ray_client_liveness(self._sample_client, self._actor)
+        if batch_size != self._sample_batch_size:
             raise ValueError(
                 "The distributed replay schema is bound to sample batch size "
                 f"{self._sample_batch_size}, got {batch_size}."
             )
+        if handled:
+            return result
         request = TensorDict(
             {
                 "batch_size": torch.tensor(

torchrl/data/replay_buffers/replay_buffers.py

@@ -610,6 +610,20 @@ def _distributed_control_client(self):
         """Return a restricted control endpoint for length and write count."""
         return self._distributed_service.control_client()
 
+    def _distributed_bound_extend_client(self):
+        """Return the pre-bound extend endpoint without transferring payload data."""
+        service = self._distributed_service
+        if service.extend_transport is None:
+            return None
+        return service.extend_client()
+
+    def _distributed_bound_sample_client(self):
+        """Return the pre-bound sample endpoint and its fixed batch size."""
+        service = self._distributed_service
+        if service.sample_transport is None:
+            return None, None
+        return service.sample_client(), service._sample_batch_size
+
     def _distributed_service_client(self):
         """Create an independently routed client for the private service."""
         service = getattr(self, "_distributed_service", None)