[Feature] Add generation-stamped replay slots for round-robin writer

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
 ~~~~~~~~~~~~~~~~~~~~~
 

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/