From e1842b2c189d6f877d7bf0973a2a5d6ea61db638 Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Fri, 24 Jul 2026 12:43:21 +0530
Subject: [PATCH 01/13] [Feature] Add opt-in generation tracking to RoundRobin
 writer

---
 test/rb/test_writers.py                | 104 +++++++++++++++++++++
 torchrl/data/replay_buffers/writers.py | 122 ++++++++++++++++++++++++-
 2 files changed, 223 insertions(+), 3 deletions(-)

diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index ffe47a51ffb..fda1c6c3cf0 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -403,6 +403,110 @@ def test_roundrobin_dumps_loads_write_count(self, tmp_path):
         assert writer2._write_count == 23
 
 
+class TestWriterGeneration:
+    def test_generation_disabled_by_default(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        assert rb._writer.supports_generation is False
+        assert rb._writer._generation is None
+        with pytest.raises(RuntimeError, match="does not track generations"):
+            rb._writer._get_generation(0)
+
+    def test_generation_increments_on_reuse(self):
+        size = 4
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb.extend(torch.arange(size))
+        assert rb._writer.supports_generation is True
+        assert (rb._writer._generation == 1).all()
+        rb.extend(torch.arange(size, size + 3))
+        torch.testing.assert_close(rb._writer._generation, torch.tensor([2, 2, 2, 1]))
+
+    def test_generation_wraparound(self):
+        size = 5
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb.extend(torch.arange(2 * size))
+        assert (rb._writer._generation == 2).all()
+
+    def test_generation_add(self):
+        size = 3
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        for i in range(size + 1):
+            rb.add(torch.tensor(i))
+        torch.testing.assert_close(rb._writer._generation, torch.tensor([2, 1, 1]))
+
+    def test_generation_get(self):
+        size = 4
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb.extend(torch.arange(size + 2))
+        gen = rb._writer._get_generation(torch.tensor([0, 3]))
+        torch.testing.assert_close(gen, torch.tensor([2, 1]))
+
+    def test_generation_tensordict_writer(self):
+        size = 4
+        rb = TensorDictReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=TensorDictRoundRobinWriter(track_generation=True),
+        )
+        rb.extend(TensorDict({"a": torch.arange(2 * size)}, [2 * size]))
+        assert (rb._writer._generation == 2).all()
+
+    def test_generation_write_at(self):
+        storage = LazyTensorStorage(4)
+        writer = RoundRobinWriter(track_generation=True)
+        writer.register_storage(storage)
+        writer.extend(torch.arange(4))
+        writer.write_at(torch.tensor([0, 1]), torch.tensor([10, 11]))
+        torch.testing.assert_close(writer._generation, torch.tensor([2, 2, 1, 1]))
+
+    def test_generation_empty_resets(self):
+        storage = LazyTensorStorage(4)
+        writer = RoundRobinWriter(track_generation=True)
+        writer.register_storage(storage)
+        writer.extend(torch.arange(4))
+        writer._empty()
+        assert writer._generation is None
+
+    def test_generation_state_dict_roundtrip(self):
+        size = 4
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb.extend(torch.arange(size + 1))
+        sd = rb.state_dict()
+        rb2 = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb2.load_state_dict(sd)
+        assert rb2._writer.supports_generation is True
+        torch.testing.assert_close(rb2._writer._generation, rb._writer._generation)
+
+    def test_generation_dumps_loads(self, tmp_path):
+        writer = RoundRobinWriter(track_generation=True)
+        writer._cursor = 2
+        writer._write_count = 9
+        writer._generation = torch.tensor([3, 2, 2, 1])
+        writer.dumps(tmp_path)
+        writer2 = RoundRobinWriter()
+        writer2.loads(tmp_path)
+        assert writer2._cursor == 2
+        assert writer2._write_count == 9
+        assert writer2._track_generation is True
+        torch.testing.assert_close(writer2._generation, torch.tensor([3, 2, 2, 1]))
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 9726e9861dd..f894394bea6 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -53,6 +53,11 @@ def __init__(self, compilable: bool = False) -> None:
     def register_storage(self, storage: Storage) -> None:
         self._storage = storage
 
+    @property
+    def supports_generation(self) -> bool:
+        """Whether this writer maintains per-slot generation stamps."""
+        return False
+
     @abstractmethod
     def add(self, data: Any) -> int:
         """Inserts one piece of data at an appropriate index, and returns that index."""
@@ -153,18 +158,95 @@ class RoundRobinWriter(Writer):
             If ``True``, the writer cannot be shared between multiple processes.
             Defaults to ``False``.
 
+    Keyword Args:
+        track_generation (bool, optional): if ``True``, the writer maintains a
+            per-slot generation counter that is incremented every time a physical
+            storage slot is written. Samplers can then stamp sampled indices with
+            the generation they observed, allowing delayed priority updates to
+            detect whether a slot has since been reused. Defaults to ``False``.
+
     """
 
-    def __init__(self, compilable: bool = False) -> None:
+    def __init__(
+        self, compilable: bool = False, *, track_generation: bool = False
+    ) -> None:
         super().__init__(compilable=compilable)
         self._cursor = 0
         self._write_count  # noqa
+        self._track_generation = track_generation
+        self._generation = None
+
+    @property
+    def supports_generation(self) -> bool:
+        return self._track_generation
+
+    def _ensure_generation(self, size: int) -> None:
+        if self._generation is None:
+            generation = torch.zeros(int(size), dtype=torch.long)
+            if not self._compilable:
+                generation.share_memory_()
+            self._generation = generation
+
+    def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
+        if not self._track_generation:
+            return
+        if _is_int(index):
+            self._ensure_generation(
+                self._storage._max_size_along_dim0(single_data=data)
+            )
+            self._generation[int(index)] += 1
+        else:
+            self._ensure_generation(
+                self._storage._max_size_along_dim0(batched_data=data)
+            )
+            index = torch.as_tensor(index, dtype=torch.long).reshape(-1)
+            self._generation.index_put_(
+                (index,), torch.ones_like(index), accumulate=True
+            )
+
+    def _get_generation(self, index: int | torch.Tensor) -> torch.Tensor:
+        if not self._track_generation:
+            raise RuntimeError(
+                "This writer does not track generations. Construct it with "
+                "track_generation=True to use generation-stamped indices."
+            )
+        if self._generation is None:
+            raise RuntimeError("No data has been written yet; generation is undefined.")
+        if (
+            isinstance(index, torch.Tensor)
+            and index.ndim
+            and self._storage.ndim > 1
+            and index.shape[-1] == self._storage.ndim
+        ):
+            index = index[..., 0]
+        if _is_int(index):
+            return self._generation[int(index)]
+        return self._generation[torch.as_tensor(index, dtype=torch.long)]
 
     def dumps(self, path):
         path = Path(path).absolute()
         path.mkdir(exist_ok=True)
+        metadata = {
+            "cursor": self._cursor,
+            "write_count": self._write_count,
+            "track_generation": self._track_generation,
+        }
+        generation = self._generation
+        if generation is not None:
+            try:
+                MemoryMappedTensor.from_filename(
+                    filename=path / "generation.memmap",
+                    shape=generation.shape,
+                    dtype=generation.dtype,
+                ).copy_(generation)
+            except FileNotFoundError:
+                MemoryMappedTensor.from_tensor(
+                    generation, filename=path / "generation.memmap"
+                )
+            metadata["generation_shape"] = list(generation.shape)
+            metadata["generation_dtype"] = str(generation.dtype)
         with open(path / "metadata.json", "w") as file:
-            json.dump({"cursor": self._cursor, "write_count": self._write_count}, file)
+            json.dump(metadata, file)
 
     def loads(self, path):
         path = Path(path).absolute()
@@ -174,6 +256,19 @@ def loads(self, path):
             write_count = metadata.get("write_count")
             if write_count is not None:
                 self._write_count = write_count
+            self._track_generation = metadata.get(
+                "track_generation", self._track_generation
+            )
+            generation_shape = metadata.get("generation_shape")
+        if generation_shape is not None:
+            generation = MemoryMappedTensor.from_filename(
+                filename=path / "generation.memmap",
+                dtype=_STRDTYPE2DTYPE[metadata["generation_dtype"]],
+                shape=torch.Size(generation_shape),
+            ).clone()
+            if not self._compilable:
+                generation.share_memory_()
+            self._generation = generation
 
     def add(self, data: Any) -> int | torch.Tensor:
         index = self._cursor
@@ -186,6 +281,7 @@ def add(self, data: Any) -> int | torch.Tensor:
         # Replicate index requires the shape of the storage to be known
         # Other than that, a "flat" (1d) index is ok to write the data
         self._storage.set(_cursor, data)
+        self._bump_generation(_cursor, data)
         index = self._replicate_index(index)
         self._mark_update_entities(index)
         return index
@@ -214,6 +310,7 @@ def extend(self, data: Sequence) -> torch.Tensor:
         # Replicate index requires the shape of the storage to be known
         # Other than that, a "flat" (1d) index is ok to write the data
         self._storage.set(index, data)
+        self._bump_generation(index, data)
         index = self._replicate_index(index)
         self._mark_update_entities(index)
         return index
@@ -229,6 +326,7 @@ def write_at(self, index: int | torch.Tensor, data: Any) -> int | torch.Tensor:
             batch_size = index.numel()
         self._write_count += batch_size
         self._storage.set(index, data, set_cursor=False)
+        self._bump_generation(index, data)
         self._update_storage_len_for_write_at(index)
         index = self._replicate_index(index)
         self._mark_update_entities(index)
@@ -249,16 +347,31 @@ def _update_storage_len_for_write_at(self, index: int | torch.Tensor) -> None:
         )
 
     def state_dict(self) -> dict[str, Any]:
-        return {"_cursor": self._cursor, "_write_count": self._write_count}
+        state_dict = {"_cursor": self._cursor, "_write_count": self._write_count}
+        if self._track_generation:
+            state_dict["_track_generation"] = True
+            if self._generation is not None:
+                state_dict["_generation"] = self._generation.clone()
+        return state_dict
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         self._cursor = state_dict["_cursor"]
         write_count = state_dict.get("_write_count")
         if write_count is not None:
             self._write_count = write_count
+        self._track_generation = state_dict.get(
+            "_track_generation", self._track_generation
+        )
+        generation = state_dict.get("_generation")
+        if generation is not None:
+            generation = generation.clone()
+            if not self._compilable:
+                generation.share_memory_()
+            self._generation = generation
 
     def _empty(self, empty_write_count: bool = True) -> None:
         self._cursor = 0
+        self._generation = None
         if empty_write_count:
             self._write_count = 0
 
@@ -363,6 +476,7 @@ def add(self, data: Any) -> int | torch.Tensor:
                 ),
             )
         self._storage.set(index, data)
+        self._bump_generation(index, data)
         index = self._replicate_index(index)
         self._mark_update_entities(index)
         return index
@@ -392,6 +506,7 @@ def extend(self, data: Sequence) -> torch.Tensor:
         # Replicate index requires the shape of the storage to be known
         # Other than that, a "flat" (1d) index is ok to write the data
         self._storage.set(index, data)
+        self._bump_generation(index, data)
         index = self._replicate_index(index)
         self._mark_update_entities(index)
         return index
@@ -407,6 +522,7 @@ def write_at(self, index: int | torch.Tensor, data: Any) -> int | torch.Tensor:
         if not is_tensorclass(data):
             data.set("index", expand_as_right(index_tensor, data))
         self._storage.set(index_tensor, data, set_cursor=False)
+        self._bump_generation(index_tensor, data)
         self._update_storage_len_for_write_at(index_tensor)
         index = self._replicate_index(index_tensor)
         self._mark_update_entities(index)

From fa9e72254156c56cdf336a29ed712902cb354505 Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Fri, 24 Jul 2026 13:06:10 +0530
Subject: [PATCH 02/13] feat: Surface index_generation in replay sample info

---
 test/rb/test_writers.py                       | 59 +++++++++++++++++++
 torchrl/data/replay_buffers/replay_buffers.py |  8 +++
 torchrl/data/replay_buffers/writers.py        |  4 +-
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index fda1c6c3cf0..1a983b57004 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -506,6 +506,65 @@ def test_generation_dumps_loads(self, tmp_path):
         assert writer2._track_generation is True
         torch.testing.assert_close(writer2._generation, torch.tensor([3, 2, 2, 1]))
 
+    def test_sample_returns_generation(self):
+        size = 8
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb.extend(torch.arange(size))
+        _, info = rb.sample(4, return_info=True)
+        assert "index_generation" in info
+        gen = torch.as_tensor(info["index_generation"])
+        idx = torch.as_tensor(info["index"])
+        assert gen.shape == idx.shape
+        torch.testing.assert_close(gen, rb._writer._generation[idx])
+
+    def test_sample_no_generation_when_disabled(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(8))
+        rb.extend(torch.arange(8))
+        _, info = rb.sample(4, return_info=True)
+        assert "index_generation" not in info
+
+    def test_tensordict_sample_has_generation_key(self):
+        size = 8
+        rb = TensorDictReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=TensorDictRoundRobinWriter(track_generation=True),
+        )
+        rb.extend(TensorDict({"a": torch.arange(size)}, [size]))
+        sample = rb.sample(4)
+        assert "index_generation" in sample.keys()
+        assert sample["index_generation"].shape[0] == 4
+
+    def test_wraparound_race_detectable(self):
+        size = 8
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb.extend(torch.arange(size))
+        _, info = rb.sample(4, return_info=True)
+        sampled_index = torch.as_tensor(info["index"])
+        sampled_generation = torch.as_tensor(info["index_generation"])
+        rb.extend(torch.arange(size, 2 * size))
+        current = rb._writer._get_generation(sampled_index)
+        assert (current != sampled_generation).all()
+
+    def test_partial_reuse_detectable(self):
+        size = 8
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generation=True),
+        )
+        rb.extend(torch.arange(size))
+        _, info = rb.sample(size, return_info=True)
+        idx = torch.as_tensor(info["index"])
+        gen = torch.as_tensor(info["index_generation"])
+        rb.extend(torch.arange(size, size + 3))
+        stale = rb._writer._get_generation(idx) != gen
+        torch.testing.assert_close(stale, idx < 3)
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
diff --git a/torchrl/data/replay_buffers/replay_buffers.py b/torchrl/data/replay_buffers/replay_buffers.py
index 068aaf3e90f..6924a42a4a7 100644
--- a/torchrl/data/replay_buffers/replay_buffers.py
+++ b/torchrl/data/replay_buffers/replay_buffers.py
@@ -1534,6 +1534,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
+            if self._writer.supports_generation:
+                info["index_generation"] = self._writer._get_generation(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2157,6 +2159,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
+            if self._writer.supports_generation:
+                info["index_generation"] = self._writer._get_generation(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2634,6 +2638,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
+            if self._writer.supports_generation:
+                info["index_generation"] = self._writer._get_generation(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -3003,6 +3009,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
+            if self._writer.supports_generation:
+                info["index_generation"] = self._writer._get_generation(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index f894394bea6..8ce22373b37 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -212,7 +212,9 @@ def _get_generation(self, index: int | torch.Tensor) -> torch.Tensor:
             )
         if self._generation is None:
             raise RuntimeError("No data has been written yet; generation is undefined.")
-        if (
+        if isinstance(index, tuple):
+            index = index[0]
+        elif (
             isinstance(index, torch.Tensor)
             and index.ndim
             and self._storage.ndim > 1

From 7e10c67a069c163d50726d387e106eabe510d37e Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Fri, 24 Jul 2026 14:29:01 +0530
Subject: [PATCH 03/13] feat: default-on slot generation tracking for
 round-robin writer

---
 test/rb/test_ensemble.py                      |   2 +-
 test/rb/test_samplers.py                      |   4 +-
 test/rb/test_storages.py                      |   2 +-
 test/rb/test_writers.py                       | 169 ++++++++++--------
 torchrl/data/replay_buffers/replay_buffers.py |  16 +-
 torchrl/data/replay_buffers/writers.py        | 123 +++++++------
 6 files changed, 169 insertions(+), 147 deletions(-)

diff --git a/test/rb/test_ensemble.py b/test/rb/test_ensemble.py
index 2da22a87128..6030729387f 100644
--- a/test/rb/test_ensemble.py
+++ b/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):
diff --git a/test/rb/test_samplers.py b/test/rb/test_samplers.py
index 2e410f160a6..9abb6f14581 100644
--- a/test/rb/test_samplers.py
+++ b/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"
diff --git a/test/rb/test_storages.py b/test/rb/test_storages.py
index 228733e8a8f..76932a87cb0 100644
--- a/test/rb/test_storages.py
+++ b/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"
diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index 1a983b57004..ad89bfc7bb9 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -404,97 +404,116 @@ def test_roundrobin_dumps_loads_write_count(self, tmp_path):
 
 
 class TestWriterGeneration:
-    def test_generation_disabled_by_default(self):
+    def test_default_writer_tracks_generations(self):
         rb = ReplayBuffer(storage=LazyTensorStorage(10))
-        assert rb._writer.supports_generation is False
-        assert rb._writer._generation is None
-        with pytest.raises(RuntimeError, match="does not track generations"):
-            rb._writer._get_generation(0)
+        assert rb._writer.tracks_generations is True
+        index = rb.extend(torch.arange(10))
+        gen = rb._writer.generations_of(index)
+        assert gen.dtype == torch.int64
+        assert gen.shape == index.shape
+        assert (gen == 1).all()
+
+    def test_non_tracking_writer_reports_minus_one(self):
+        writer = TensorDictMaxValueWriter(rank_key="key")
+        assert writer.tracks_generations is False
+        gen = writer.generations_of(torch.arange(4))
+        torch.testing.assert_close(gen, torch.full((4,), -1, dtype=torch.int64))
 
     def test_generation_increments_on_reuse(self):
         size = 4
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
-        )
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(torch.arange(size))
-        assert rb._writer.supports_generation is True
-        assert (rb._writer._generation == 1).all()
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size)),
+            torch.ones(size, dtype=torch.int64),
+        )
         rb.extend(torch.arange(size, size + 3))
-        torch.testing.assert_close(rb._writer._generation, torch.tensor([2, 2, 2, 1]))
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size)), torch.tensor([2, 2, 2, 1])
+        )
 
     def test_generation_wraparound(self):
         size = 5
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
-        )
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(torch.arange(2 * size))
-        assert (rb._writer._generation == 2).all()
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size)),
+            torch.full((size,), 2, dtype=torch.int64),
+        )
 
     def test_generation_add(self):
         size = 3
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
-        )
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
         for i in range(size + 1):
             rb.add(torch.tensor(i))
-        torch.testing.assert_close(rb._writer._generation, torch.tensor([2, 1, 1]))
-
-    def test_generation_get(self):
-        size = 4
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size)), torch.tensor([2, 1, 1])
         )
-        rb.extend(torch.arange(size + 2))
-        gen = rb._writer._get_generation(torch.tensor([0, 3]))
-        torch.testing.assert_close(gen, torch.tensor([2, 1]))
+
+    def test_generations_of_returns_zero_for_unwritten(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(4))
+        rb.extend(torch.arange(2))
+        gen = rb._writer.generations_of(torch.arange(4))
+        torch.testing.assert_close(gen, torch.tensor([1, 1, 0, 0]))
 
     def test_generation_tensordict_writer(self):
         size = 4
-        rb = TensorDictReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=TensorDictRoundRobinWriter(track_generation=True),
-        )
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(TensorDict({"a": torch.arange(2 * size)}, [2 * size]))
-        assert (rb._writer._generation == 2).all()
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size)),
+            torch.full((size,), 2, dtype=torch.int64),
+        )
 
     def test_generation_write_at(self):
         storage = LazyTensorStorage(4)
-        writer = RoundRobinWriter(track_generation=True)
+        writer = RoundRobinWriter()
         writer.register_storage(storage)
         writer.extend(torch.arange(4))
         writer.write_at(torch.tensor([0, 1]), torch.tensor([10, 11]))
-        torch.testing.assert_close(writer._generation, torch.tensor([2, 2, 1, 1]))
+        torch.testing.assert_close(
+            writer.generations_of(torch.arange(4)), torch.tensor([2, 2, 1, 1])
+        )
 
-    def test_generation_empty_resets(self):
-        storage = LazyTensorStorage(4)
-        writer = RoundRobinWriter(track_generation=True)
-        writer.register_storage(storage)
-        writer.extend(torch.arange(4))
-        writer._empty()
-        assert writer._generation is None
+    def test_empty_is_monotonic(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        index = rb.extend(torch.arange(10))
+        before = rb._writer.generations_of(index)
+        rb.empty()
+        rb.extend(torch.arange(10))
+        after = rb._writer.generations_of(index)
+        assert (after > before).all()
+
+    def test_empty_invalidates_handles_immediately(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        index = rb.extend(torch.arange(10))
+        gen = rb._writer.generations_of(index)
+        rb.empty()
+        assert (rb._writer.generations_of(index) != gen).all()
 
     def test_generation_state_dict_roundtrip(self):
         size = 4
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
-        )
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(torch.arange(size + 1))
         sd = rb.state_dict()
-        rb2 = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
+        rb2 = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb2.load_state_dict(sd)
+        torch.testing.assert_close(
+            rb2._writer.generations_of(torch.arange(size)),
+            rb._writer.generations_of(torch.arange(size)),
         )
+
+    def test_legacy_state_dict_without_generation_loads(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(torch.arange(5))
+        sd = rb.state_dict()
+        del sd["_writer"]["_generation"]
+        rb2 = ReplayBuffer(storage=LazyTensorStorage(10))
         rb2.load_state_dict(sd)
-        assert rb2._writer.supports_generation is True
-        torch.testing.assert_close(rb2._writer._generation, rb._writer._generation)
+        assert rb2._writer._cursor == 5
 
     def test_generation_dumps_loads(self, tmp_path):
-        writer = RoundRobinWriter(track_generation=True)
+        writer = RoundRobinWriter()
         writer._cursor = 2
         writer._write_count = 9
         writer._generation = torch.tensor([3, 2, 2, 1])
@@ -503,35 +522,33 @@ def test_generation_dumps_loads(self, tmp_path):
         writer2.loads(tmp_path)
         assert writer2._cursor == 2
         assert writer2._write_count == 9
-        assert writer2._track_generation is True
-        torch.testing.assert_close(writer2._generation, torch.tensor([3, 2, 2, 1]))
+        torch.testing.assert_close(
+            writer2.generations_of(torch.arange(4)), torch.tensor([3, 2, 2, 1])
+        )
 
     def test_sample_returns_generation(self):
         size = 8
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
-        )
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(torch.arange(size))
         _, info = rb.sample(4, return_info=True)
         assert "index_generation" in info
         gen = torch.as_tensor(info["index_generation"])
         idx = torch.as_tensor(info["index"])
         assert gen.shape == idx.shape
-        torch.testing.assert_close(gen, rb._writer._generation[idx])
+        torch.testing.assert_close(gen, rb._writer.generations_of(idx))
 
-    def test_sample_no_generation_when_disabled(self):
-        rb = ReplayBuffer(storage=LazyTensorStorage(8))
-        rb.extend(torch.arange(8))
+    def test_non_tracking_sample_has_no_generation(self):
+        rb = TensorDictReplayBuffer(
+            storage=LazyTensorStorage(10),
+            writer=TensorDictMaxValueWriter(rank_key="key"),
+        )
+        rb.extend(TensorDict({"key": torch.arange(10), "a": torch.arange(10)}, [10]))
         _, info = rb.sample(4, return_info=True)
         assert "index_generation" not in info
 
     def test_tensordict_sample_has_generation_key(self):
         size = 8
-        rb = TensorDictReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=TensorDictRoundRobinWriter(track_generation=True),
-        )
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(TensorDict({"a": torch.arange(size)}, [size]))
         sample = rb.sample(4)
         assert "index_generation" in sample.keys()
@@ -539,30 +556,24 @@ def test_tensordict_sample_has_generation_key(self):
 
     def test_wraparound_race_detectable(self):
         size = 8
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
-        )
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(torch.arange(size))
         _, info = rb.sample(4, return_info=True)
         sampled_index = torch.as_tensor(info["index"])
         sampled_generation = torch.as_tensor(info["index_generation"])
         rb.extend(torch.arange(size, 2 * size))
-        current = rb._writer._get_generation(sampled_index)
+        current = rb._writer.generations_of(sampled_index)
         assert (current != sampled_generation).all()
 
     def test_partial_reuse_detectable(self):
         size = 8
-        rb = ReplayBuffer(
-            storage=LazyTensorStorage(size),
-            writer=RoundRobinWriter(track_generation=True),
-        )
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
         rb.extend(torch.arange(size))
         _, info = rb.sample(size, return_info=True)
         idx = torch.as_tensor(info["index"])
         gen = torch.as_tensor(info["index_generation"])
         rb.extend(torch.arange(size, size + 3))
-        stale = rb._writer._get_generation(idx) != gen
+        stale = rb._writer.generations_of(idx) != gen
         torch.testing.assert_close(stale, idx < 3)
 
 
diff --git a/torchrl/data/replay_buffers/replay_buffers.py b/torchrl/data/replay_buffers/replay_buffers.py
index 6924a42a4a7..4251005123b 100644
--- a/torchrl/data/replay_buffers/replay_buffers.py
+++ b/torchrl/data/replay_buffers/replay_buffers.py
@@ -1534,8 +1534,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
-            if self._writer.supports_generation:
-                info["index_generation"] = self._writer._get_generation(index)
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2159,8 +2159,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
-            if self._writer.supports_generation:
-                info["index_generation"] = self._writer._get_generation(index)
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2638,8 +2638,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
-            if self._writer.supports_generation:
-                info["index_generation"] = self._writer._get_generation(index)
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -3009,8 +3009,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
             if self._sample_unit is not None:
                 index, info = self._sample_unit.expand(index, info, self._storage)
             info["index"] = index
-            if self._writer.supports_generation:
-                info["index_generation"] = self._writer._get_generation(index)
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(index)
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 8ce22373b37..00f6d677623 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -39,6 +39,10 @@ def tree_leaves(data):  # noqa: D103
 from torchrl.data.replay_buffers.storages import Storage
 from torchrl.data.replay_buffers.utils import _is_int, _reduce
 
+# Storage capacities at or above this value are treated as unbounded (lazy
+# storages report a sentinel max size), triggering dynamic generation growth.
+_GENERATION_UNBOUNDED = 2**40
+
 
 class Writer(ABC):
     """A ReplayBuffer base Writer class."""
@@ -50,13 +54,19 @@ def __init__(self, compilable: bool = False) -> None:
         self._storage = None
         self._compilable = compilable
 
+    #: Whether this writer type stamps storage slots with a reuse generation.
+    tracks_generations: bool = False
+
     def register_storage(self, storage: Storage) -> None:
         self._storage = storage
 
-    @property
-    def supports_generation(self) -> bool:
-        """Whether this writer maintains per-slot generation stamps."""
-        return False
+    def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
+        """Returns the generation stamp for each physical slot in ``index``.
+
+        Writers that do not track slot reuse report ``-1``.
+        """
+        index = torch.as_tensor(index)
+        return torch.full(index.shape, -1, dtype=torch.int64, device=index.device)
 
     @abstractmethod
     def add(self, data: Any) -> int:
@@ -153,77 +163,86 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 class RoundRobinWriter(Writer):
     """A RoundRobin Writer class for composable replay buffers.
 
+    Round-robin writers stamp every physical storage slot with an int64
+    generation counter, incremented once per write. Sampling returns the
+    generation observed for each index (see :meth:`generations_of`), so a
+    captured ``(index, generation)`` pair identifies one specific record and
+    becomes detectably stale once its slot is reused or the buffer is emptied.
+
     Args:
         compilable (bool, optional): whether the writer is compilable.
             If ``True``, the writer cannot be shared between multiple processes.
             Defaults to ``False``.
 
-    Keyword Args:
-        track_generation (bool, optional): if ``True``, the writer maintains a
-            per-slot generation counter that is incremented every time a physical
-            storage slot is written. Samplers can then stamp sampled indices with
-            the generation they observed, allowing delayed priority updates to
-            detect whether a slot has since been reused. Defaults to ``False``.
-
     """
 
-    def __init__(
-        self, compilable: bool = False, *, track_generation: bool = False
-    ) -> None:
+    tracks_generations: bool = True
+
+    def __init__(self, compilable: bool = False) -> None:
         super().__init__(compilable=compilable)
         self._cursor = 0
         self._write_count  # noqa
-        self._track_generation = track_generation
         self._generation = None
 
-    @property
-    def supports_generation(self) -> bool:
-        return self._track_generation
-
-    def _ensure_generation(self, size: int) -> None:
-        if self._generation is None:
-            generation = torch.zeros(int(size), dtype=torch.long)
-            if not self._compilable:
-                generation.share_memory_()
-            self._generation = generation
+    def _ensure_generation(self, capacity: int, min_size: int) -> None:
+        # Bounded storages allocate to capacity once (stable shape, so the
+        # ``torch.compile`` extend/sample path does not recompile); lazy storages
+        # report a sentinel capacity and instead grow geometrically.
+        generation = self._generation
+        current = 0 if generation is None else generation.numel()
+        if current >= min_size:
+            return
+        size = (
+            capacity if capacity < _GENERATION_UNBOUNDED else max(min_size, current * 2)
+        )
+        new_generation = torch.zeros(size, dtype=torch.int64)
+        if generation is not None:
+            new_generation[:current] = generation
+        if not self._compilable:
+            new_generation.share_memory_()
+        self._generation = new_generation
 
     def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
-        if not self._track_generation:
-            return
         if _is_int(index):
-            self._ensure_generation(
-                self._storage._max_size_along_dim0(single_data=data)
-            )
+            capacity = self._storage._max_size_along_dim0(single_data=data)
+            self._ensure_generation(capacity, int(index) + 1)
             self._generation[int(index)] += 1
         else:
-            self._ensure_generation(
-                self._storage._max_size_along_dim0(batched_data=data)
-            )
             index = torch.as_tensor(index, dtype=torch.long).reshape(-1)
+            if index.numel() == 0:
+                return
+            capacity = self._storage._max_size_along_dim0(batched_data=data)
+            min_size = (
+                capacity if capacity < _GENERATION_UNBOUNDED else int(index.max()) + 1
+            )
+            self._ensure_generation(capacity, min_size)
             self._generation.index_put_(
                 (index,), torch.ones_like(index), accumulate=True
             )
 
-    def _get_generation(self, index: int | torch.Tensor) -> torch.Tensor:
-        if not self._track_generation:
-            raise RuntimeError(
-                "This writer does not track generations. Construct it with "
-                "track_generation=True to use generation-stamped indices."
-            )
-        if self._generation is None:
-            raise RuntimeError("No data has been written yet; generation is undefined.")
+    def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
+        """Returns the int64 generation stamp for each physical slot in ``index``.
+
+        For multidimensional storages the index is reduced to its first (slot)
+        dimension. Slots that have never been written report ``0``.
+        """
         if isinstance(index, tuple):
             index = index[0]
         elif (
             isinstance(index, torch.Tensor)
             and index.ndim
+            and self._storage is not None
             and self._storage.ndim > 1
             and index.shape[-1] == self._storage.ndim
         ):
             index = index[..., 0]
-        if _is_int(index):
-            return self._generation[int(index)]
-        return self._generation[torch.as_tensor(index, dtype=torch.long)]
+        index = torch.as_tensor(index, dtype=torch.long)
+        if self._generation is None:
+            return torch.zeros(index.shape, dtype=torch.int64, device=index.device)
+        idx = index.to(self._generation.device)
+        n = self._generation.numel()
+        gen = self._generation[idx.clamp(max=n - 1)]
+        return torch.where(idx < n, gen, torch.zeros_like(gen))
 
     def dumps(self, path):
         path = Path(path).absolute()
@@ -231,7 +250,6 @@ def dumps(self, path):
         metadata = {
             "cursor": self._cursor,
             "write_count": self._write_count,
-            "track_generation": self._track_generation,
         }
         generation = self._generation
         if generation is not None:
@@ -258,9 +276,6 @@ def loads(self, path):
             write_count = metadata.get("write_count")
             if write_count is not None:
                 self._write_count = write_count
-            self._track_generation = metadata.get(
-                "track_generation", self._track_generation
-            )
             generation_shape = metadata.get("generation_shape")
         if generation_shape is not None:
             generation = MemoryMappedTensor.from_filename(
@@ -350,10 +365,8 @@ def _update_storage_len_for_write_at(self, index: int | torch.Tensor) -> None:
 
     def state_dict(self) -> dict[str, Any]:
         state_dict = {"_cursor": self._cursor, "_write_count": self._write_count}
-        if self._track_generation:
-            state_dict["_track_generation"] = True
-            if self._generation is not None:
-                state_dict["_generation"] = self._generation.clone()
+        if self._generation is not None:
+            state_dict["_generation"] = self._generation.clone()
         return state_dict
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
@@ -361,9 +374,6 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         write_count = state_dict.get("_write_count")
         if write_count is not None:
             self._write_count = write_count
-        self._track_generation = state_dict.get(
-            "_track_generation", self._track_generation
-        )
         generation = state_dict.get("_generation")
         if generation is not None:
             generation = generation.clone()
@@ -373,7 +383,8 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 
     def _empty(self, empty_write_count: bool = True) -> None:
         self._cursor = 0
-        self._generation = None
+        if self._generation is not None:
+            self._generation += 1
         if empty_write_count:
             self._write_count = 0
 

From 8125e23c1456877e06184288a1b5258361432a47 Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Fri, 24 Jul 2026 14:39:19 +0530
Subject: [PATCH 04/13] feat: align slot generation to 0-based reuse counter

---
 test/rb/test_writers.py                | 18 +++++++++---------
 torchrl/data/replay_buffers/writers.py | 20 +++++++++++---------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index ad89bfc7bb9..86401d17f33 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -411,7 +411,7 @@ def test_default_writer_tracks_generations(self):
         gen = rb._writer.generations_of(index)
         assert gen.dtype == torch.int64
         assert gen.shape == index.shape
-        assert (gen == 1).all()
+        assert (gen == 0).all()
 
     def test_non_tracking_writer_reports_minus_one(self):
         writer = TensorDictMaxValueWriter(rank_key="key")
@@ -425,11 +425,11 @@ def test_generation_increments_on_reuse(self):
         rb.extend(torch.arange(size))
         torch.testing.assert_close(
             rb._writer.generations_of(torch.arange(size)),
-            torch.ones(size, dtype=torch.int64),
+            torch.zeros(size, dtype=torch.int64),
         )
         rb.extend(torch.arange(size, size + 3))
         torch.testing.assert_close(
-            rb._writer.generations_of(torch.arange(size)), torch.tensor([2, 2, 2, 1])
+            rb._writer.generations_of(torch.arange(size)), torch.tensor([1, 1, 1, 0])
         )
 
     def test_generation_wraparound(self):
@@ -438,7 +438,7 @@ def test_generation_wraparound(self):
         rb.extend(torch.arange(2 * size))
         torch.testing.assert_close(
             rb._writer.generations_of(torch.arange(size)),
-            torch.full((size,), 2, dtype=torch.int64),
+            torch.full((size,), 1, dtype=torch.int64),
         )
 
     def test_generation_add(self):
@@ -447,14 +447,14 @@ def test_generation_add(self):
         for i in range(size + 1):
             rb.add(torch.tensor(i))
         torch.testing.assert_close(
-            rb._writer.generations_of(torch.arange(size)), torch.tensor([2, 1, 1])
+            rb._writer.generations_of(torch.arange(size)), torch.tensor([1, 0, 0])
         )
 
-    def test_generations_of_returns_zero_for_unwritten(self):
+    def test_generations_of_unwritten_reports_minus_one(self):
         rb = ReplayBuffer(storage=LazyTensorStorage(4))
         rb.extend(torch.arange(2))
         gen = rb._writer.generations_of(torch.arange(4))
-        torch.testing.assert_close(gen, torch.tensor([1, 1, 0, 0]))
+        torch.testing.assert_close(gen, torch.tensor([0, 0, -1, -1]))
 
     def test_generation_tensordict_writer(self):
         size = 4
@@ -462,7 +462,7 @@ def test_generation_tensordict_writer(self):
         rb.extend(TensorDict({"a": torch.arange(2 * size)}, [2 * size]))
         torch.testing.assert_close(
             rb._writer.generations_of(torch.arange(size)),
-            torch.full((size,), 2, dtype=torch.int64),
+            torch.full((size,), 1, dtype=torch.int64),
         )
 
     def test_generation_write_at(self):
@@ -472,7 +472,7 @@ def test_generation_write_at(self):
         writer.extend(torch.arange(4))
         writer.write_at(torch.tensor([0, 1]), torch.tensor([10, 11]))
         torch.testing.assert_close(
-            writer.generations_of(torch.arange(4)), torch.tensor([2, 2, 1, 1])
+            writer.generations_of(torch.arange(4)), torch.tensor([1, 1, 0, 0])
         )
 
     def test_empty_is_monotonic(self):
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 00f6d677623..8b1bc64e2e4 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -164,10 +164,11 @@ class RoundRobinWriter(Writer):
     """A RoundRobin Writer class for composable replay buffers.
 
     Round-robin writers stamp every physical storage slot with an int64
-    generation counter, incremented once per write. Sampling returns the
-    generation observed for each index (see :meth:`generations_of`), so a
-    captured ``(index, generation)`` pair identifies one specific record and
-    becomes detectably stale once its slot is reused or the buffer is emptied.
+    generation counter: the first write of a slot has generation ``0`` and every
+    reuse increments it. Sampling returns the generation observed for each index
+    (see :meth:`generations_of`), so a captured ``(index, generation)`` pair
+    identifies one specific record and becomes detectably stale once its slot is
+    reused or the buffer is emptied.
 
     Args:
         compilable (bool, optional): whether the writer is compilable.
@@ -195,7 +196,7 @@ def _ensure_generation(self, capacity: int, min_size: int) -> None:
         size = (
             capacity if capacity < _GENERATION_UNBOUNDED else max(min_size, current * 2)
         )
-        new_generation = torch.zeros(size, dtype=torch.int64)
+        new_generation = torch.full((size,), -1, dtype=torch.int64)
         if generation is not None:
             new_generation[:current] = generation
         if not self._compilable:
@@ -223,8 +224,9 @@ def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
     def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
         """Returns the int64 generation stamp for each physical slot in ``index``.
 
-        For multidimensional storages the index is reduced to its first (slot)
-        dimension. Slots that have never been written report ``0``.
+        The first write of a slot has generation ``0`` and every reuse increments
+        it. For multidimensional storages the index is reduced to its first (slot)
+        dimension. Slots that have never been written report ``-1``.
         """
         if isinstance(index, tuple):
             index = index[0]
@@ -238,11 +240,11 @@ def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
             index = index[..., 0]
         index = torch.as_tensor(index, dtype=torch.long)
         if self._generation is None:
-            return torch.zeros(index.shape, dtype=torch.int64, device=index.device)
+            return torch.full(index.shape, -1, dtype=torch.int64, device=index.device)
         idx = index.to(self._generation.device)
         n = self._generation.numel()
         gen = self._generation[idx.clamp(max=n - 1)]
-        return torch.where(idx < n, gen, torch.zeros_like(gen))
+        return torch.where(idx < n, gen, torch.full_like(gen, -1))
 
     def dumps(self, path):
         path = Path(path).absolute()

From bb3af3d8252359f9f0877eca8672f33071b4b1ca Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Fri, 24 Jul 2026 21:00:16 +0530
Subject: [PATCH 05/13] docs: drop verbose generation docstrings from
 RoundRobinWriter

---
 torchrl/data/replay_buffers/writers.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 8b1bc64e2e4..a7adbcc2da2 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -163,13 +163,6 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 class RoundRobinWriter(Writer):
     """A RoundRobin Writer class for composable replay buffers.
 
-    Round-robin writers stamp every physical storage slot with an int64
-    generation counter: the first write of a slot has generation ``0`` and every
-    reuse increments it. Sampling returns the generation observed for each index
-    (see :meth:`generations_of`), so a captured ``(index, generation)`` pair
-    identifies one specific record and becomes detectably stale once its slot is
-    reused or the buffer is emptied.
-
     Args:
         compilable (bool, optional): whether the writer is compilable.
             If ``True``, the writer cannot be shared between multiple processes.
@@ -222,12 +215,6 @@ def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
             )
 
     def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
-        """Returns the int64 generation stamp for each physical slot in ``index``.
-
-        The first write of a slot has generation ``0`` and every reuse increments
-        it. For multidimensional storages the index is reduced to its first (slot)
-        dimension. Slots that have never been written report ``-1``.
-        """
         if isinstance(index, tuple):
             index = index[0]
         elif (

From ff4fb4ac7428d0e87783c20c16339e3d3828edba Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Sun, 26 Jul 2026 01:59:17 +0530
Subject: [PATCH 06/13] fix: allocate writer generation on the storage device

---
 test/rb/test_writers.py                | 86 ++++++++++++++++++++++++++
 torchrl/data/replay_buffers/writers.py | 49 ++++++++++++---
 2 files changed, 127 insertions(+), 8 deletions(-)

diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index 86401d17f33..d475478d147 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -576,6 +576,92 @@ def test_partial_reuse_detectable(self):
         stale = rb._writer.generations_of(idx) != gen
         torch.testing.assert_close(stale, idx < 3)
 
+    @pytest.mark.parametrize("device", get_default_devices())
+    def test_generation_on_storage_device(self, device):
+        size = 8
+        rb = ReplayBuffer(storage=LazyTensorStorage(size, device=device))
+        rb.extend(torch.arange(size, device=device))
+        assert rb._writer._generation.device.type == device.type
+        _, info = rb.sample(4, return_info=True)
+        gen = info["index_generation"]
+        idx = torch.as_tensor(info["index"])
+        assert gen.device == idx.device
+        torch.testing.assert_close(gen, rb._writer.generations_of(idx))
+        rb.extend(torch.arange(size, 2 * size, device=device))
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size, device=device)),
+            torch.ones(size, dtype=torch.int64, device=device),
+        )
+
+    @pytest.mark.parametrize("device", get_default_devices())
+    def test_generation_add_on_storage_device(self, device):
+        size = 3
+        rb = ReplayBuffer(storage=LazyTensorStorage(size, device=device))
+        for i in range(size + 1):
+            rb.add(torch.tensor(i, device=device))
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size, device=device)),
+            torch.tensor([1, 0, 0], device=device),
+        )
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+    def test_generation_cuda_data_into_cuda_storage(self):
+        size = 8
+        rb = ReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb.extend(torch.arange(size, device="cuda"))
+        assert rb._writer._generation.device.type == "cuda"
+        _, info = rb.sample(4, return_info=True)
+        idx = torch.as_tensor(info["index"])
+        assert info["index_generation"].device == idx.device
+        rb.extend(torch.arange(size, 2 * size, device="cuda"))
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size, device="cuda")),
+            torch.ones(size, dtype=torch.int64, device="cuda"),
+        )
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+    def test_generation_cpu_data_into_cuda_storage(self):
+        size = 4
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb.extend(TensorDict({"a": torch.arange(2 * size)}, [2 * size]))
+        assert rb._writer._generation.device.type == "cuda"
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size)),
+            torch.ones(size, dtype=torch.int64),
+        )
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+    def test_generation_state_dict_roundtrip_cuda(self):
+        size = 4
+        rb = ReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb.extend(torch.arange(size + 1, device="cuda"))
+        rb2 = ReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb2.load_state_dict(rb.state_dict())
+        index = torch.arange(size, device="cuda")
+        assert rb2._writer._generation.device.type == "cuda"
+        torch.testing.assert_close(
+            rb2._writer.generations_of(index), rb._writer.generations_of(index)
+        )
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+    def test_generation_dumps_loads_cuda(self, tmp_path):
+        writer = RoundRobinWriter()
+        writer.register_storage(LazyTensorStorage(4, device="cuda"))
+        writer._generation = torch.tensor([3, 2, 2, 1], device="cuda")
+        writer.dumps(tmp_path)
+        writer2 = RoundRobinWriter()
+        writer2.register_storage(LazyTensorStorage(4, device="cuda"))
+        writer2.loads(tmp_path)
+        assert writer2._generation.device.type == "cuda"
+        torch.testing.assert_close(
+            writer2.generations_of(torch.arange(4, device="cuda")),
+            torch.tensor([3, 2, 2, 1], device="cuda"),
+        )
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index a7adbcc2da2..44af2777bbb 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -19,7 +19,7 @@
 from tensordict import is_tensor_collection, MemoryMappedTensor, TensorDictBase
 from tensordict.utils import expand_as_right, is_tensorclass
 from torch import multiprocessing as mp
-from torchrl._utils import _STRDTYPE2DTYPE
+from torchrl._utils import _make_ordinal_device, _STRDTYPE2DTYPE
 
 try:
     from torch.compiler import disable as compile_disable
@@ -178,28 +178,56 @@ def __init__(self, compilable: bool = False) -> None:
         self._write_count  # noqa
         self._generation = None
 
-    def _ensure_generation(self, capacity: int, min_size: int) -> None:
+    def register_storage(self, storage: Storage) -> None:
+        super().register_storage(storage)
+        self._align_generation_device()
+
+    def _generation_device(self, index: int | torch.Tensor) -> torch.device:
+        # The generation buffer follows the storage: sampled indices are built on
+        # the storage device, so lookups stay sync-free on the sampling path.
+        device = getattr(self._storage, "device", None)
+        if device is None or device == "auto":
+            if isinstance(index, torch.Tensor):
+                return _make_ordinal_device(index.device)
+            return torch.device("cpu")
+        return _make_ordinal_device(torch.device(device))
+
+    def _align_generation_device(self) -> None:
+        generation = self._generation
+        if generation is None:
+            return
+        device = self._generation_device(generation)
+        if generation.device != device:
+            self._generation = generation.to(device)
+
+    def _ensure_generation(
+        self, capacity: int, min_size: int, device: torch.device
+    ) -> None:
         # Bounded storages allocate to capacity once (stable shape, so the
         # ``torch.compile`` extend/sample path does not recompile); lazy storages
         # report a sentinel capacity and instead grow geometrically.
         generation = self._generation
+        if generation is not None and generation.device != device:
+            generation = generation.to(device)
+            self._generation = generation
         current = 0 if generation is None else generation.numel()
         if current >= min_size:
             return
         size = (
             capacity if capacity < _GENERATION_UNBOUNDED else max(min_size, current * 2)
         )
-        new_generation = torch.full((size,), -1, dtype=torch.int64)
+        new_generation = torch.full((size,), -1, dtype=torch.int64, device=device)
         if generation is not None:
             new_generation[:current] = generation
-        if not self._compilable:
+        if not self._compilable and new_generation.device.type == "cpu":
             new_generation.share_memory_()
         self._generation = new_generation
 
     def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
+        device = self._generation_device(index)
         if _is_int(index):
             capacity = self._storage._max_size_along_dim0(single_data=data)
-            self._ensure_generation(capacity, int(index) + 1)
+            self._ensure_generation(capacity, int(index) + 1, device)
             self._generation[int(index)] += 1
         else:
             index = torch.as_tensor(index, dtype=torch.long).reshape(-1)
@@ -209,7 +237,8 @@ def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
             min_size = (
                 capacity if capacity < _GENERATION_UNBOUNDED else int(index.max()) + 1
             )
-            self._ensure_generation(capacity, min_size)
+            self._ensure_generation(capacity, min_size, device)
+            index = index.to(device)
             self._generation.index_put_(
                 (index,), torch.ones_like(index), accumulate=True
             )
@@ -231,7 +260,8 @@ def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
         idx = index.to(self._generation.device)
         n = self._generation.numel()
         gen = self._generation[idx.clamp(max=n - 1)]
-        return torch.where(idx < n, gen, torch.full_like(gen, -1))
+        gen = torch.where(idx < n, gen, torch.full_like(gen, -1))
+        return gen.to(index.device)
 
     def dumps(self, path):
         path = Path(path).absolute()
@@ -242,6 +272,7 @@ def dumps(self, path):
         }
         generation = self._generation
         if generation is not None:
+            generation = generation.cpu()
             try:
                 MemoryMappedTensor.from_filename(
                     filename=path / "generation.memmap",
@@ -275,6 +306,7 @@ def loads(self, path):
             if not self._compilable:
                 generation.share_memory_()
             self._generation = generation
+            self._align_generation_device()
 
     def add(self, data: Any) -> int | torch.Tensor:
         index = self._cursor
@@ -366,9 +398,10 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         generation = state_dict.get("_generation")
         if generation is not None:
             generation = generation.clone()
-            if not self._compilable:
+            if not self._compilable and generation.device.type == "cpu":
                 generation.share_memory_()
             self._generation = generation
+            self._align_generation_device()
 
     def _empty(self, empty_write_count: bool = True) -> None:
         self._cursor = 0

From 164108daa25263bbca0e1dd22677b6694a22713c Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Sun, 26 Jul 2026 01:59:28 +0530
Subject: [PATCH 07/13] fix: keep unwritten generation slots at -1 on empty

---
 test/rb/test_writers.py                | 8 ++++++++
 torchrl/data/replay_buffers/writers.py | 6 ++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index d475478d147..5095e4f284c 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -491,6 +491,14 @@ def test_empty_invalidates_handles_immediately(self):
         rb.empty()
         assert (rb._writer.generations_of(index) != gen).all()
 
+    def test_empty_preserves_unwritten_sentinel(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(4))
+        rb.extend(torch.arange(2))
+        rb.empty()
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(4)), torch.tensor([1, 1, -1, -1])
+        )
+
     def test_generation_state_dict_roundtrip(self):
         size = 4
         rb = ReplayBuffer(storage=LazyTensorStorage(size))
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 44af2777bbb..a3d9b3612fe 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -405,8 +405,10 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 
     def _empty(self, empty_write_count: bool = True) -> None:
         self._cursor = 0
-        if self._generation is not None:
-            self._generation += 1
+        generation = self._generation
+        if generation is not None:
+            # never-written slots keep the -1 sentinel
+            generation[generation >= 0] += 1
         if empty_write_count:
             self._write_count = 0
 

From d4b48141735789f2aaf80ad51d7de5ead63e8576 Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Sun, 26 Jul 2026 12:30:33 +0530
Subject: [PATCH 08/13] docs: document generation bumps on writes

---
 torchrl/data/replay_buffers/writers.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index a3d9b3612fe..59932b42773 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -63,7 +63,9 @@ def register_storage(self, storage: Storage) -> None:
     def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
         """Returns the generation stamp for each physical slot in ``index``.
 
-        Writers that do not track slot reuse report ``-1``.
+        The stamp advances once per write to that slot, so a single ``extend``
+        that wraps the storage advances a reused slot once per write it
+        receives. Writers that do not track slot reuse report ``-1``.
         """
         index = torch.as_tensor(index)
         return torch.full(index.shape, -1, dtype=torch.int64, device=index.device)
@@ -354,7 +356,11 @@ def extend(self, data: Sequence) -> torch.Tensor:
         return index
 
     def write_at(self, index: int | torch.Tensor, data: Any) -> int | torch.Tensor:
-        """Writes data at explicit storage indices without moving the cursor."""
+        """Writes data at explicit storage indices without moving the cursor.
+
+        The generation of every written slot is bumped, so handles previously
+        handed out for those slots are stale once this returns.
+        """
         if _is_int(index):
             batch_size = 1
         else:

From 6b2d2805be9355de7a07fd48c2c93c5122cab1e3 Mon Sep 17 00:00:00 2001
From: harryfrzz <harikrishnac005@gmail.com>
Date: Sun, 26 Jul 2026 12:30:33 +0530
Subject: [PATCH 09/13] test: cover generation bump on wrapping extend

---
 test/rb/test_writers.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index 5095e4f284c..1b218471311 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -441,6 +441,15 @@ def test_generation_wraparound(self):
             torch.full((size,), 1, dtype=torch.int64),
         )
 
+    def test_generation_extend_wrapping_twice(self):
+        size = 4
+        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        # slots 0 and 1 are written three times, slots 2 and 3 twice
+        rb.extend(torch.arange(2 * size + 2))
+        torch.testing.assert_close(
+            rb._writer.generations_of(torch.arange(size)), torch.tensor([2, 2, 1, 1])
+        )
+
     def test_generation_add(self):
         size = 3
         rb = ReplayBuffer(storage=LazyTensorStorage(size))

From 5bef5f7897dd83201a9eab18f19057fc3615f344 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vincentmoens@gmail.com>
Date: Tue, 4 Aug 2026 17:01:04 +0100
Subject: [PATCH 10/13] [BugFix] Generation stamps: make opt-in, move state to
 the storage, document

Addresses review feedback.

- Opt-in. `RoundRobinWriter(track_generations=True)` now gates the feature;
  `tracks_generations` is an instance property rather than a class attribute set
  to True. Previously every replay buffer allocated an int64 per slot and every
  `sample()` grew an `"index_generation"` entry -- which for tensordict buffers
  is a new sample key. The four pre-existing tests this PR had to patch to
  `exclude("index", "index_generation")` are reverted, which is the check that
  the BC break is gone.

- The counters live on the storage, not on the writer. Two buffers sharing one
  storage overwrite each other's slots, so a per-writer counter let buffer A's
  handles read as live after buffer B overwrote them -- defeating the one thing
  the feature is for. The buffer is attached to the storage; a writer registered
  against a storage that already has one adopts it. A storage-less writer (as
  `dumps`/`loads` and `load_state_dict` can be) keeps it locally until a storage
  is registered.

- Dropped the `_GENERATION_UNBOUNDED = 2**40` sentinel. Rather than trying to
  classify a capacity as bounded or not, allocate eagerly below an explicit
  allocation limit (stable shape, no torch.compile recompile) and grow
  geometrically above it. Correct for genuinely unbounded storages
  (`torch.iinfo(torch.int64).max`) and for implausibly large bounded ones alike.

- Dropped the `share_memory_()` calls. The buffer is replaced, not mutated, when
  it grows, so a shared mapping silently stops tracking after the first growth;
  cross-process visibility was never actually supported. Now documented as
  process-local instead of half-implemented.

- `generations_of` no longer guesses that a 1-D tensor whose length matches
  `storage.ndim` is a coordinate vector: with `ndim == 2`,
  `generations_of(tensor([3, 7]))` silently returned one generation where the
  caller asked for two. Only `ndim >= 2` tensors are read as coordinate batches;
  pass a tuple for a single cell.

- Documentation. A "Detecting overwritten slots: generation stamps" section in
  data_replaybuffers.rst covering why slot indices go stale, the semantics
  (one stamp per write, `-1` means no usable stamp, monotonic across `empty()`,
  not an ordering across slots) and the implementation notes (opt-in, storage-
  owned, allocation policy, process-local, multidim indexing, checkpointing),
  with a runnable example. Cross-referenced from the Writers section, and
  `Writer.generations_of` / `tracks_generations` / `RoundRobinWriter` carry full
  docstrings.

- Tests for each: opt-in default leaves `sample()` untouched, shared-storage
  visibility, and the multidim 1-D index case.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 docs/source/reference/data_replaybuffers.rst |  89 ++++++++++
 docs/source/reference/data_samplers.rst      |   5 +-
 test/rb/test_ensemble.py                     |   2 +-
 test/rb/test_samplers.py                     |   4 +-
 test/rb/test_storages.py                     |   2 +-
 test/rb/test_writers.py                      | 165 +++++++++++++++---
 torchrl/data/replay_buffers/writers.py       | 172 +++++++++++++++----
 7 files changed, 377 insertions(+), 62 deletions(-)

diff --git a/docs/source/reference/data_replaybuffers.rst b/docs/source/reference/data_replaybuffers.rst
index 68c2473923d..50c4e5e3d2c 100644
--- a/docs/source/reference/data_replaybuffers.rst
+++ b/docs/source/reference/data_replaybuffers.rst
@@ -242,6 +242,95 @@ capacity without scanning the full storage on every write. This mode supports
 random sampling. Prefetching, prioritized replay and multidimensional storages
 are rejected explicitly.
 
+Detecting overwritten slots: generation stamps
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. _ref_buffers_generations:
+
+A replay buffer index is a *physical slot number*, not a handle on a piece of
+data. A round-robin writer reuses slots, so an index sampled at one point in
+time may name completely different data a moment later. That matters whenever
+something outside the buffer holds an index across a write:
+
+- asynchronous training, where an inference worker samples, computes, and only
+  then writes results back at the index it was given;
+- prioritized replay, where priorities are updated after the forward pass;
+- any conditional write ("update this record only if it is still the one I
+  read").
+
+Generation stamps make that staleness detectable. With
+``track_generations=True``, the writer keeps one counter per storage slot and
+advances it on every write to that slot. Comparing the stamp you captured
+against the current stamp answers "is this still my data?":
+
+    >>> import torch
+    >>> from torchrl.data import LazyTensorStorage, ReplayBuffer
+    >>> from torchrl.data import RoundRobinWriter
+    >>> rb = ReplayBuffer(
+    ...     storage=LazyTensorStorage(8),
+    ...     writer=RoundRobinWriter(track_generations=True),
+    ... )
+    >>> _ = rb.extend(torch.arange(8))
+    >>> _, info = rb.sample(4, return_info=True)
+    >>> index, generation = info["index"], info["index_generation"]
+    >>> _ = rb.extend(torch.arange(8, 11))   # overwrites slots 0, 1, 2
+    >>> stale = rb.writer.generations_of(index) != generation
+    >>> # `index[stale]` no longer holds the sampled data
+
+:meth:`~torchrl.data.ReplayBuffer.sample` adds ``"index_generation"`` to its
+``info`` (and, for tensordict buffers, to the sample itself) whenever the writer
+tracks generations, alongside the existing ``"index"``.
+
+Semantics
+^^^^^^^^^
+
+- **One stamp per write, not per ``extend`` call.** A single ``extend`` that
+  wraps the storage advances a reused slot once for each write it receives, so
+  a slot written twice in one call advances by two.
+- **``-1`` means "no usable stamp"**: a never-written slot, an
+  out-of-range index, or a writer that does not track generations. It is not
+  "generation zero".
+- **Monotonic across** :meth:`~torchrl.data.ReplayBuffer.empty`. Emptying
+  advances every written slot's stamp rather than resetting it, so handles taken
+  before the ``empty()`` correctly read as stale. Never-written slots keep
+  ``-1``.
+- **Stamps are for detection, not for ordering across slots.** Two slots'
+  stamps are independent counters; a higher stamp on slot 3 than on slot 7 says
+  nothing about write order between them.
+
+Implementation notes
+^^^^^^^^^^^^^^^^^^^^
+
+- **Opt-in.** The default is ``track_generations=False``: enabling it allocates
+  one ``int64`` per storage slot and adds a key to the sampler output, neither
+  of which should be imposed on buffers that do not need it.
+- **The counters live on the storage, not on the writer.** Two buffers sharing
+  one storage overwrite each other's slots, so a per-writer counter would let
+  one buffer's handles read as live after the other overwrote them. The buffer
+  is attached to the storage object, and a writer registered against a storage
+  that already has one adopts it rather than replacing it.
+- **Allocation.** Storages small enough to allocate up front get a single
+  allocation, so the buffer's shape never changes and the ``torch.compile``
+  extend/sample path does not recompile. Larger and unbounded storages
+  (``ListStorage`` with no ``max_size`` reports ``torch.iinfo(torch.int64).max``)
+  grow geometrically on demand instead.
+- **Process-local.** The counters are not shared across processes: the buffer is
+  replaced rather than mutated when it grows, so a shared mapping would silently
+  stop tracking after the first growth. A slot overwritten by another process is
+  not reflected. Cross-process staleness detection needs a storage-owned,
+  fixed-size mapping and is not implemented yet.
+- **Multidimensional storages.** A generation stamps a whole dim-0 slot. A 1-D
+  index tensor is therefore always read as a batch of slot indices; to identify
+  a single cell of an ``ndim > 1`` storage, pass the ``tuple`` of per-dimension
+  indices that :meth:`~torchrl.data.ReplayBuffer.extend` returns.
+- **Checkpointing.** Stamps are part of ``state_dict``/``dumps`` when tracking
+  is on, and a checkpoint written without them (or by an older version) loads
+  fine -- tracking simply starts from scratch.
+
+The relevant APIs are :attr:`~torchrl.data.Writer.tracks_generations` and
+:meth:`~torchrl.data.Writer.generations_of`, and the ``track_generations``
+argument of :class:`~torchrl.data.RoundRobinWriter`.
+
 Trajectory boundaries
 ~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/docs/source/reference/data_samplers.rst b/docs/source/reference/data_samplers.rst
index 22204df3e5c..b0257827aba 100644
--- a/docs/source/reference/data_samplers.rst
+++ b/docs/source/reference/data_samplers.rst
@@ -32,7 +32,10 @@ Samplers control how data is retrieved from the replay buffer storage.
 Writers
 -------
 
-Writers control how data is written to the storage.
+Writers control how data is written to the storage. Writers that reuse
+storage slots can stamp each slot with a reuse counter so consumers holding
+an index can detect that it was overwritten -- see
+:ref:`Detecting overwritten slots <ref_buffers_generations>`.
 
 .. autosummary::
     :toctree: generated/
diff --git a/test/rb/test_ensemble.py b/test/rb/test_ensemble.py
index 6030729387f..2da22a87128 100644
--- a/test/rb/test_ensemble.py
+++ b/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", "index_generation") == 1).all()
+            assert (s.exclude("index") == 1).all()
             assert s.numel() == 4
         else:
             for leaf in tree_iter(s):
diff --git a/test/rb/test_samplers.py b/test/rb/test_samplers.py
index 9abb6f14581..2e410f160a6 100644
--- a/test/rb/test_samplers.py
+++ b/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", "index_generation") == 1).all()
+            assert (s.exclude("index") == 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", "index_generation") == 0).all()
+        assert (s.exclude("index") == 0).all()
 
     def test_sampler_without_rep_dumps_loads(self, tmpdir):
         d0 = tmpdir + "/save0"
diff --git a/test/rb/test_storages.py b/test/rb/test_storages.py
index 76932a87cb0..228733e8a8f 100644
--- a/test/rb/test_storages.py
+++ b/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", "index_generation") == 1).all()
+        assert (s.exclude("index") == 1).all()
 
     @pytest.mark.skipif(
         TORCH_VERSION < version.parse("2.5.0"), reason="requires Torch >= 2.5.0"
diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index 1b218471311..7c31e670073 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -404,8 +404,33 @@ def test_roundrobin_dumps_loads_write_count(self, tmp_path):
 
 
 class TestWriterGeneration:
-    def test_default_writer_tracks_generations(self):
+    def test_tracking_is_opt_in(self):
+        # generation tracking allocates one int64 per slot and adds a key to the
+        # sampler output, so an unconfigured buffer must be untouched by it
         rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        assert rb._writer.tracks_generations is False
+        index = rb.extend(torch.arange(10))
+        torch.testing.assert_close(
+            rb._writer.generations_of(index),
+            torch.full((10,), -1, dtype=torch.int64),
+        )
+        assert getattr(rb._storage, "_slot_generations", None) is None
+        _, info = rb.sample(4, return_info=True)
+        assert "index_generation" not in info
+
+    def test_default_sample_is_unchanged_by_the_feature(self):
+        # the regression this guards: a buffer that never asked for generations
+        # must keep its exact sample() key set
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(TensorDict({"a": torch.arange(10)}, [10]))
+        sample = rb.sample(4)
+        assert "index_generation" not in sample.keys()
+
+    def test_enabled_writer_tracks_generations(self):
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(10),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         assert rb._writer.tracks_generations is True
         index = rb.extend(torch.arange(10))
         gen = rb._writer.generations_of(index)
@@ -413,15 +438,55 @@ def test_default_writer_tracks_generations(self):
         assert gen.shape == index.shape
         assert (gen == 0).all()
 
+    def test_generations_live_on_the_storage(self):
+        # two buffers sharing a storage overwrite each other's slots, so they
+        # must observe the same stamps -- a per-writer counter would let one
+        # buffer's handles look live after the other overwrote the slot
+        storage = LazyTensorStorage(4)
+        rb_a = ReplayBuffer(
+            storage=storage, writer=RoundRobinWriter(track_generations=True)
+        )
+        rb_b = ReplayBuffer(
+            storage=storage, writer=RoundRobinWriter(track_generations=True)
+        )
+        index = rb_a.extend(torch.arange(4))
+        gen = rb_a._writer.generations_of(index)
+        # rb_b overwrites slots 0 and 1; rb_a must see them go stale
+        rb_b.extend(torch.arange(10, 12))
+        after = rb_a._writer.generations_of(index)
+        torch.testing.assert_close(
+            after != gen, torch.tensor([True, True, False, False])
+        )
+
     def test_non_tracking_writer_reports_minus_one(self):
         writer = TensorDictMaxValueWriter(rank_key="key")
         assert writer.tracks_generations is False
         gen = writer.generations_of(torch.arange(4))
         torch.testing.assert_close(gen, torch.full((4,), -1, dtype=torch.int64))
 
+    def test_multidim_storage_1d_index_is_a_batch_of_slots(self):
+        # with storage.ndim == 2, a 1-D tensor of length 2 is two slot indices,
+        # not one (row, col) coordinate: guessing wrong silently returns one
+        # generation where the caller asked for two
+        rb = TensorDictReplayBuffer(
+            storage=LazyTensorStorage(4, ndim=2),
+            writer=TensorDictRoundRobinWriter(track_generations=True),
+        )
+        rb.extend(TensorDict({"a": torch.zeros(4, 3)}, [4, 3]))
+        assert rb._storage.ndim == 2
+        gen = rb._writer.generations_of(torch.tensor([1, 3]))
+        assert gen.shape == (2,)
+        torch.testing.assert_close(gen, torch.zeros(2, dtype=torch.int64))
+        # the tuple form still addresses a single cell by its dim-0 slot
+        single = rb._writer.generations_of((torch.tensor(1), torch.tensor(2)))
+        torch.testing.assert_close(single, torch.zeros((), dtype=torch.int64))
+
     def test_generation_increments_on_reuse(self):
         size = 4
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size))
         torch.testing.assert_close(
             rb._writer.generations_of(torch.arange(size)),
@@ -434,7 +499,10 @@ def test_generation_increments_on_reuse(self):
 
     def test_generation_wraparound(self):
         size = 5
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(2 * size))
         torch.testing.assert_close(
             rb._writer.generations_of(torch.arange(size)),
@@ -443,7 +511,10 @@ def test_generation_wraparound(self):
 
     def test_generation_extend_wrapping_twice(self):
         size = 4
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         # slots 0 and 1 are written three times, slots 2 and 3 twice
         rb.extend(torch.arange(2 * size + 2))
         torch.testing.assert_close(
@@ -452,7 +523,10 @@ def test_generation_extend_wrapping_twice(self):
 
     def test_generation_add(self):
         size = 3
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         for i in range(size + 1):
             rb.add(torch.tensor(i))
         torch.testing.assert_close(
@@ -460,14 +534,20 @@ def test_generation_add(self):
         )
 
     def test_generations_of_unwritten_reports_minus_one(self):
-        rb = ReplayBuffer(storage=LazyTensorStorage(4))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(4),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(2))
         gen = rb._writer.generations_of(torch.arange(4))
         torch.testing.assert_close(gen, torch.tensor([0, 0, -1, -1]))
 
     def test_generation_tensordict_writer(self):
         size = 4
-        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(size))
+        rb = TensorDictReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(TensorDict({"a": torch.arange(2 * size)}, [2 * size]))
         torch.testing.assert_close(
             rb._writer.generations_of(torch.arange(size)),
@@ -476,7 +556,7 @@ def test_generation_tensordict_writer(self):
 
     def test_generation_write_at(self):
         storage = LazyTensorStorage(4)
-        writer = RoundRobinWriter()
+        writer = RoundRobinWriter(track_generations=True)
         writer.register_storage(storage)
         writer.extend(torch.arange(4))
         writer.write_at(torch.tensor([0, 1]), torch.tensor([10, 11]))
@@ -485,7 +565,10 @@ def test_generation_write_at(self):
         )
 
     def test_empty_is_monotonic(self):
-        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(10),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         index = rb.extend(torch.arange(10))
         before = rb._writer.generations_of(index)
         rb.empty()
@@ -494,14 +577,20 @@ def test_empty_is_monotonic(self):
         assert (after > before).all()
 
     def test_empty_invalidates_handles_immediately(self):
-        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(10),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         index = rb.extend(torch.arange(10))
         gen = rb._writer.generations_of(index)
         rb.empty()
         assert (rb._writer.generations_of(index) != gen).all()
 
     def test_empty_preserves_unwritten_sentinel(self):
-        rb = ReplayBuffer(storage=LazyTensorStorage(4))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(4),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(2))
         rb.empty()
         torch.testing.assert_close(
@@ -510,10 +599,16 @@ def test_empty_preserves_unwritten_sentinel(self):
 
     def test_generation_state_dict_roundtrip(self):
         size = 4
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size + 1))
         sd = rb.state_dict()
-        rb2 = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb2 = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb2.load_state_dict(sd)
         torch.testing.assert_close(
             rb2._writer.generations_of(torch.arange(size)),
@@ -521,21 +616,27 @@ def test_generation_state_dict_roundtrip(self):
         )
 
     def test_legacy_state_dict_without_generation_loads(self):
-        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(10),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(5))
         sd = rb.state_dict()
         del sd["_writer"]["_generation"]
-        rb2 = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb2 = ReplayBuffer(
+            storage=LazyTensorStorage(10),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb2.load_state_dict(sd)
         assert rb2._writer._cursor == 5
 
     def test_generation_dumps_loads(self, tmp_path):
-        writer = RoundRobinWriter()
+        writer = RoundRobinWriter(track_generations=True)
         writer._cursor = 2
         writer._write_count = 9
         writer._generation = torch.tensor([3, 2, 2, 1])
         writer.dumps(tmp_path)
-        writer2 = RoundRobinWriter()
+        writer2 = RoundRobinWriter(track_generations=True)
         writer2.loads(tmp_path)
         assert writer2._cursor == 2
         assert writer2._write_count == 9
@@ -545,7 +646,10 @@ def test_generation_dumps_loads(self, tmp_path):
 
     def test_sample_returns_generation(self):
         size = 8
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size))
         _, info = rb.sample(4, return_info=True)
         assert "index_generation" in info
@@ -565,7 +669,10 @@ def test_non_tracking_sample_has_no_generation(self):
 
     def test_tensordict_sample_has_generation_key(self):
         size = 8
-        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(size))
+        rb = TensorDictReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(TensorDict({"a": torch.arange(size)}, [size]))
         sample = rb.sample(4)
         assert "index_generation" in sample.keys()
@@ -573,7 +680,10 @@ def test_tensordict_sample_has_generation_key(self):
 
     def test_wraparound_race_detectable(self):
         size = 8
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size))
         _, info = rb.sample(4, return_info=True)
         sampled_index = torch.as_tensor(info["index"])
@@ -584,7 +694,10 @@ def test_wraparound_race_detectable(self):
 
     def test_partial_reuse_detectable(self):
         size = 8
-        rb = ReplayBuffer(storage=LazyTensorStorage(size))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size))
         _, info = rb.sample(size, return_info=True)
         idx = torch.as_tensor(info["index"])
@@ -596,7 +709,10 @@ def test_partial_reuse_detectable(self):
     @pytest.mark.parametrize("device", get_default_devices())
     def test_generation_on_storage_device(self, device):
         size = 8
-        rb = ReplayBuffer(storage=LazyTensorStorage(size, device=device))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size, device=device),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size, device=device))
         assert rb._writer._generation.device.type == device.type
         _, info = rb.sample(4, return_info=True)
@@ -613,7 +729,10 @@ def test_generation_on_storage_device(self, device):
     @pytest.mark.parametrize("device", get_default_devices())
     def test_generation_add_on_storage_device(self, device):
         size = 3
-        rb = ReplayBuffer(storage=LazyTensorStorage(size, device=device))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size, device=device),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         for i in range(size + 1):
             rb.add(torch.tensor(i, device=device))
         torch.testing.assert_close(
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 59932b42773..15c11266dc5 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -39,9 +39,20 @@ def tree_leaves(data):  # noqa: D103
 from torchrl.data.replay_buffers.storages import Storage
 from torchrl.data.replay_buffers.utils import _is_int, _reduce
 
-# Storage capacities at or above this value are treated as unbounded (lazy
-# storages report a sentinel max size), triggering dynamic generation growth.
-_GENERATION_UNBOUNDED = 2**40
+# Generation buffers for storages up to this many slots are allocated in one
+# shot, so their shape is stable and the ``torch.compile`` extend/sample path
+# does not recompile. Larger (or effectively unbounded -- ``ListStorage`` with
+# no ``max_size`` reports ``torch.iinfo(torch.int64).max``) capacities grow
+# geometrically on demand instead of trying to allocate the whole thing.
+_GENERATION_EAGER_ALLOC_LIMIT = 2**20
+_GENERATION_MIN_ALLOC = 1024
+
+# Attribute under which the per-slot generation buffer is stored *on the
+# storage*. It belongs to the storage, not to the writer: two buffers sharing
+# one storage overwrite each other's slots, so a per-writer counter would let
+# buffer A's handles look live after buffer B overwrote the slot -- exactly the
+# staleness the feature exists to detect.
+_SLOT_GENERATIONS_ATTR = "_slot_generations"
 
 
 class Writer(ABC):
@@ -54,7 +65,10 @@ def __init__(self, compilable: bool = False) -> None:
         self._storage = None
         self._compilable = compilable
 
-    #: Whether this writer type stamps storage slots with a reuse generation.
+    #: Whether this writer stamps storage slots with a reuse generation. Always
+    #: ``False`` unless the writer both supports generation tracking and was
+    #: constructed with it enabled (see
+    #: :class:`~torchrl.data.RoundRobinWriter`).
     tracks_generations: bool = False
 
     def register_storage(self, storage: Storage) -> None:
@@ -63,9 +77,27 @@ def register_storage(self, storage: Storage) -> None:
     def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
         """Returns the generation stamp for each physical slot in ``index``.
 
-        The stamp advances once per write to that slot, so a single ``extend``
-        that wraps the storage advances a reused slot once per write it
-        receives. Writers that do not track slot reuse report ``-1``.
+        A slot's stamp advances once per write it receives, so a single
+        ``extend`` that wraps the storage advances a reused slot once per write.
+        Comparing a stamp captured at sampling time against the current stamp
+        tells you whether the slot still holds the data you sampled.
+
+        Writers that do not track slot reuse -- and writers constructed with
+        ``track_generations=False``, which is the default -- report ``-1``
+        everywhere. Never-written slots also report ``-1``, so ``-1`` means
+        "no usable stamp" rather than "generation zero".
+
+        Args:
+            index (int or torch.Tensor): dim-0 slot indices. A 1-D tensor is
+                always read as a batch of slot indices; for a storage with
+                ``ndim > 1``, pass a ``tuple`` of per-dimension indices (as
+                :meth:`~torchrl.data.ReplayBuffer.extend` returns) to identify
+                a single cell -- only its dim-0 component is used, since a
+                generation stamps a whole dim-0 slot.
+
+        Returns:
+            torch.Tensor: ``int64`` stamps shaped like the dim-0 component of
+            ``index``, on ``index``'s device.
         """
         index = torch.as_tensor(index)
         return torch.full(index.shape, -1, dtype=torch.int64, device=index.device)
@@ -170,18 +202,74 @@ class RoundRobinWriter(Writer):
             If ``True``, the writer cannot be shared between multiple processes.
             Defaults to ``False``.
 
-    """
+    Keyword Args:
+        track_generations (bool, optional): if ``True``, stamp every storage
+            slot with a counter that advances each time the slot is written, so
+            a consumer holding an index can tell whether the slot still holds
+            the data it sampled. Reads are exposed through
+            :meth:`generations_of`, and :meth:`~torchrl.data.ReplayBuffer.sample`
+            adds an ``"index_generation"`` entry to its ``info`` (and, for
+            tensordict buffers, to the sample). Defaults to ``False``: enabling
+            it allocates one ``int64`` slot per storage slot and adds a key to
+            the sampler output, so it is opt-in.
 
-    tracks_generations: bool = True
+    .. note::
+        The generation buffer lives on the *storage*, not on the writer, so two
+        buffers sharing one storage observe each other's writes. It is
+        process-local: a slot overwritten in another process is not reflected
+        here. See :ref:`ref_buffers_generations`.
 
-    def __init__(self, compilable: bool = False) -> None:
+    Examples:
+        >>> import torch
+        >>> from torchrl.data import LazyTensorStorage, ReplayBuffer, RoundRobinWriter
+        >>> rb = ReplayBuffer(
+        ...     storage=LazyTensorStorage(4),
+        ...     writer=RoundRobinWriter(track_generations=True),
+        ... )
+        >>> index = rb.extend(torch.arange(4))
+        >>> rb.writer.generations_of(index)
+        tensor([0, 0, 0, 0])
+        >>> _ = rb.extend(torch.arange(4, 6))  # overwrites slots 0 and 1
+        >>> rb.writer.generations_of(index)
+        tensor([1, 1, 0, 0])
+    """
+
+    def __init__(
+        self, compilable: bool = False, *, track_generations: bool = False
+    ) -> None:
         super().__init__(compilable=compilable)
         self._cursor = 0
         self._write_count  # noqa
-        self._generation = None
+        self._track_generations = track_generations
+        # Holds the buffer until a storage is registered (dumps/loads and
+        # load_state_dict can both run on a storage-less writer).
+        self._pending_generation = None
+
+    @property
+    def tracks_generations(self) -> bool:
+        return self._track_generations
+
+    @property
+    def _generation(self) -> torch.Tensor | None:
+        if self._storage is None:
+            return self._pending_generation
+        return getattr(self._storage, _SLOT_GENERATIONS_ATTR, None)
+
+    @_generation.setter
+    def _generation(self, value: torch.Tensor | None) -> None:
+        if self._storage is None:
+            self._pending_generation = value
+        else:
+            setattr(self._storage, _SLOT_GENERATIONS_ATTR, value)
 
     def register_storage(self, storage: Storage) -> None:
         super().register_storage(storage)
+        pending, self._pending_generation = self._pending_generation, None
+        # A buffer restored from a checkpoint carries its stamps in the writer;
+        # a storage already shared with another buffer carries the live ones and
+        # wins, so the two writers cannot disagree about a slot's generation.
+        if pending is not None and self._generation is None:
+            self._generation = pending
         self._align_generation_device()
 
     def _generation_device(self, index: int | torch.Tensor) -> torch.device:
@@ -205,9 +293,6 @@ def _align_generation_device(self) -> None:
     def _ensure_generation(
         self, capacity: int, min_size: int, device: torch.device
     ) -> None:
-        # Bounded storages allocate to capacity once (stable shape, so the
-        # ``torch.compile`` extend/sample path does not recompile); lazy storages
-        # report a sentinel capacity and instead grow geometrically.
         generation = self._generation
         if generation is not None and generation.device != device:
             generation = generation.to(device)
@@ -215,17 +300,25 @@ def _ensure_generation(
         current = 0 if generation is None else generation.numel()
         if current >= min_size:
             return
-        size = (
-            capacity if capacity < _GENERATION_UNBOUNDED else max(min_size, current * 2)
-        )
+        if capacity <= _GENERATION_EAGER_ALLOC_LIMIT:
+            # One allocation covering every slot: the shape never changes again.
+            size = capacity
+        else:
+            # Too large (or unbounded) to allocate up front -- grow geometrically
+            # and stay within the storage's capacity.
+            size = min(capacity, max(min_size, 2 * current, _GENERATION_MIN_ALLOC))
         new_generation = torch.full((size,), -1, dtype=torch.int64, device=device)
         if generation is not None:
             new_generation[:current] = generation
-        if not self._compilable and new_generation.device.type == "cpu":
-            new_generation.share_memory_()
+        # Deliberately not shared across processes: the buffer is replaced (not
+        # mutated) whenever it grows, so a shared mapping would silently stop
+        # tracking after the first growth. Cross-process staleness detection
+        # needs a storage-owned, fixed-size mapping -- see the docs.
         self._generation = new_generation
 
     def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
+        if not self._track_generations:
+            return
         device = self._generation_device(index)
         if _is_int(index):
             capacity = self._storage._max_size_along_dim0(single_data=data)
@@ -236,9 +329,12 @@ def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
             if index.numel() == 0:
                 return
             capacity = self._storage._max_size_along_dim0(batched_data=data)
-            min_size = (
-                capacity if capacity < _GENERATION_UNBOUNDED else int(index.max()) + 1
-            )
+            if capacity <= _GENERATION_EAGER_ALLOC_LIMIT:
+                min_size = capacity
+            else:
+                # Only reached for capacities we cannot allocate up front, so the
+                # device sync from ``.max()`` is not on the common extend path.
+                min_size = int(index.max()) + 1
             self._ensure_generation(capacity, min_size, device)
             index = index.to(device)
             self._generation.index_put_(
@@ -246,11 +342,17 @@ def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
             )
 
     def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
+        if not self._track_generations:
+            return super().generations_of(index)
         if isinstance(index, tuple):
             index = index[0]
         elif (
             isinstance(index, torch.Tensor)
-            and index.ndim
+            # Only a batch of coordinate vectors, i.e. ndim >= 2, can be
+            # unambiguously distinguished from a batch of dim-0 indices: a 1-D
+            # tensor of length storage.ndim is far more likely to be several
+            # slot indices than one coordinate. Pass a tuple for the latter.
+            and index.ndim >= 2
             and self._storage is not None
             and self._storage.ndim > 1
             and index.shape[-1] == self._storage.ndim
@@ -272,7 +374,7 @@ def dumps(self, path):
             "cursor": self._cursor,
             "write_count": self._write_count,
         }
-        generation = self._generation
+        generation = self._generation if self._track_generations else None
         if generation is not None:
             generation = generation.cpu()
             try:
@@ -305,8 +407,6 @@ def loads(self, path):
                 dtype=_STRDTYPE2DTYPE[metadata["generation_dtype"]],
                 shape=torch.Size(generation_shape),
             ).clone()
-            if not self._compilable:
-                generation.share_memory_()
             self._generation = generation
             self._align_generation_device()
 
@@ -392,7 +492,7 @@ def _update_storage_len_for_write_at(self, index: int | torch.Tensor) -> None:
 
     def state_dict(self) -> dict[str, Any]:
         state_dict = {"_cursor": self._cursor, "_write_count": self._write_count}
-        if self._generation is not None:
+        if self._track_generations and self._generation is not None:
             state_dict["_generation"] = self._generation.clone()
         return state_dict
 
@@ -403,17 +503,16 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
             self._write_count = write_count
         generation = state_dict.get("_generation")
         if generation is not None:
-            generation = generation.clone()
-            if not self._compilable and generation.device.type == "cpu":
-                generation.share_memory_()
-            self._generation = generation
+            self._generation = generation.clone()
             self._align_generation_device()
 
     def _empty(self, empty_write_count: bool = True) -> None:
         self._cursor = 0
-        generation = self._generation
+        generation = self._generation if self._track_generations else None
         if generation is not None:
-            # never-written slots keep the -1 sentinel
+            # Emptying invalidates every handle, so stamps advance rather than
+            # reset -- a reset would make pre-empty handles look live again.
+            # Never-written slots keep the -1 sentinel.
             generation[generation >= 0] += 1
         if empty_write_count:
             self._write_count = 0
@@ -503,7 +602,12 @@ def __repr__(self):
 
 
 class TensorDictRoundRobinWriter(RoundRobinWriter):
-    """A RoundRobin Writer class for composable, tensordict-based replay buffers."""
+    """A RoundRobin Writer class for composable, tensordict-based replay buffers.
+
+    Takes the same arguments as :class:`RoundRobinWriter`, including
+    ``track_generations``. When enabled, ``"index_generation"`` is written into
+    the sampled tensordict alongside ``"index"``.
+    """
 
     def add(self, data: Any) -> int | torch.Tensor:
         index = self._cursor

From 7dc7c0e8589af93c041572a204d8cb298fea118a Mon Sep 17 00:00:00 2001
From: Vincent Moens <vincentmoens@gmail.com>
Date: Tue, 4 Aug 2026 17:22:49 +0100
Subject: [PATCH 11/13] [Test] Correct the multidim generations_of test
 expectation

Running it locally: LazyTensorStorage(4, ndim=2) with [4, 3] data gives a dim-0
capacity of 1, so the extend wrapped and the stamps were not all 0. Use
LazyTensorStorage(12, ndim=2) (dim-0 capacity 3) with in-range indices, and also
assert the [N, ndim] coordinate batch extend() returns is read as coordinates.

test/rb: 4201 passed, 1029 skipped -- baseline is 4175/1025, so +26 new tests
and no regressions.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 test/rb/test_writers.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index 7c31e670073..27f81ab4c2c 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -469,12 +469,20 @@ def test_multidim_storage_1d_index_is_a_batch_of_slots(self):
         # not one (row, col) coordinate: guessing wrong silently returns one
         # generation where the caller asked for two
         rb = TensorDictReplayBuffer(
-            storage=LazyTensorStorage(4, ndim=2),
+            storage=LazyTensorStorage(12, ndim=2),
             writer=TensorDictRoundRobinWriter(track_generations=True),
         )
-        rb.extend(TensorDict({"a": torch.zeros(4, 3)}, [4, 3]))
+        index = rb.extend(TensorDict({"a": torch.zeros(4, 3)}, [4, 3]))
         assert rb._storage.ndim == 2
-        gen = rb._writer.generations_of(torch.tensor([1, 3]))
+        # extend returns a [N, ndim] coordinate batch: read as coordinates
+        assert index.ndim == 2 and index.shape[-1] == 2
+        torch.testing.assert_close(
+            rb._writer.generations_of(index),
+            torch.zeros(index.shape[0], dtype=torch.int64),
+        )
+        # a 1-D tensor of length ndim is two slot indices, not one coordinate:
+        # the old heuristic collapsed this to a single 0-dim stamp
+        gen = rb._writer.generations_of(torch.tensor([1, 2]))
         assert gen.shape == (2,)
         torch.testing.assert_close(gen, torch.zeros(2, dtype=torch.int64))
         # the tuple form still addresses a single cell by its dim-0 slot

From 6d3af1aec3bb114fa9144514a717f9e04f909828 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vincentmoens@gmail.com>
Date: Fri, 7 Aug 2026 15:17:06 +0100
Subject: [PATCH 12/13] [BugFix] Fix replay generation tracking edge cases

---
 docs/source/reference/data_replaybuffers.rst |  2 +-
 test/rb/test_writers.py                      | 63 ++++++++++++++++++--
 test/test_configs.py                         | 10 +++-
 torchrl/data/replay_buffers/writers.py       | 21 +++++--
 torchrl/trainers/algorithms/configs/data.py  |  6 +-
 5 files changed, 87 insertions(+), 15 deletions(-)

diff --git a/docs/source/reference/data_replaybuffers.rst b/docs/source/reference/data_replaybuffers.rst
index 50c4e5e3d2c..3ccf6d34735 100644
--- a/docs/source/reference/data_replaybuffers.rst
+++ b/docs/source/reference/data_replaybuffers.rst
@@ -243,7 +243,7 @@ random sampling. Prefetching, prioritized replay and multidimensional storages
 are rejected explicitly.
 
 Detecting overwritten slots: generation stamps
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. _ref_buffers_generations:
 
diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index 27f81ab4c2c..2a99124a34a 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -458,6 +458,36 @@ def test_generations_live_on_the_storage(self):
             after != gen, torch.tensor([True, True, False, False])
         )
 
+    def test_non_tracking_writer_updates_shared_generations(self):
+        storage = LazyTensorStorage(4)
+        tracking_rb = ReplayBuffer(
+            storage=storage, writer=RoundRobinWriter(track_generations=True)
+        )
+        non_tracking_rb = ReplayBuffer(storage=storage)
+        index = tracking_rb.extend(torch.arange(4))
+        generation = tracking_rb.writer.generations_of(index)
+
+        non_tracking_rb.extend(torch.arange(10, 12))
+
+        torch.testing.assert_close(
+            tracking_rb.writer.generations_of(index) != generation,
+            torch.tensor([True, True, False, False]),
+        )
+        assert non_tracking_rb.writer.tracks_generations is False
+
+    def test_non_tracking_writer_empty_invalidates_shared_generations(self):
+        storage = LazyTensorStorage(4)
+        tracking_rb = ReplayBuffer(
+            storage=storage, writer=RoundRobinWriter(track_generations=True)
+        )
+        non_tracking_rb = ReplayBuffer(storage=storage)
+        index = tracking_rb.extend(torch.arange(4))
+        generation = tracking_rb.writer.generations_of(index)
+
+        non_tracking_rb.empty()
+
+        assert (tracking_rb.writer.generations_of(index) != generation).all()
+
     def test_non_tracking_writer_reports_minus_one(self):
         writer = TensorDictMaxValueWriter(rank_key="key")
         assert writer.tracks_generations is False
@@ -550,6 +580,15 @@ def test_generations_of_unwritten_reports_minus_one(self):
         gen = rb._writer.generations_of(torch.arange(4))
         torch.testing.assert_close(gen, torch.tensor([0, 0, -1, -1]))
 
+    def test_generations_of_out_of_range_reports_minus_one(self):
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(4),
+            writer=RoundRobinWriter(track_generations=True),
+        )
+        rb.extend(torch.arange(4))
+        gen = rb.writer.generations_of(torch.tensor([-5, -1, 4, 100]))
+        torch.testing.assert_close(gen, torch.full((4,), -1, dtype=torch.int64))
+
     def test_generation_tensordict_writer(self):
         size = 4
         rb = TensorDictReplayBuffer(
@@ -752,7 +791,10 @@ def test_generation_add_on_storage_device(self, device):
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
     def test_generation_cuda_data_into_cuda_storage(self):
         size = 8
-        rb = ReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size, device="cuda"),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size, device="cuda"))
         assert rb._writer._generation.device.type == "cuda"
         _, info = rb.sample(4, return_info=True)
@@ -768,7 +810,10 @@ def test_generation_cuda_data_into_cuda_storage(self):
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
     def test_generation_cpu_data_into_cuda_storage(self):
         size = 4
-        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb = TensorDictReplayBuffer(
+            storage=LazyTensorStorage(size, device="cuda"),
+            writer=TensorDictRoundRobinWriter(track_generations=True),
+        )
         rb.extend(TensorDict({"a": torch.arange(2 * size)}, [2 * size]))
         assert rb._writer._generation.device.type == "cuda"
         torch.testing.assert_close(
@@ -780,9 +825,15 @@ def test_generation_cpu_data_into_cuda_storage(self):
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
     def test_generation_state_dict_roundtrip_cuda(self):
         size = 4
-        rb = ReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(size, device="cuda"),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb.extend(torch.arange(size + 1, device="cuda"))
-        rb2 = ReplayBuffer(storage=LazyTensorStorage(size, device="cuda"))
+        rb2 = ReplayBuffer(
+            storage=LazyTensorStorage(size, device="cuda"),
+            writer=RoundRobinWriter(track_generations=True),
+        )
         rb2.load_state_dict(rb.state_dict())
         index = torch.arange(size, device="cuda")
         assert rb2._writer._generation.device.type == "cuda"
@@ -793,11 +844,11 @@ def test_generation_state_dict_roundtrip_cuda(self):
     @pytest.mark.gpu
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
     def test_generation_dumps_loads_cuda(self, tmp_path):
-        writer = RoundRobinWriter()
+        writer = RoundRobinWriter(track_generations=True)
         writer.register_storage(LazyTensorStorage(4, device="cuda"))
         writer._generation = torch.tensor([3, 2, 2, 1], device="cuda")
         writer.dumps(tmp_path)
-        writer2 = RoundRobinWriter()
+        writer2 = RoundRobinWriter(track_generations=True)
         writer2.register_storage(LazyTensorStorage(4, device="cuda"))
         writer2.loads(tmp_path)
         assert writer2._generation.device.type == "cuda"
diff --git a/test/test_configs.py b/test/test_configs.py
index 5189ef16e07..20f3824e838 100644
--- a/test/test_configs.py
+++ b/test/test_configs.py
@@ -441,14 +441,16 @@ def test_round_robin_writer_config(self):
         from hydra.utils import instantiate
         from torchrl.trainers.algorithms.configs.data import RoundRobinWriterConfig
 
-        cfg = RoundRobinWriterConfig(compilable=True)
+        cfg = RoundRobinWriterConfig(compilable=True, track_generations=True)
         assert cfg._target_ == "torchrl.data.replay_buffers.RoundRobinWriter"
         assert cfg.compilable is True
+        assert cfg.track_generations is True
 
         # Test instantiation
         writer = instantiate(cfg)
         assert isinstance(writer, RoundRobinWriter)
         assert writer._compilable is True
+        assert writer.tracks_generations is True
 
     def test_sampler_config(self):
         """Test basic SamplerConfig."""
@@ -652,14 +654,18 @@ def test_tensor_dict_round_robin_writer_config(self):
             TensorDictRoundRobinWriterConfig,
         )
 
-        cfg = TensorDictRoundRobinWriterConfig(compilable=True)
+        cfg = TensorDictRoundRobinWriterConfig(
+            compilable=True, track_generations=True
+        )
         assert cfg._target_ == "torchrl.data.replay_buffers.TensorDictRoundRobinWriter"
         assert cfg.compilable is True
+        assert cfg.track_generations is True
 
         # Test instantiation
         writer = instantiate(cfg)
         assert isinstance(writer, TensorDictRoundRobinWriter)
         assert writer._compilable is True
+        assert writer.tracks_generations is True
 
     @pytest.mark.skipif(not _has_hydra, reason="Hydra is not installed")
     def test_immutable_dataset_writer_config(self):
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 15c11266dc5..55e635a08ab 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -197,6 +197,8 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 class RoundRobinWriter(Writer):
     """A RoundRobin Writer class for composable replay buffers.
 
+    See also :class:`~torchrl.trainers.algorithms.configs.RoundRobinWriterConfig`.
+
     Args:
         compilable (bool, optional): whether the writer is compilable.
             If ``True``, the writer cannot be shared between multiple processes.
@@ -317,7 +319,10 @@ def _ensure_generation(
         self._generation = new_generation
 
     def _bump_generation(self, index: int | torch.Tensor, data: Any) -> None:
-        if not self._track_generations:
+        # A writer that did not opt into generation tracking must still update a
+        # buffer installed on a shared storage by another writer. It does not
+        # allocate the buffer itself or expose generations in its samples.
+        if not self._track_generations and self._generation is None:
             return
         device = self._generation_device(index)
         if _is_int(index):
@@ -363,8 +368,11 @@ def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
             return torch.full(index.shape, -1, dtype=torch.int64, device=index.device)
         idx = index.to(self._generation.device)
         n = self._generation.numel()
-        gen = self._generation[idx.clamp(max=n - 1)]
-        gen = torch.where(idx < n, gen, torch.full_like(gen, -1))
+        if not n:
+            return torch.full(index.shape, -1, dtype=torch.int64, device=index.device)
+        gen = self._generation[idx.clamp(min=0, max=n - 1)]
+        valid = (idx >= 0) & (idx < n)
+        gen = torch.where(valid, gen, torch.full_like(gen, -1))
         return gen.to(index.device)
 
     def dumps(self, path):
@@ -508,7 +516,9 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 
     def _empty(self, empty_write_count: bool = True) -> None:
         self._cursor = 0
-        generation = self._generation if self._track_generations else None
+        # Emptying through any writer invalidates handles held by tracking
+        # writers that share this storage.
+        generation = self._generation
         if generation is not None:
             # Emptying invalidates every handle, so stamps advance rather than
             # reset -- a reset would make pre-empty handles look live again.
@@ -604,6 +614,9 @@ def __repr__(self):
 class TensorDictRoundRobinWriter(RoundRobinWriter):
     """A RoundRobin Writer class for composable, tensordict-based replay buffers.
 
+    See also
+    :class:`~torchrl.trainers.algorithms.configs.TensorDictRoundRobinWriterConfig`.
+
     Takes the same arguments as :class:`RoundRobinWriter`, including
     ``track_generations``. When enabled, ``"index_generation"`` is written into
     the sampled tensordict alongside ``"index"``.
diff --git a/torchrl/trainers/algorithms/configs/data.py b/torchrl/trainers/algorithms/configs/data.py
index e0bdf78ef8b..5a6c5acad3e 100644
--- a/torchrl/trainers/algorithms/configs/data.py
+++ b/torchrl/trainers/algorithms/configs/data.py
@@ -31,10 +31,11 @@ def __post_init__(self) -> None:
 
 @dataclass
 class RoundRobinWriterConfig(WriterConfig):
-    """Configuration for round-robin writer that distributes data across multiple storages."""
+    """Hydra configuration for :class:`~torchrl.data.RoundRobinWriter`."""
 
     _target_: str = "torchrl.data.replay_buffers.RoundRobinWriter"
     compilable: bool = False
+    track_generations: bool = False
 
     def __post_init__(self) -> None:
         """Post-initialization hook for round-robin writer configurations."""
@@ -94,10 +95,11 @@ class TensorDictMaxValueWriterConfig(WriterConfig):
 
 @dataclass
 class TensorDictRoundRobinWriterConfig(WriterConfig):
-    """Configuration for TensorDict round-robin writer."""
+    """Hydra configuration for :class:`~torchrl.data.TensorDictRoundRobinWriter`."""
 
     _target_: str = "torchrl.data.replay_buffers.TensorDictRoundRobinWriter"
     compilable: bool = False
+    track_generations: bool = False
 
 
 @dataclass

From 1fd72be08b1145b7a299538e3e890bb59ddec0e7 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vincentmoens@gmail.com>
Date: Fri, 7 Aug 2026 15:27:35 +0100
Subject: [PATCH 13/13] [CI] Fix generation tracking lint

---
 test/test_configs.py                   | 4 +---
 torchrl/data/replay_buffers/writers.py | 2 +-
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/test/test_configs.py b/test/test_configs.py
index 20f3824e838..7aef1f3e5a7 100644
--- a/test/test_configs.py
+++ b/test/test_configs.py
@@ -654,9 +654,7 @@ def test_tensor_dict_round_robin_writer_config(self):
             TensorDictRoundRobinWriterConfig,
         )
 
-        cfg = TensorDictRoundRobinWriterConfig(
-            compilable=True, track_generations=True
-        )
+        cfg = TensorDictRoundRobinWriterConfig(compilable=True, track_generations=True)
         assert cfg._target_ == "torchrl.data.replay_buffers.TensorDictRoundRobinWriter"
         assert cfg.compilable is True
         assert cfg.track_generations is True
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 55e635a08ab..5e9e19457d5 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -614,7 +614,7 @@ def __repr__(self):
 class TensorDictRoundRobinWriter(RoundRobinWriter):
     """A RoundRobin Writer class for composable, tensordict-based replay buffers.
 
-    See also
+    See Also:
     :class:`~torchrl.trainers.algorithms.configs.TensorDictRoundRobinWriterConfig`.
 
     Takes the same arguments as :class:`RoundRobinWriter`, including