[BugFix] Buffer prefetching off by one (#3871)

Author: Agade09

test/rb/test_rb_core.py

@@ -623,6 +623,26 @@ def test_replay_buffer_iter(size, drop_last):
         assert i == (size - 1) // 3
 
 
+def test_replay_buffer_prefetch_queue_length():
+    """Test that the prefetch queue maintains the correct length.
+
+    This test verifies that after sampling from a replay buffer with prefetching
+    enabled, the prefetch queue has exactly `prefetch` items computing in the
+    background (no off-by-one error).
+    """
+    torch.manual_seed(0)
+
+    rb = ReplayBuffer(storage=ListStorage(max_size=100), batch_size=2, prefetch=2)
+
+    rb.extend(torch.arange(100))
+
+    _ = rb.sample()
+
+    assert (
+        len(rb._prefetch_queue) == 2
+    ), f"Expected prefetch queue to have 2 items, but got {len(rb._prefetch_queue)}."
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/data/replay_buffers/replay_buffers.py

@@ -1154,14 +1154,17 @@ def sample(self, batch_size: int | None = None, return_info: bool = False) -> An
             result = self._sample(batch_size)
         else:
             with self._futures_lock:
+                if len(self._prefetch_queue):
+                    result = self._prefetch_queue.popleft().result()
+                else:
+                    result = self._sample(batch_size)
                 while (
                     len(self._prefetch_queue)
                     < min(self._sampler._remaining_batches, self._prefetch_cap)
                     and not self._sampler.ran_out
-                ) or not len(self._prefetch_queue):
+                ):
                     fut = self._prefetch_executor.submit(self._sample, batch_size)
                     self._prefetch_queue.append(fut)
-                result = self._prefetch_queue.popleft().result()
 
         if return_info:
             out, info = result