From 3ec6726dfe0de811d37c0e11ad4ae33a92448ef4 Mon Sep 17 00:00:00 2001
From: Achintya P <apaninga@berkeley.edu>
Date: Thu, 23 Jul 2026 17:01:49 -0700
Subject: [PATCH 1/4] [Test] Spec tests for generation-stamped replay slots

Executable spec for step 1 of the conditional replay-update RFC:
round-robin writers stamp each slot with an int64 generation exposed
via writer.generations_of(index); first write is generation 0, every
slot reuse increments it, empty() never revives old handles, and
generations persist through state_dict and dumps with legacy
checkpoints still loading. Sampling exposes the stamps as an
index_generation info entry / TensorDict key aligned with index.

Tests are expected to fail until the implementation lands.
---
 test/rb/test_rb_core.py | 47 ++++++++++++++++++++++
 test/rb/test_writers.py | 89 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 136 insertions(+)

diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py
index a5763e9aa1d..883cbd5daba 100644
--- a/test/rb/test_rb_core.py
+++ b/test/rb/test_rb_core.py
@@ -1196,6 +1196,53 @@ def test_stats_with_non_counting_writer(self):
         assert stats["capacity"] == 10
 
 
+class TestSampleGenerationInfo:
+    """Executable spec for exposing slot generations at sampling time (RFC step 1).
+
+    Contract pinned by this class:
+
+    - ``sample(return_info=True)`` returns an ``"index_generation"`` entry in
+      the info dict, aligned element-for-element with ``info["index"]`` and
+      equal to the writer's generation for those slots at sampling time.
+    - ``TensorDictReplayBuffer`` samples carry an ``"index_generation"`` key
+      alongside the existing ``"index"`` key.
+    - After a wraparound, previously captured (index, generation) handles
+      disagree with the writer's current generations exactly on the reused
+      slots, which is what makes staleness detectable.
+    """
+
+    def test_sample_info_contains_index_generation(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10), batch_size=4)
+        rb.extend(torch.arange(10))
+        _, info = rb.sample(return_info=True)
+        assert "index_generation" in info
+        index = torch.as_tensor(info["index"]).reshape(-1)
+        generations = torch.as_tensor(info["index_generation"]).reshape(-1)
+        assert generations.shape == index.shape
+        torch.testing.assert_close(generations, rb._writer.generations_of(index))
+
+    def test_tensordict_sample_carries_index_generation(self):
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10), batch_size=4)
+        rb.extend(TensorDict({"obs": torch.randn(10, 3)}, batch_size=[10]))
+        rb.extend(TensorDict({"obs": torch.randn(4, 3)}, batch_size=[4]))
+        sample = rb.sample()
+        assert "index_generation" in sample.keys()
+        batch = sample.batch_size[0]
+        index = sample.get("index").reshape(batch, -1)[:, 0]
+        generations = sample.get("index_generation").reshape(batch, -1)[:, 0]
+        torch.testing.assert_close(generations, rb._writer.generations_of(index))
+
+    def test_stale_handles_detectable_after_wraparound(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10), batch_size=4)
+        index = rb.extend(torch.arange(10))
+        generations = rb._writer.generations_of(index)
+        rb.extend(torch.arange(4))
+        live = rb._writer.generations_of(index) == generations
+        assert live.sum() == 6
+        assert not live[:4].any()
+        assert live[4:].all()
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py
index ffe47a51ffb..a23726cf6d3 100644
--- a/test/rb/test_writers.py
+++ b/test/rb/test_writers.py
@@ -403,6 +403,95 @@ def test_roundrobin_dumps_loads_write_count(self, tmp_path):
         assert writer2._write_count == 23
 
 
+class TestSlotGenerations:
+    """Executable spec for generation-stamped replay slots (RFC step 1).
+
+    Contract pinned by this class:
+
+    - Round-robin writers maintain one int64 generation counter per storage
+      slot, exposed through ``writer.generations_of(index)`` which accepts an
+      index tensor and returns a same-shaped int64 tensor.
+    - The first write of a slot has generation 0; every reuse of a slot
+      (round-robin wraparound or rewrite through ``add``/``extend``)
+      increments that slot's generation.
+    - ``empty()`` never revives previously handed-out (index, generation)
+      pairs: generations are monotonically nondecreasing across the buffer's
+      lifetime, including through ``empty()``.
+    - Generations persist through ``state_dict``/``load_state_dict`` and
+      ``dumps``/``loads``; checkpoints created before the feature still load.
+    """
+
+    def test_first_writes_start_at_generation_zero(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        index = rb.extend(torch.arange(10))
+        generations = rb._writer.generations_of(index)
+        assert generations.dtype == torch.int64
+        assert generations.shape == index.shape
+        assert (generations == 0).all()
+
+    def test_wraparound_increments_reused_slots_only(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(torch.arange(10))
+        reused_index = rb.extend(torch.arange(4))
+        assert (rb._writer.generations_of(reused_index) == 1).all()
+        untouched = rb._writer.generations_of(torch.arange(4, 10))
+        assert (untouched == 0).all()
+
+    def test_add_reuse_increments_generation(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(2))
+        for value in range(5):
+            rb.add(torch.full((3,), float(value)))
+        assert rb._writer.generations_of(torch.tensor([0])).item() == 2
+        assert rb._writer.generations_of(torch.tensor([1])).item() == 1
+
+    def test_empty_never_revives_old_handles(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        index = rb.extend(torch.arange(10))
+        generations_before = rb._writer.generations_of(index)
+        rb.empty()
+        rb.extend(torch.arange(10))
+        generations_after = rb._writer.generations_of(index)
+        assert (generations_after > generations_before).all()
+
+    def test_generations_survive_state_dict_roundtrip(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(torch.arange(10))
+        index = rb.extend(torch.arange(4))
+        sd = rb.state_dict()
+        rb2 = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb2.load_state_dict(sd)
+        torch.testing.assert_close(
+            rb2._writer.generations_of(torch.arange(10)),
+            rb._writer.generations_of(torch.arange(10)),
+        )
+        assert (rb2._writer.generations_of(index) == 1).all()
+
+    def test_generations_survive_dumps_loads(self, tmp_path):
+        rb = ReplayBuffer(storage=LazyMemmapStorage(10, scratch_dir=tmp_path / "data"))
+        rb.extend(torch.arange(10))
+        rb.extend(torch.arange(4))
+        rb._writer.dumps(tmp_path / "writer")
+        writer2 = RoundRobinWriter()
+        writer2.loads(tmp_path / "writer")
+        torch.testing.assert_close(
+            writer2.generations_of(torch.arange(10)),
+            rb._writer.generations_of(torch.arange(10)),
+        )
+
+    def test_legacy_state_dict_without_generations_loads(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb.extend(torch.arange(5))
+        sd = rb.state_dict()
+        sd["_writer"] = {
+            key: value
+            for key, value in sd["_writer"].items()
+            if key in ("_cursor", "_write_count")
+        }
+        rb2 = ReplayBuffer(storage=LazyTensorStorage(10))
+        rb2.load_state_dict(sd)
+        assert rb2._writer._cursor == 5
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

From 0940fd720c723bb369c7b07f2f6c17ad1ee8cfab Mon Sep 17 00:00:00 2001
From: Achintya P <apaninga@berkeley.edu>
Date: Fri, 24 Jul 2026 00:28:24 -0700
Subject: [PATCH 2/4] [Feature] Add generation-stamped replay buffer slots

Round-robin writers now stamp each storage slot with an int64
generation counter: the first write of a slot has generation 0 and
every round-robin reuse increments it, so a captured
(index, generation) pair identifies one specific record and becomes
detectably stale once the slot is recycled or the buffer is emptied.
Generations are exposed through writer.generations_of(index), persist
through state_dict and dumps/loads (legacy checkpoints still load),
and are reported at sampling time as an index_generation info entry,
which TensorDict buffers surface as a sample key next to index.
Writers that do not track reuse report -1. Positional write_at and
raw __setitem__ writes patch records in place and do not advance
generations. A wraparound-heavy extend benchmark covers the hot-path
cost.

First step of the conditional replay-update RFC (#4040).
Addresses #4041.
---
 benchmarks/test_replaybuffer_benchmark.py     |  29 ++++
 test/rb/test_ensemble.py                      |   2 +-
 test/rb/test_samplers.py                      |   4 +-
 test/rb/test_storages.py                      |   2 +-
 torchrl/data/replay_buffers/replay_buffers.py |  20 +++
 torchrl/data/replay_buffers/writers.py        | 149 +++++++++++++++++-
 6 files changed, 196 insertions(+), 10 deletions(-)

diff --git a/benchmarks/test_replaybuffer_benchmark.py b/benchmarks/test_replaybuffer_benchmark.py
index 8bdfa341c97..cfef1aba75f 100644
--- a/benchmarks/test_replaybuffer_benchmark.py
+++ b/benchmarks/test_replaybuffer_benchmark.py
@@ -402,6 +402,35 @@ def test_rb_populate(benchmark, rb, storage, sampler, size):
     )
 
 
+class create_wraparound_rb:
+    """Builds a full buffer so every timed extend reuses slots and bumps generations."""
+
+    def __init__(self, size=10_000, batch=1_000):
+        self.size = size
+        self.batch = batch
+
+    def __call__(self):
+        rb = ReplayBuffer(storage=LazyTensorStorage(self.size))
+        data = TensorDict({"a": torch.zeros(self.batch, 5)}, batch_size=[self.batch])
+        while rb.write_count < self.size:
+            rb.extend(data)
+        return ((rb, data), {})
+
+
+def extend_wraparound(rb, data):
+    for _ in range(10):
+        rb.extend(data)
+
+
+def test_rb_extend_generation_stamping(benchmark):
+    benchmark.pedantic(
+        extend_wraparound,
+        setup=create_wraparound_rb(),
+        iterations=1,
+        rounds=50,
+    )
+
+
 class create_compiled_tensor_rb:
     def __init__(
         self, rb, storage, sampler, storage_size, data_size, iters, compilable=False
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/torchrl/data/replay_buffers/replay_buffers.py b/torchrl/data/replay_buffers/replay_buffers.py
index a5d781e71df..4ae45c85e42 100644
--- a/torchrl/data/replay_buffers/replay_buffers.py
+++ b/torchrl/data/replay_buffers/replay_buffers.py
@@ -1516,6 +1516,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         with self._replay_lock if not is_comp else nc, self._write_lock if not is_comp else nc:
             index, info = self._sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -1555,6 +1559,10 @@ def sample(self, batch_size: int | None = None, return_info: bool = False) -> An
         Returns:
             A batch of data selected in the replay buffer.
             A tuple containing this batch and info if return_info flag is set to True.
+            The info entries include ``"index"``, the storage slots the batch
+            was read from, and ``"index_generation"``, the generation of each
+            slot at sampling time (see
+            :meth:`~torchrl.data.replay_buffers.RoundRobinWriter.generations_of`).
         """
         if (
             batch_size is not None
@@ -2135,6 +2143,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         ):
             index, info = self.prioritized_sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2564,6 +2576,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         with self._replay_lock if not is_comp else nc, self._write_lock if not is_comp else nc:
             index, info = self._sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
@@ -2929,6 +2945,10 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]:
         ):
             index, info = self.prioritized_sampler.sample(self._storage, batch_size)
             info["index"] = index
+            if self._writer.tracks_generations:
+                info["index_generation"] = self._writer.generations_of(
+                    index[0] if isinstance(index, tuple) else index
+                )
             data = self._storage.get(_storage_index(index, self._storage))
         if not isinstance(index, INT_CLASSES):
             data = self._collate_fn(data)
diff --git a/torchrl/data/replay_buffers/writers.py b/torchrl/data/replay_buffers/writers.py
index 9726e9861dd..7ba26a53a08 100644
--- a/torchrl/data/replay_buffers/writers.py
+++ b/torchrl/data/replay_buffers/writers.py
@@ -45,6 +45,7 @@ class Writer(ABC):
 
     _storage: Storage
     _rng: torch.Generator | None = None
+    tracks_generations: bool = False
 
     def __init__(self, compilable: bool = False) -> None:
         self._storage = None
@@ -53,6 +54,18 @@ def __init__(self, compilable: bool = False) -> None:
     def register_storage(self, storage: Storage) -> None:
         self._storage = storage
 
+    def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
+        """Returns the slot generations for the given indices.
+
+        Writers that do not track slot reuse (``tracks_generations=False``)
+        report ``-1`` for every slot. See
+        :meth:`RoundRobinWriter.generations_of` for the tracking semantics.
+        """
+        index = torch.as_tensor(index, dtype=torch.long)
+        if index.ndim > 1 and index.shape[-1]:
+            index = index[..., 0]
+        return torch.full_like(index, -1)
+
     @abstractmethod
     def add(self, data: Any) -> int:
         """Inserts one piece of data at an appropriate index, and returns that index."""
@@ -155,16 +168,40 @@ class RoundRobinWriter(Writer):
 
     """
 
+    tracks_generations: bool = True
+
     def __init__(self, compilable: bool = False) -> None:
         super().__init__(compilable=compilable)
         self._cursor = 0
         self._write_count  # noqa
+        self._slot_generations = None
 
     def dumps(self, path):
         path = Path(path).absolute()
         path.mkdir(exist_ok=True)
+        generations = self._slot_generations
+        if generations is not None:
+            try:
+                MemoryMappedTensor.from_filename(
+                    filename=path / "slot_generations.memmap",
+                    shape=generations.shape,
+                    dtype=generations.dtype,
+                ).copy_(generations)
+            except FileNotFoundError:
+                MemoryMappedTensor.from_tensor(
+                    generations, filename=path / "slot_generations.memmap"
+                )
         with open(path / "metadata.json", "w") as file:
-            json.dump({"cursor": self._cursor, "write_count": self._write_count}, file)
+            json.dump(
+                {
+                    "cursor": self._cursor,
+                    "write_count": self._write_count,
+                    "slot_generations_size": None
+                    if generations is None
+                    else generations.numel(),
+                },
+                file,
+            )
 
     def loads(self, path):
         path = Path(path).absolute()
@@ -174,15 +211,22 @@ def loads(self, path):
             write_count = metadata.get("write_count")
             if write_count is not None:
                 self._write_count = write_count
+            generations_size = metadata.get("slot_generations_size")
+        if generations_size is not None:
+            self._slot_generations = MemoryMappedTensor.from_filename(
+                filename=path / "slot_generations.memmap",
+                shape=torch.Size([generations_size]),
+                dtype=torch.int64,
+            ).clone()
 
     def add(self, data: Any) -> int | torch.Tensor:
         index = self._cursor
         _cursor = self._cursor
+        max_size_along0 = self._storage._max_size_along_dim0(single_data=data)
         # we need to update the cursor first to avoid race conditions between workers
-        self._cursor = (self._cursor + 1) % self._storage._max_size_along_dim0(
-            single_data=data
-        )
+        self._cursor = (self._cursor + 1) % max_size_along0
         self._write_count += 1
+        self._bump_generations(_cursor, max_size_along0)
         # 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)
@@ -211,6 +255,7 @@ def extend(self, data: Sequence) -> torch.Tensor:
         # we need to update the cursor first to avoid race conditions between workers
         self._cursor = (batch_size + cur_size) % max_size_along0
         self._write_count += batch_size
+        self._bump_generations(index, max_size_along0)
         # 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)
@@ -219,7 +264,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.
+
+        Positional writes patch a record in place and therefore do not change
+        the slot generation reported by :meth:`generations_of`.
+        """
         if _is_int(index):
             batch_size = 1
         else:
@@ -248,19 +297,105 @@ def _update_storage_len_for_write_at(self, index: int | torch.Tensor) -> None:
             max(len(self._storage), max_index + 1), self._storage.max_size
         )
 
+    def _ensure_generations(self, capacity: int) -> torch.Tensor:
+        generations = self._slot_generations
+        if generations is None:
+            generations = torch.full((capacity,), -1, dtype=torch.int64)
+            self._slot_generations = generations
+        elif generations.numel() < capacity:
+            grown = torch.full((capacity,), -1, dtype=torch.int64)
+            grown[: generations.numel()] = generations
+            self._slot_generations = grown
+            generations = grown
+        return generations
+
+    _UNBOUNDED_CAPACITY = 2**48
+
+    def _bump_generations(self, index: int | torch.Tensor, capacity: int) -> None:
+        if capacity >= self._UNBOUNDED_CAPACITY:
+            if isinstance(index, torch.Tensor):
+                capacity = int(index.max()) + 1
+            else:
+                capacity = int(index) + 1
+            generations = self._slot_generations
+            if generations is not None:
+                capacity = max(capacity, generations.numel())
+        generations = self._ensure_generations(capacity)
+        if isinstance(index, torch.Tensor):
+            index = index.to(generations.device)
+        generations[index] += 1
+
+    def generations_of(self, index: int | torch.Tensor) -> torch.Tensor:
+        """Returns the current generation of the given storage slots.
+
+        A slot's generation counts how many times it has been filled through
+        :meth:`add` or :meth:`extend`: the first write of a slot has
+        generation ``0`` and every round-robin reuse increments it, so a
+        previously captured ``(index, generation)`` pair identifies one
+        specific record and becomes detectably stale once the slot is
+        recycled. Emptying the buffer invalidates all outstanding pairs.
+        Slots that were never written report ``-1``. When a single
+        :meth:`extend` call wraps the storage and writes a slot more than
+        once, the slot's generation advances by one, not once per write.
+
+        Args:
+            index (int or torch.Tensor): storage slot indices. Indices
+                carrying a trailing coordinate dimension (as returned by
+                writes to multidimensional storages) are reduced to their
+                first, round-robin dimension.
+
+        Returns:
+            An ``int64`` tensor of generations with the same shape as the
+            (reduced) index.
+
+        Examples:
+            >>> import torch
+            >>> from torchrl.data import LazyTensorStorage, ReplayBuffer
+            >>> rb = ReplayBuffer(storage=LazyTensorStorage(3))
+            >>> first = rb.extend(torch.arange(3))
+            >>> reused = rb.extend(torch.arange(2))
+            >>> rb._writer.generations_of(reused)
+            tensor([1, 1])
+            >>> rb._writer.generations_of(first)
+            tensor([1, 1, 0])
+        """
+        index = torch.as_tensor(index, dtype=torch.long)
+        if index.ndim > 1 and index.shape[-1]:
+            index = index[..., 0]
+        generations = self._slot_generations
+        if generations is None:
+            return torch.full_like(index, -1)
+        index = index.to(generations.device)
+        in_range = index < generations.numel()
+        if bool(in_range.all()):
+            return generations[index]
+        out = torch.full_like(index, -1)
+        out[in_range] = generations[index[in_range]]
+        return out
+
     def state_dict(self) -> dict[str, Any]:
-        return {"_cursor": self._cursor, "_write_count": self._write_count}
+        generations = self._slot_generations
+        return {
+            "_cursor": self._cursor,
+            "_write_count": self._write_count,
+            "_slot_generations": None if generations is None else generations.clone(),
+        }
 
     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
+        generations = state_dict.get("_slot_generations")
+        if generations is not None:
+            self._slot_generations = generations.clone()
 
     def _empty(self, empty_write_count: bool = True) -> None:
         self._cursor = 0
         if empty_write_count:
             self._write_count = 0
+        if self._slot_generations is not None:
+            self._slot_generations += 1
 
     # TODO: Workaround for PyTorch nightly regression where compiler can't handle
     # method calls on objects returned from _attached_entities_iter()
@@ -355,6 +490,7 @@ def add(self, data: Any) -> int | torch.Tensor:
         max_size_along_dim0 = self._storage._max_size_along_dim0(single_data=data)
         self._cursor = (index + 1) % max_size_along_dim0
         self._write_count += 1
+        self._bump_generations(index, max_size_along_dim0)
         if not is_tensorclass(data):
             data.set(
                 "index",
@@ -381,6 +517,7 @@ def extend(self, data: Sequence) -> torch.Tensor:
         # we need to update the cursor first to avoid race conditions between workers
         self._cursor = (batch_size + cur_size) % max_size_along_dim0
         self._write_count += batch_size
+        self._bump_generations(index, max_size_along_dim0)
         # storage must convert the data to the appropriate format if needed
         if not is_tensorclass(data):
             data.set(

From 8ff3698cfbfbcc33cd279fd59e5ff3bb3f782d6f Mon Sep 17 00:00:00 2001
From: Achintya P <apaninga@berkeley.edu>
Date: Thu, 23 Jul 2026 17:04:25 -0700
Subject: [PATCH 3/4] [Test] Spec tests for ReplayBuffer.update_if_present

Executable spec for step 2 of the conditional replay-update RFC:
rb.update_if_present(index=, generation=, patch=) applies every patch
key to records whose (index, generation) is still live, skips reused
or emptied slots without touching their content, and returns a result
with an updated mask aligned to the input order plus updated/stale
counts. The whole patch is validated before any write (KeyError for
unknown keys, ValueError for shape or dtype mismatches, storage
untouched in both cases), handles survive repeated updates, nested
keys are supported, ListStorage raises a capability error, multidim
storages round-trip, sampled handles flow straight into the call, and
a concurrent writer/updater stress test pins non-torn multi-key
visibility. RayReplayBuffer delegates the call to the actor.

Tests are expected to fail until the implementation lands.
---
 test/rb/test_rb_core.py        | 251 +++++++++++++++++++++++++++++++++
 test/rb/test_rb_distributed.py |  41 ++++++
 2 files changed, 292 insertions(+)

diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py
index 883cbd5daba..464906b42d5 100644
--- a/test/rb/test_rb_core.py
+++ b/test/rb/test_rb_core.py
@@ -8,6 +8,7 @@
 import contextlib
 import functools
 import json
+import threading
 
 import pytest
 import torch
@@ -1243,6 +1244,256 @@ def test_stale_handles_detectable_after_wraparound(self):
         assert live[4:].all()
 
 
+class TestUpdateIfPresent:
+    """Executable spec for ReplayBuffer.update_if_present (RFC step 2).
+
+    Contract pinned by this class:
+
+    - Signature: ``rb.update_if_present(index=..., generation=..., patch=...)``
+      with keyword-only arguments. ``patch`` maps tensordict keys (flat or
+      nested) to tensors whose leading dimension equals ``len(index)``.
+    - Records whose (index, generation) pair is still live receive every
+      patch key; records whose slot was reused or emptied are skipped and
+      their current content is never modified.
+    - The result exposes ``updated`` (bool tensor aligned with the input
+      index order), ``updated_count`` and ``stale_count``.
+    - Updating a record does not consume its handle: the same (index,
+      generation) pair keeps working until the slot is rewritten.
+    - The whole patch is validated before any write: an unknown key raises
+      ``KeyError``, a shape or dtype mismatch raises ``ValueError``, and in
+      both cases storage is left byte-for-byte untouched, even when other
+      keys of the same patch were valid.
+    - Storages that cannot validate generations raise a capability error
+      mentioning "conditional" instead of writing through raw indices.
+    """
+
+    def _make_rb(self, size=10):
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(size), batch_size=4)
+        data = TensorDict(
+            {
+                "obs": torch.arange(size, dtype=torch.float32)
+                .unsqueeze(-1)
+                .expand(size, 3)
+                .clone(),
+                "info": {"label": torch.zeros(size, dtype=torch.int64)},
+            },
+            batch_size=[size],
+        )
+        index = rb.extend(data)
+        generation = rb._writer.generations_of(index)
+        return rb, data, index, generation
+
+    def test_updates_live_records(self):
+        rb, _, index, generation = self._make_rb()
+        patch = {"obs": torch.full((10, 3), 42.0)}
+        result = rb.update_if_present(index=index, generation=generation, patch=patch)
+        assert result.updated.dtype == torch.bool
+        assert result.updated.all()
+        assert result.updated_count == 10
+        assert result.stale_count == 0
+        torch.testing.assert_close(rb[:]["obs"], patch["obs"])
+
+    def test_stale_records_skipped_and_unmodified(self):
+        rb, _, index, generation = self._make_rb()
+        overwrite = TensorDict(
+            {
+                "obs": torch.full((4, 3), -1.0),
+                "info": {"label": torch.ones(4, dtype=torch.int64)},
+            },
+            batch_size=[4],
+        )
+        rb.extend(overwrite)
+        result = rb.update_if_present(
+            index=index,
+            generation=generation,
+            patch={"obs": torch.full((10, 3), 42.0)},
+        )
+        assert result.updated.tolist() == [False] * 4 + [True] * 6
+        assert result.updated_count == 6
+        assert result.stale_count == 4
+        torch.testing.assert_close(rb[:]["obs"][:4], overwrite["obs"])
+        torch.testing.assert_close(rb[:]["obs"][4:], torch.full((6, 3), 42.0))
+
+    def test_mask_aligns_with_input_order(self):
+        rb, _, index, generation = self._make_rb()
+        rb.extend(
+            TensorDict(
+                {
+                    "obs": torch.zeros(4, 3),
+                    "info": {"label": torch.zeros(4, dtype=torch.int64)},
+                },
+                batch_size=[4],
+            )
+        )
+        permutation = torch.tensor([7, 0, 5, 2, 9, 1])
+        result = rb.update_if_present(
+            index=index[permutation],
+            generation=generation[permutation],
+            patch={"obs": torch.full((6, 3), 42.0)},
+        )
+        expected_live = (permutation >= 4).tolist()
+        assert result.updated.tolist() == expected_live
+
+    def test_handle_survives_repeated_updates(self):
+        rb, _, index, generation = self._make_rb()
+        for value in (1.0, 2.0):
+            result = rb.update_if_present(
+                index=index,
+                generation=generation,
+                patch={"obs": torch.full((10, 3), value)},
+            )
+            assert result.updated.all()
+        torch.testing.assert_close(rb[:]["obs"], torch.full((10, 3), 2.0))
+
+    def test_unknown_key_raises_and_storage_untouched(self):
+        rb, data, index, generation = self._make_rb()
+        with pytest.raises(KeyError):
+            rb.update_if_present(
+                index=index,
+                generation=generation,
+                patch={"not_a_key": torch.zeros(10, 3)},
+            )
+        torch.testing.assert_close(rb[:]["obs"], data["obs"])
+
+    def test_invalid_shape_or_dtype_raises_before_any_write(self):
+        rb, data, index, generation = self._make_rb()
+        with pytest.raises(ValueError):
+            rb.update_if_present(
+                index=index,
+                generation=generation,
+                patch={
+                    "obs": torch.full((10, 3), 42.0),
+                    ("info", "label"): torch.zeros(10, 5, dtype=torch.int64),
+                },
+            )
+        with pytest.raises(ValueError):
+            rb.update_if_present(
+                index=index,
+                generation=generation,
+                patch={"obs": torch.zeros(10, 3, dtype=torch.int64)},
+            )
+        torch.testing.assert_close(rb[:]["obs"], data["obs"])
+        torch.testing.assert_close(rb[:]["info", "label"], data["info", "label"])
+
+    def test_nested_key_patch(self):
+        rb, _, index, generation = self._make_rb()
+        result = rb.update_if_present(
+            index=index,
+            generation=generation,
+            patch={("info", "label"): torch.full((10,), 7, dtype=torch.int64)},
+        )
+        assert result.updated.all()
+        assert (rb[:]["info", "label"] == 7).all()
+
+    def test_capability_error_on_list_storage(self):
+        rb = ReplayBuffer(storage=ListStorage(10))
+        index = rb.extend([torch.randn(3) for _ in range(5)])
+        with pytest.raises(
+            (RuntimeError, TypeError, NotImplementedError), match="(?i)conditional"
+        ):
+            rb.update_if_present(
+                index=torch.as_tensor(index),
+                generation=torch.zeros(5, dtype=torch.int64),
+                patch={"obs": torch.zeros(5, 3)},
+            )
+
+    def test_empty_invalidates_handles(self):
+        rb, _, index, generation = self._make_rb()
+        rb.empty()
+        rb.extend(
+            TensorDict(
+                {
+                    "obs": torch.zeros(10, 3),
+                    "info": {"label": torch.zeros(10, dtype=torch.int64)},
+                },
+                batch_size=[10],
+            )
+        )
+        result = rb.update_if_present(
+            index=index,
+            generation=generation,
+            patch={"obs": torch.full((10, 3), 42.0)},
+        )
+        assert not result.updated.any()
+        assert result.stale_count == 10
+        assert (rb[:]["obs"] == 0).all()
+
+    def test_sampled_handles_roundtrip(self):
+        rb, _, _, _ = self._make_rb()
+        sample = rb.sample()
+        batch = sample.batch_size[0]
+        index = sample.get("index").reshape(batch, -1)[:, 0]
+        generation = sample.get("index_generation").reshape(batch, -1)[:, 0]
+        marker = torch.full((batch, 3), 123.0)
+        result = rb.update_if_present(
+            index=index, generation=generation, patch={"obs": marker}
+        )
+        assert result.updated.all()
+        torch.testing.assert_close(rb[:]["obs"][index], marker)
+
+    def test_multidim_storage_roundtrip(self):
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(6, ndim=2), batch_size=4)
+        data = TensorDict(
+            {"obs": torch.arange(6, dtype=torch.float32).reshape(2, 3)},
+            batch_size=[2, 3],
+        )
+        index = rb.extend(data)
+        generation = rb._writer.generations_of(index)
+        result = rb.update_if_present(
+            index=index,
+            generation=generation,
+            patch={"obs": torch.full_like(data["obs"], 42.0)},
+        )
+        assert result.updated.all()
+        assert (rb[:]["obs"] == 42.0).all()
+
+    def test_concurrent_updates_do_not_tear_records(self):
+        rb = TensorDictReplayBuffer(storage=LazyTensorStorage(64), batch_size=8)
+        rb.extend(
+            TensorDict({"a": torch.zeros(64), "b": torch.zeros(64)}, batch_size=[64])
+        )
+        stop = threading.Event()
+        errors = []
+
+        def writer_loop():
+            value = 1.0
+            try:
+                while not stop.is_set():
+                    rb.extend(
+                        TensorDict(
+                            {
+                                "a": torch.full((8,), value),
+                                "b": torch.full((8,), value),
+                            },
+                            batch_size=[8],
+                        )
+                    )
+                    value += 1.0
+            except Exception as err:
+                errors.append(err)
+
+        thread = threading.Thread(target=writer_loop)
+        thread.start()
+        try:
+            for step in range(200):
+                sample = rb.sample()
+                batch = sample.batch_size[0]
+                index = sample.get("index").reshape(batch, -1)[:, 0]
+                generation = sample.get("index_generation").reshape(batch, -1)[:, 0]
+                marker = torch.full((batch,), 10_000.0 + step)
+                rb.update_if_present(
+                    index=index,
+                    generation=generation,
+                    patch={"a": marker, "b": marker},
+                )
+                content = rb[:]
+                torch.testing.assert_close(content["a"], content["b"])
+        finally:
+            stop.set()
+            thread.join(timeout=10)
+        assert not errors
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
diff --git a/test/rb/test_rb_distributed.py b/test/rb/test_rb_distributed.py
index 6e2ea620130..e829d87bf6a 100644
--- a/test/rb/test_rb_distributed.py
+++ b/test/rb/test_rb_distributed.py
@@ -213,6 +213,47 @@ def test_ray_rb_stats(self):
         finally:
             rb.close()
 
+    def test_ray_rb_update_if_present(self):
+        """Spec: update_if_present is delegated to the actor in one RPC.
+
+        The remote buffer is a TensorDictReplayBuffer so samples carry the
+        index and index_generation keys; the conditional update validates
+        and writes inside the actor, and stale handles created by a
+        wraparound are skipped exactly as in the local contract.
+        """
+        from torchrl.data import TensorDictReplayBuffer
+
+        rb = RayReplayBuffer(
+            replay_buffer_cls=TensorDictReplayBuffer,
+            storage=partial(LazyTensorStorage, 10),
+            batch_size=4,
+            ray_init_config={"num_cpus": 1},
+        )
+        try:
+            index = rb.extend(TensorDict({"x": torch.zeros(10, 2)}, batch_size=10))
+            index = torch.as_tensor(index).reshape(-1)
+            sample = rb.sample()
+            batch = sample.batch_size[0]
+            sampled_index = sample.get("index").reshape(batch, -1)[:, 0]
+            generation = sample.get("index_generation").reshape(batch, -1)[:, 0]
+            marker = torch.full((batch, 2), 42.0)
+            result = rb.update_if_present(
+                index=sampled_index, generation=generation, patch={"x": marker}
+            )
+            assert result.updated.all()
+            assert result.updated_count == batch
+            rb.extend(TensorDict({"x": torch.ones(4, 2)}, batch_size=4))
+            stale = rb.update_if_present(
+                index=index[:4],
+                generation=torch.zeros(4, dtype=torch.int64),
+                patch={"x": torch.full((4, 2), -5.0)},
+            )
+            assert not stale.updated.any()
+            assert stale.stale_count == 4
+            assert (rb[:4]["x"] == 1.0).all()
+        finally:
+            rb.close()
+
     def test_ray_rb_iter(self):
         rb = RayReplayBuffer(
             storage=partial(LazyTensorStorage, 100),

From 11715f694a0ca02262507fe069e556d46ad7021c Mon Sep 17 00:00:00 2001
From: Achintya P <apaninga@berkeley.edu>
Date: Fri, 24 Jul 2026 00:47:21 -0700
Subject: [PATCH 4/4] [Feature] Add ReplayBuffer.update_if_present for
 generation-safe conditional updates

Adds a best-effort conditional mutation API for stored replay fields.
update_if_present(index=, generation=, patch=) applies a patch only to
records whose (index, generation) pair still matches the writer's
current slot generation, skipping records whose slot was recycled or
emptied instead of corrupting them, and returns a
ConditionalUpdateResult with a per-record updated mask plus
updated/stale counts. The whole patch is validated (key existence,
shape, dtype) before any write; validation failures leave storage
untouched. The generation comparison and the patch write share one
replay-lock acquisition, giving per-record atomicity against
concurrent extends, and updating a record does not consume its handle.

Tensor storages advertise supports_conditional_update; unsupported
backends such as ListStorage raise a capability error instead of
performing an unsafe raw-index write. RayReplayBuffer delegates the
call to the actor in a single RPC (validation and write run inside the
actor under its own lock); the distributed transport raises a clear
capability error. Nested keys and multidimensional storages are
supported.

Second step of the conditional replay-update RFC. Closes #4040.
---
 docs/source/reference/data_replaybuffers.rst  |  32 +++++
 torchrl/data/__init__.py                      |   2 +
 torchrl/data/replay_buffers/__init__.py       |   2 +
 torchrl/data/replay_buffers/ray_buffer.py     |  24 ++++
 torchrl/data/replay_buffers/replay_buffers.py | 117 +++++++++++++++++-
 torchrl/data/replay_buffers/storages.py       |  56 +++++++++
 6 files changed, 232 insertions(+), 1 deletion(-)

diff --git a/docs/source/reference/data_replaybuffers.rst b/docs/source/reference/data_replaybuffers.rst
index 8ca46081ad5..a1e52707262 100644
--- a/docs/source/reference/data_replaybuffers.rst
+++ b/docs/source/reference/data_replaybuffers.rst
@@ -49,6 +49,38 @@ discovery and buffer lifecycle.
     RemoteTensorDictReplayBuffer
 
 
+Conditional record updates
+--------------------------
+
+Round-robin writers recycle storage slots, so a physical index captured at
+sampling time can point to a different record by the time an asynchronous
+computation writes back. Replay slots therefore carry a generation counter:
+samples expose it as an ``"index_generation"`` entry next to ``"index"``, and
+:meth:`~torchrl.data.ReplayBuffer.update_if_present` applies a patch only to
+records whose ``(index, generation)`` pair is still live, skipping recycled
+slots instead of corrupting them. This supports algorithms that refresh
+stored fields after sampling, such as recurrent-state refreshes or
+asynchronously computed labels, without pinning the buffer or racing against
+collection.
+
+.. code-block:: python
+
+    sample = buffer.sample()
+    refreshed = compute_refreshed_state(sample)
+    result = buffer.update_if_present(
+        index=sample["index"],
+        generation=sample["index_generation"],
+        patch={"recurrent_state": refreshed},
+    )
+    print(f"updated {result.updated_count}, skipped {result.stale_count} stale records")
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template.rst
+
+    ConditionalUpdateResult
+
+
 Offline-to-online helpers
 -------------------------
 
diff --git a/torchrl/data/__init__.py b/torchrl/data/__init__.py
index 1fe24678119..1cb14713df2 100644
--- a/torchrl/data/__init__.py
+++ b/torchrl/data/__init__.py
@@ -34,6 +34,7 @@
 from .replay_buffers import (
     CompressedListStorage,
     CompressedListStorageCheckpointer,
+    ConditionalUpdateResult,
     ConsumingSampler,
     DEFAULT_DONE_KEYS,
     filter_trajectories,
@@ -176,6 +177,7 @@
     "RandomSampler",
     "RayReplayBuffer",
     "RemoteTensorDictReplayBuffer",
+    "ConditionalUpdateResult",
     "ReplayBuffer",
     "ReplayBufferEnsemble",
     "RewardData",
diff --git a/torchrl/data/replay_buffers/__init__.py b/torchrl/data/replay_buffers/__init__.py
index 1d258b8d78b..0d3df5426cb 100644
--- a/torchrl/data/replay_buffers/__init__.py
+++ b/torchrl/data/replay_buffers/__init__.py
@@ -24,6 +24,7 @@
 )
 from .ray_buffer import RayReplayBuffer
 from .replay_buffers import (
+    ConditionalUpdateResult,
     PrioritizedReplayBuffer,
     RemoteTensorDictReplayBuffer,
     ReplayBuffer,
@@ -94,6 +95,7 @@
     "RayReplayBuffer",
     "PrioritizedReplayBuffer",
     "RemoteTensorDictReplayBuffer",
+    "ConditionalUpdateResult",
     "ReplayBuffer",
     "ReplayBufferEnsemble",
     "TensorDictPrioritizedReplayBuffer",
diff --git a/torchrl/data/replay_buffers/ray_buffer.py b/torchrl/data/replay_buffers/ray_buffer.py
index 2ac1ffa3f05..67fdf21c3c8 100644
--- a/torchrl/data/replay_buffers/ray_buffer.py
+++ b/torchrl/data/replay_buffers/ray_buffer.py
@@ -77,6 +77,13 @@ def write_count(self):
     def stats(self):
         return ray.get(self._actor.stats.remote())
 
+    def update_if_present(self, *, index, generation, patch):
+        return ray.get(
+            self._actor.update_if_present.remote(
+                index=index, generation=generation, patch=patch
+            )
+        )
+
     @property
     def dim_extend(self):
         return ray.get(self._actor._getattr.remote("dim_extend"))
@@ -187,6 +194,12 @@ def stats(self, *, timeout: float | None = None) -> dict[str, int | float | bool
         snapshot = self._stats(timeout=timeout)
         return {key: value.item() for key, value in snapshot.items()}
 
+    def update_if_present(self, *, index, generation, patch):
+        raise RuntimeError(
+            "Conditional updates are not supported by the distributed replay "
+            "transport. Use transport='ray' for update_if_present."
+        )
+
     def extend(self, data: TensorDictBase, *, timeout: float | None = None):
         if self._extend_client is None:
             self._extend_client, result, handled = ray.get(
@@ -556,6 +569,17 @@ def stats(self) -> dict[str, int | float | bool]:
         """
         return self._client.stats()
 
+    def update_if_present(self, *, index, generation, patch):
+        """Conditionally updates live records through a single actor round-trip.
+
+        Validation, the generation comparison and the patch write all run
+        inside the replay-buffer actor under its own lock.
+        See :meth:`~torchrl.data.ReplayBuffer.update_if_present`.
+        """
+        return self._client.update_if_present(
+            index=index, generation=generation, patch=patch
+        )
+
     @property
     def dim_extend(self):
         return self._client.dim_extend
diff --git a/torchrl/data/replay_buffers/replay_buffers.py b/torchrl/data/replay_buffers/replay_buffers.py
index 4ae45c85e42..3f7288476c1 100644
--- a/torchrl/data/replay_buffers/replay_buffers.py
+++ b/torchrl/data/replay_buffers/replay_buffers.py
@@ -11,7 +11,7 @@
 import textwrap
 import threading
 import warnings
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from multiprocessing.context import get_spawning_popen
 from pathlib import Path
@@ -33,6 +33,7 @@
     is_tensorclass,
     LazyStackedTensorDict,
     NestedKey,
+    TensorClass,
     TensorDict,
     TensorDictBase,
     unravel_key,
@@ -123,6 +124,30 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
+class ConditionalUpdateResult(TensorClass["nocast"]):
+    """Result of :meth:`ReplayBuffer.update_if_present`.
+
+    Attributes:
+        updated (torch.Tensor): boolean mask aligned with the order of the
+            indices passed to the update. ``True`` marks records that were
+            still live and received the patch; ``False`` marks stale records
+            whose slot had been reused or emptied and whose content was left
+            untouched.
+    """
+
+    updated: torch.Tensor
+
+    @property
+    def updated_count(self) -> int:
+        """Number of records that were live and patched."""
+        return int(self.updated.sum().item())
+
+    @property
+    def stale_count(self) -> int:
+        """Number of records that were stale and skipped."""
+        return int(self.updated.numel()) - self.updated_count
+
+
 class ReplayBuffer(metaclass=_RayServiceMetaClass):
     """A generic, composable replay buffer class.
 
@@ -998,6 +1023,96 @@ def stats(self) -> dict[str, int | float | bool]:
             stats["utilization"] = float(size) / capacity if capacity else 0.0
         return stats
 
+    def update_if_present(
+        self,
+        *,
+        index: torch.Tensor,
+        generation: torch.Tensor,
+        patch: Mapping[NestedKey, torch.Tensor] | TensorDictBase,
+    ) -> ConditionalUpdateResult:
+        """Conditionally updates stored records that are still live.
+
+        Replay slots are recycled by round-robin writers, so a physical index
+        captured at sampling time can point to a different record by the time
+        an asynchronous computation writes back. This method applies ``patch``
+        only to records whose ``(index, generation)`` pair still matches the
+        writer's current slot generation, skipping records whose slot was
+        reused or emptied since the handle was captured. Skipped records are
+        never modified.
+
+        The whole patch is validated (key existence, shape and dtype) before
+        any write happens; a validation failure leaves the storage untouched.
+        Updating a record refreshes its content, not its identity: the same
+        handle keeps working until the slot is rewritten by ``add``,
+        ``extend`` or ``empty``.
+
+        Keyword Args:
+            index (torch.Tensor): storage indices, as returned by
+                :meth:`extend` or found in the sample under ``"index"``.
+            generation (torch.Tensor): slot generations captured with the
+                indices, as found in the sample under ``"index_generation"``.
+            patch (mapping of NestedKey to torch.Tensor, or TensorDictBase):
+                the fields to overwrite for live records. Leading dimension
+                must match the number of records addressed by ``index``.
+
+        Returns:
+            A :class:`ConditionalUpdateResult` whose ``updated`` mask is
+            aligned with the input index order, with ``updated_count`` and
+            ``stale_count`` conveniences.
+
+        Raises:
+            RuntimeError: if the storage or writer does not support
+                conditional updates (for example :class:`ListStorage`).
+            KeyError: if a patch key does not exist in the storage.
+            ValueError: if a patch entry has an incompatible shape or dtype.
+
+        Examples:
+            >>> import torch
+            >>> from tensordict import TensorDict
+            >>> from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
+            >>> rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10), batch_size=4)
+            >>> rb.extend(TensorDict({"obs": torch.zeros(10, 3)}, batch_size=[10]))
+            >>> sample = rb.sample()
+            >>> result = rb.update_if_present(
+            ...     index=sample["index"],
+            ...     generation=sample["index_generation"],
+            ...     patch={"obs": torch.ones(4, 3)},
+            ... )
+            >>> print(result.updated_count, result.stale_count)
+            4 0
+        """
+        storage = self._storage
+        if not getattr(storage, "supports_conditional_update", False) or not getattr(
+            self._writer, "tracks_generations", False
+        ):
+            raise RuntimeError(
+                f"Conditional updates are not supported by {type(storage).__name__} "
+                f"with {type(self._writer).__name__}: the storage must support "
+                "conditional updates and the writer must track slot generations."
+            )
+        index = torch.as_tensor(index, dtype=torch.long)
+        dim0 = index[..., 0] if index.ndim > 1 else index.reshape(-1)
+        generation = torch.as_tensor(generation, dtype=torch.long).reshape(-1)
+        if generation.numel() != dim0.numel():
+            raise ValueError(
+                f"index and generation must address the same number of records, "
+                f"got {dim0.numel()} indices and {generation.numel()} generations."
+            )
+        if isinstance(patch, TensorDictBase):
+            patch = dict(patch.items(include_nested=True, leaves_only=True))
+        else:
+            patch = dict(patch)
+        normalized = storage._validate_conditional_patch(index, patch)
+        with self._replay_lock, self._write_lock:
+            live = self._writer.generations_of(dim0) == generation
+            if live.any():
+                live_index = index[live]
+                storage._apply_conditional_patch(
+                    live_index,
+                    {key: value[live] for key, value in normalized.items()},
+                )
+        return ConditionalUpdateResult(updated=live, batch_size=live.shape)
+
     def __repr__(self) -> str:
         from torchrl.envs.transforms import Compose
 
diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py
index fe8836980ae..a3ad1e41566 100644
--- a/torchrl/data/replay_buffers/storages.py
+++ b/torchrl/data/replay_buffers/storages.py
@@ -30,6 +30,7 @@
     is_tensor_collection,
     lazy_stack,
     LazyStackedTensorDict,
+    NestedKey,
     TensorDict,
     TensorDictBase,
 )
@@ -181,6 +182,7 @@ class Storage:
 
     ndim = 1
     max_size: int
+    supports_conditional_update: bool = False
     _default_checkpointer: StorageCheckpointerBase = StorageCheckpointerBase
     _rng: torch.Generator | None = None
 
@@ -712,6 +714,7 @@ class TensorStorage(Storage):
 
     _storage = None
     _default_checkpointer = TensorStorageCheckpointer
+    supports_conditional_update = True
 
     def __init__(
         self,
@@ -906,6 +909,59 @@ def flatten(self):
             )
         )
 
+    def _conditional_patch_leaf(self, key: NestedKey) -> torch.Tensor:
+        storage = getattr(self, "_storage", None)
+        if storage is None or not self.initialized:
+            raise RuntimeError(
+                "Conditional updates require an initialized storage. Write some "
+                "data to the buffer before calling update_if_present."
+            )
+        leaf = None
+        if is_tensor_collection(storage):
+            leaf = storage.get(key, default=None)
+        if leaf is None:
+            raise KeyError(
+                f"Key {key} does not exist in the storage. Conditional patches "
+                "can only target existing tensor fields of a tensordict storage."
+            )
+        return leaf
+
+    def _validate_conditional_patch(
+        self, index: torch.Tensor, patch: dict[NestedKey, torch.Tensor]
+    ) -> dict[NestedKey, torch.Tensor]:
+        n_coords = index.shape[-1] if index.ndim > 1 else 1
+        n_rows = index.shape[0] if index.ndim > 1 else index.numel()
+        normalized = {}
+        for key, value in patch.items():
+            leaf = self._conditional_patch_leaf(key)
+            value = torch.as_tensor(value)
+            if value.dtype != leaf.dtype:
+                raise ValueError(
+                    f"dtype mismatch for patch key {key}: got {value.dtype}, "
+                    f"the storage holds {leaf.dtype}."
+                )
+            feature_shape = leaf.shape[n_coords:]
+            try:
+                value = value.reshape((n_rows, *feature_shape))
+            except RuntimeError:
+                raise ValueError(
+                    f"shape mismatch for patch key {key}: got {tuple(value.shape)}, "
+                    f"expected {n_rows} records with feature shape {tuple(feature_shape)}."
+                )
+            normalized[key] = value.to(leaf.device)
+        return normalized
+
+    def _apply_conditional_patch(
+        self, index: torch.Tensor, patch: dict[NestedKey, torch.Tensor]
+    ) -> None:
+        if index.ndim > 1:
+            coords = tuple(index.unbind(-1))
+        else:
+            coords = (index,)
+        for key, value in patch.items():
+            leaf = self._conditional_patch_leaf(key)
+            leaf[coords] = value
+
     def __getstate__(self):
         state = super().__getstate__()
         if get_spawning_popen() is None: