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_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_samplers.py b/test/rb/test_samplers.py index 2e410f160a6..9abb6f14581 100644 --- a/test/rb/test_samplers.py +++ b/test/rb/test_samplers.py @@ -243,7 +243,7 @@ def test_sampler_without_rep_state_dict(self, backend): replay_buffer.extend(transition.clone()) for _ in range(n_samples): s = replay_buffer.sample(batch_size=1) - assert (s.exclude("index") == 1).all() + assert (s.exclude("index", "index_generation") == 1).all() replay_buffer.extend(torch.zeros_like(transition)) @@ -257,7 +257,7 @@ def test_sampler_without_rep_state_dict(self, backend): new_replay_buffer.load_state_dict(state_dict) s = new_replay_buffer.sample(batch_size=1) - assert (s.exclude("index") == 0).all() + assert (s.exclude("index", "index_generation") == 0).all() def test_sampler_without_rep_dumps_loads(self, tmpdir): d0 = tmpdir + "/save0" diff --git a/test/rb/test_storages.py b/test/rb/test_storages.py index 228733e8a8f..76932a87cb0 100644 --- a/test/rb/test_storages.py +++ b/test/rb/test_storages.py @@ -306,7 +306,7 @@ def test_storage_state_dict(self, storage_in, storage_out, init_out, backend): new_replay_buffer.load_state_dict(state_dict) s = new_replay_buffer.sample() - assert (s.exclude("index") == 1).all() + assert (s.exclude("index", "index_generation") == 1).all() @pytest.mark.skipif( TORCH_VERSION < version.parse("2.5.0"), reason="requires Torch >= 2.5.0" diff --git a/test/rb/test_writers.py b/test/rb/test_writers.py index 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) 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(