[Feature] Add generation-stamped replay buffer slots

benchmarks/test_replaybuffer_benchmark.py

@@ -402,6 +402,35 @@ def test_rb_populate(benchmark, rb, storage, sampler, size):
     )
 
 
+class create_wraparound_rb:
+    """Builds a full buffer so every timed extend reuses slots and bumps generations."""
+
+    def __init__(self, size=10_000, batch=1_000):
+        self.size = size
+        self.batch = batch
+
+    def __call__(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(self.size))
+        data = TensorDict({"a": torch.zeros(self.batch, 5)}, batch_size=[self.batch])
+        while rb.write_count < self.size:
+            rb.extend(data)
+        return ((rb, data), {})
+
+
+def extend_wraparound(rb, data):
+    for _ in range(10):
+        rb.extend(data)
+
+
+def test_rb_extend_generation_stamping(benchmark):
+    benchmark.pedantic(
+        extend_wraparound,
+        setup=create_wraparound_rb(),
+        iterations=1,
+        rounds=50,
+    )
+
+
 class create_compiled_tensor_rb:
     def __init__(
         self, rb, storage, sampler, storage_size, data_size, iters, compilable=False

test/rb/test_ensemble.py

@@ -437,7 +437,7 @@ def test_rb_multidim(self, datatype, datadim, rbtype, storage_cls, sampler_cls):
         s = rb.sample()
         assert str(rb)
         if datatype in ("tensordict", "tensorclass"):
-            assert (s.exclude("index") == 1).all()
+            assert (s.exclude("index", "index_generation") == 1).all()
             assert s.numel() == 4
         else:
             for leaf in tree_iter(s):

test/rb/test_rb_core.py

@@ -1196,6 +1196,53 @@ def test_stats_with_non_counting_writer(self):
         assert stats["capacity"] == 10
 
 
+class TestSampleGenerationInfo:
+    """Executable spec for exposing slot generations at sampling time (RFC step 1).
+
+    Contract pinned by this class:
+
+    - ``sample(return_info=True)`` returns an ``"index_generation"`` entry in
+      the info dict, aligned element-for-element with ``info["index"]`` and
+      equal to the writer's generation for those slots at sampling time.
+    - ``TensorDictReplayBuffer`` samples carry an ``"index_generation"`` key
+      alongside the existing ``"index"`` key.
+    - After a wraparound, previously captured (index, generation) handles
+      disagree with the writer's current generations exactly on the reused
+      slots, which is what makes staleness detectable.
+    """
+
+    def test_sample_info_contains_index_generation(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10), batch_size=4)
+        rb.extend(torch.arange(10))
+        _, info = rb.sample(return_info=True)
+        assert "index_generation" in info
+        index = torch.as_tensor(info["index"]).reshape(-1)
+        generations = torch.as_tensor(info["index_generation"]).reshape(-1)
+        assert generations.shape == index.shape
+        torch.testing.assert_close(generations, rb._writer.generations_of(index))
+
+    def test_tensordict_sample_carries_index_generation(self):
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10), batch_size=4)
+        rb.extend(TensorDict({"obs": torch.randn(10, 3)}, batch_size=[10]))
+        rb.extend(TensorDict({"obs": torch.randn(4, 3)}, batch_size=[4]))
+        sample = rb.sample()
+        assert "index_generation" in sample.keys()
+        batch = sample.batch_size[0]
+        index = sample.get("index").reshape(batch, -1)[:, 0]
+        generations = sample.get("index_generation").reshape(batch, -1)[:, 0]
+        torch.testing.assert_close(generations, rb._writer.generations_of(index))
+
+    def test_stale_handles_detectable_after_wraparound(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10), batch_size=4)
+        index = rb.extend(torch.arange(10))
+        generations = rb._writer.generations_of(index)
+        rb.extend(torch.arange(4))
+        live = rb._writer.generations_of(index) == generations
+        assert live.sum() == 6
+        assert not live[:4].any()
+        assert live[4:].all()
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

test/rb/test_samplers.py

@@ -243,7 +243,7 @@ def test_sampler_without_rep_state_dict(self, backend):
             replay_buffer.extend(transition.clone())
         for _ in range(n_samples):
             s = replay_buffer.sample(batch_size=1)
-            assert (s.exclude("index") == 1).all()
+            assert (s.exclude("index", "index_generation") == 1).all()
 
         replay_buffer.extend(torch.zeros_like(transition))
 
@@ -257,7 +257,7 @@ def test_sampler_without_rep_state_dict(self, backend):
 
         new_replay_buffer.load_state_dict(state_dict)
         s = new_replay_buffer.sample(batch_size=1)
-        assert (s.exclude("index") == 0).all()
+        assert (s.exclude("index", "index_generation") == 0).all()
 
     def test_sampler_without_rep_dumps_loads(self, tmpdir):
         d0 = tmpdir + "/save0"

test/rb/test_storages.py

@@ -306,7 +306,7 @@ def test_storage_state_dict(self, storage_in, storage_out, init_out, backend):
 
         new_replay_buffer.load_state_dict(state_dict)
         s = new_replay_buffer.sample()
-        assert (s.exclude("index") == 1).all()
+        assert (s.exclude("index", "index_generation") == 1).all()
 
     @pytest.mark.skipif(
         TORCH_VERSION < version.parse("2.5.0"), reason="requires Torch >= 2.5.0"

test/rb/test_writers.py

@@ -403,6 +403,95 @@ def test_roundrobin_dumps_loads_write_count(self, tmp_path):
         assert writer2._write_count == 23
 
 
+class TestSlotGenerations:
+    """Executable spec for generation-stamped replay slots (RFC step 1).
+
+    Contract pinned by this class:
+
+    - Round-robin writers maintain one int64 generation counter per storage
+      slot, exposed through ``writer.generations_of(index)`` which accepts an
+      index tensor and returns a same-shaped int64 tensor.
+    - The first write of a slot has generation 0; every reuse of a slot
+      (round-robin wraparound or rewrite through ``add``/``extend``)
+      increments that slot's generation.
+    - ``empty()`` never revives previously handed-out (index, generation)
+      pairs: generations are monotonically nondecreasing across the buffer's
+      lifetime, including through ``empty()``.
+    - Generations persist through ``state_dict``/``load_state_dict`` and
+      ``dumps``/``loads``; checkpoints created before the feature still load.
+    """
+
+    def test_first_writes_start_at_generation_zero(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        index = rb.extend(torch.arange(10))
+        generations = rb._writer.generations_of(index)
+        assert generations.dtype == torch.int64
+        assert generations.shape == index.shape
+        assert (generations == 0).all()
+
+    def test_wraparound_increments_reused_slots_only(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(torch.arange(10))
+        reused_index = rb.extend(torch.arange(4))
+        assert (rb._writer.generations_of(reused_index) == 1).all()
+        untouched = rb._writer.generations_of(torch.arange(4, 10))
+        assert (untouched == 0).all()
+
+    def test_add_reuse_increments_generation(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(2))
+        for value in range(5):
+            rb.add(torch.full((3,), float(value)))
+        assert rb._writer.generations_of(torch.tensor([0])).item() == 2
+        assert rb._writer.generations_of(torch.tensor([1])).item() == 1
+
+    def test_empty_never_revives_old_handles(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        index = rb.extend(torch.arange(10))
+        generations_before = rb._writer.generations_of(index)
+        rb.empty()
+        rb.extend(torch.arange(10))
+        generations_after = rb._writer.generations_of(index)
+        assert (generations_after > generations_before).all()
+
+    def test_generations_survive_state_dict_roundtrip(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(torch.arange(10))
+        index = rb.extend(torch.arange(4))
+        sd = rb.state_dict()
+        rb2 = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb2.load_state_dict(sd)
+        torch.testing.assert_close(
+            rb2._writer.generations_of(torch.arange(10)),
+            rb._writer.generations_of(torch.arange(10)),
+        )
+        assert (rb2._writer.generations_of(index) == 1).all()
+
+    def test_generations_survive_dumps_loads(self, tmp_path):
+        rb = ReplayBuffer(storage=LazyMemmapStorage(10, scratch_dir=tmp_path / "data"))
+        rb.extend(torch.arange(10))
+        rb.extend(torch.arange(4))
+        rb._writer.dumps(tmp_path / "writer")
+        writer2 = RoundRobinWriter()
+        writer2.loads(tmp_path / "writer")
+        torch.testing.assert_close(
+            writer2.generations_of(torch.arange(10)),
+            rb._writer.generations_of(torch.arange(10)),
+        )
+
+    def test_legacy_state_dict_without_generations_loads(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(torch.arange(5))
+        sd = rb.state_dict()
+        sd["_writer"] = {
+            key: value
+            for key, value in sd["_writer"].items()
+            if key in ("_cursor", "_write_count")
+        }
+        rb2 = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb2.load_state_dict(sd)
+        assert rb2._writer._cursor == 5
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/data/replay_buffers/replay_buffers.py

@@ -1516,6 +1516,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         with self._replay_lock if not is_comp else nc, self._write_lock if not is_comp else nc:
             index, info = self._sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -1555,6 +1559,10 @@ def sample(self, batch_size: int | None = None, return_info: bool = False) -> An
         Returns:
             A batch of data selected in the replay buffer.
             A tuple containing this batch and info if return_info flag is set to True.
+            The info entries include ``"index"``, the storage slots the batch
+            was read from, and ``"index_generation"``, the generation of each
+            slot at sampling time (see
+            :meth:`~torchrl.data.replay_buffers.RoundRobinWriter.generations_of`).
         """
         if (
             batch_size is not None
@@ -2135,6 +2143,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         ):
             index, info = self.prioritized_sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2564,6 +2576,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         with self._replay_lock if not is_comp else nc, self._write_lock if not is_comp else nc:
             index, info = self._sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2929,6 +2945,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         ):
             index, info = self.prioritized_sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)