diff --git a/docs/source/reference/data_replaybuffers.rst b/docs/source/reference/data_replaybuffers.rst index 8ca46081ad5..8287afa812a 100644 --- a/docs/source/reference/data_replaybuffers.rst +++ b/docs/source/reference/data_replaybuffers.rst @@ -184,6 +184,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 `. .. autosummary:: :toctree: generated/ diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py index ffe47a51ffb..27f81ab4c2c 100644 --- a/test/rb/test_writers.py +++ b/test/rb/test_writers.py @@ -403,6 +403,410 @@ def test_roundrobin_dumps_loads_write_count(self, tmp_path): assert writer2._write_count == 23 +class TestWriterGeneration: + 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) + assert gen.dtype == torch.int64 + 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(12, ndim=2), + writer=TensorDictRoundRobinWriter(track_generations=True), + ) + index = rb.extend(TensorDict({"a": torch.zeros(4, 3)}, [4, 3])) + assert rb._storage.ndim == 2 + # 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 + 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), + writer=RoundRobinWriter(track_generations=True), + ) + rb.extend(torch.arange(size)) + torch.testing.assert_close( + rb._writer.generations_of(torch.arange(size)), + 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([1, 1, 1, 0]) + ) + + def test_generation_wraparound(self): + size = 5 + 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)), + torch.full((size,), 1, dtype=torch.int64), + ) + + def test_generation_extend_wrapping_twice(self): + size = 4 + 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( + 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), + writer=RoundRobinWriter(track_generations=True), + ) + for i in range(size + 1): + rb.add(torch.tensor(i)) + torch.testing.assert_close( + rb._writer.generations_of(torch.arange(size)), torch.tensor([1, 0, 0]) + ) + + def test_generations_of_unwritten_reports_minus_one(self): + 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), + 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)), + torch.full((size,), 1, dtype=torch.int64), + ) + + def test_generation_write_at(self): + storage = LazyTensorStorage(4) + 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])) + torch.testing.assert_close( + writer.generations_of(torch.arange(4)), torch.tensor([1, 1, 0, 0]) + ) + + def test_empty_is_monotonic(self): + rb = ReplayBuffer( + storage=LazyTensorStorage(10), + writer=RoundRobinWriter(track_generations=True), + ) + 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), + 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), + writer=RoundRobinWriter(track_generations=True), + ) + 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), + writer=RoundRobinWriter(track_generations=True), + ) + rb.extend(torch.arange(size + 1)) + sd = rb.state_dict() + 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)), + rb._writer.generations_of(torch.arange(size)), + ) + + def test_legacy_state_dict_without_generation_loads(self): + 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), + 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(track_generations=True) + writer._cursor = 2 + writer._write_count = 9 + writer._generation = torch.tensor([3, 2, 2, 1]) + writer.dumps(tmp_path) + writer2 = RoundRobinWriter(track_generations=True) + writer2.loads(tmp_path) + assert writer2._cursor == 2 + assert writer2._write_count == 9 + 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_generations=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.generations_of(idx)) + + 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=RoundRobinWriter(track_generations=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_generations=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.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_generations=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.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), + 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) + 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), + writer=RoundRobinWriter(track_generations=True), + ) + 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() pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown) diff --git a/torchrl/data/replay_buffers/replay_buffers.py b/torchrl/data/replay_buffers/replay_buffers.py index a5d781e71df..64942fbace7 100644 --- a/torchrl/data/replay_buffers/replay_buffers.py +++ b/torchrl/data/replay_buffers/replay_buffers.py @@ -1516,6 +1516,8 @@ 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) data = self._storage.get(_storage_index(index, self._storage)) if not isinstance(index, INT_CLASSES): data = self._collate_fn(data) @@ -2135,6 +2137,8 @@ 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) data = self._storage.get(_storage_index(index, self._storage)) if not isinstance(index, INT_CLASSES): data = self._collate_fn(data) @@ -2564,6 +2568,8 @@ 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) data = self._storage.get(_storage_index(index, self._storage)) if not isinstance(index, INT_CLASSES): data = self._collate_fn(data) @@ -2929,6 +2935,8 @@ 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) 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..15c11266dc5 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 @@ -39,6 +39,21 @@ def tree_leaves(data): # noqa: D103 from torchrl.data.replay_buffers.storages import Storage from torchrl.data.replay_buffers.utils import _is_int, _reduce +# 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): """A ReplayBuffer base Writer class.""" @@ -50,9 +65,43 @@ def __init__(self, compilable: bool = False) -> None: self._storage = None self._compilable = compilable + #: 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: self._storage = storage + def generations_of(self, index: int | torch.Tensor) -> torch.Tensor: + """Returns the generation stamp for each physical slot in ``index``. + + 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) + @abstractmethod def add(self, data: Any) -> int: """Inserts one piece of data at an appropriate index, and returns that index.""" @@ -153,18 +202,195 @@ 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. + + .. 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`. + + 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) -> None: + def __init__( + self, compilable: bool = False, *, track_generations: bool = False + ) -> None: super().__init__(compilable=compilable) self._cursor = 0 self._write_count # noqa + 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: + # 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: + 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 + 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 + # 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) + self._ensure_generation(capacity, int(index) + 1, device) + self._generation[int(index)] += 1 + else: + 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) + 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_( + (index,), torch.ones_like(index), accumulate=True + ) + + 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) + # 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 + ): + index = index[..., 0] + index = torch.as_tensor(index, dtype=torch.long) + if self._generation is None: + 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)) + return gen.to(index.device) def dumps(self, path): path = Path(path).absolute() path.mkdir(exist_ok=True) + metadata = { + "cursor": self._cursor, + "write_count": self._write_count, + } + generation = self._generation if self._track_generations else None + if generation is not None: + generation = generation.cpu() + 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 +400,15 @@ def loads(self, path): write_count = metadata.get("write_count") if write_count is not None: self._write_count = write_count + 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() + self._generation = generation + self._align_generation_device() def add(self, data: Any) -> int | torch.Tensor: index = self._cursor @@ -186,6 +421,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,12 +450,17 @@ 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 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: @@ -229,6 +470,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 +491,29 @@ 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_generations and 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 + generation = state_dict.get("_generation") + if generation is not None: + self._generation = generation.clone() + self._align_generation_device() def _empty(self, empty_write_count: bool = True) -> None: self._cursor = 0 + generation = self._generation if self._track_generations else None + 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. + # Never-written slots keep the -1 sentinel. + generation[generation >= 0] += 1 if empty_write_count: self._write_count = 0 @@ -347,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 @@ -363,6 +623,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 +653,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 +669,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)