From 3fffeeea0702b46b2b2e7675ef6371bc197026fe Mon Sep 17 00:00:00 2001 From: Achintya P Date: Sat, 25 Jul 2026 01:04:46 -0700 Subject: [PATCH 1/3] [Test] Spec tests for the Sequence sample unit Executable contract for piece 2 of the #4039 split: Sequence(length, episode_boundary, done_key) expands each anchor into the length records that follow it in stored-time order, wrapping ring indices across the storage seam. Boundary policies: pad keeps the anchor and marks the tail past an episode end invalid with indices clamped inside the episode; stop shifts the anchor backward to end exactly at the boundary, falling back to pad for episodes shorter than length; include_reset crosses the boundary with all entries valid. The unit adds per-record sequence_id, step_in_sequence and validity_mask info entries that surface as TensorDict sample keys, expands per-anchor sampler entries such as prioritized weights to the record count, and sample(batch_size=B) returns B*length records. Constructor validates length and the boundary policy. Tests are expected to fail until the implementation lands. --- test/rb/test_rb_core.py | 179 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 179 insertions(+) diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py index 3786a93a83b..83027b0117b 100644 --- a/test/rb/test_rb_core.py +++ b/test/rb/test_rb_core.py @@ -1300,6 +1300,185 @@ def test_prioritized_buffer_with_transition_unit(self): rb.update_tensordict_priority(sample) +class TestSequenceUnit: + """Executable spec for the Sequence sample unit (#4039, piece 2). + + Contract pinned by this class: + + - ``Sequence(length, episode_boundary="pad", done_key=("next","done"))`` + expands each anchor into the ``length`` records that follow it in + stored-time order, wrapping physical ring indices when an episode spans + the storage seam. + - Boundary policies: ``"pad"`` keeps the anchor and marks entries past the + episode end invalid, clamping their indices inside the episode; + ``"stop"`` shifts the anchor backward so the sequence ends at the + boundary (full-length, fully valid), falling back to pad behavior when + the episode is shorter than ``length``; ``"include_reset"`` crosses the + boundary with all entries valid. + - ``expand`` adds per-record ``"sequence_id"``, ``"step_in_sequence"`` and + ``"validity_mask"`` entries to ``info``, and expands per-anchor entries + such as prioritized weights to the record count. + - ``sample(batch_size=B)`` therefore returns ``B * length`` records. + """ + + def _sequence_cls(self): + from torchrl.data.replay_buffers import sample_units + + return sample_units.Sequence + + def _make_storage(self, capacity=10, done_at=(5, 9)): + rb = TensorDictReplayBuffer(storage=LazyTensorStorage(capacity), batch_size=4) + size = capacity + done = torch.zeros(size, 1, dtype=torch.bool) + for idx in done_at: + done[idx] = True + rb.extend( + TensorDict( + { + "obs": torch.arange(size, dtype=torch.float32), + ("next", "done"): done, + }, + batch_size=[size], + ) + ) + return rb + + def _expand(self, rb, unit, anchors): + index, info = unit.expand( + torch.as_tensor(anchors, dtype=torch.long), {}, rb._storage + ) + return torch.as_tensor(index), info + + def test_expansion_is_consecutive_within_episode(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand(rb, Sequence(length=4), [1]) + assert index.tolist() == [1, 2, 3, 4] + assert info["sequence_id"].tolist() == [0, 0, 0, 0] + assert info["step_in_sequence"].tolist() == [0, 1, 2, 3] + assert info["validity_mask"].all() + + def test_pad_masks_tail_and_clamps_inside_episode(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand(rb, Sequence(length=4, episode_boundary="pad"), [4]) + assert index.tolist() == [4, 5, 5, 5] + assert info["validity_mask"].tolist() == [True, True, False, False] + obs = rb[:]["obs"][index] + assert obs.tolist() == [4.0, 5.0, 5.0, 5.0] + + def test_stop_shifts_anchor_to_end_at_boundary(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand(rb, Sequence(length=4, episode_boundary="stop"), [4]) + assert index.tolist() == [2, 3, 4, 5] + assert info["validity_mask"].all() + + def test_stop_falls_back_to_pad_for_short_episode(self): + Sequence = self._sequence_cls() + rb = self._make_storage(done_at=(1, 9)) + index, info = self._expand(rb, Sequence(length=4, episode_boundary="stop"), [0]) + assert index.tolist() == [0, 1, 1, 1] + assert info["validity_mask"].tolist() == [True, True, False, False] + + def test_include_reset_crosses_boundary(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand( + rb, Sequence(length=4, episode_boundary="include_reset"), [4] + ) + assert index.tolist() == [4, 5, 6, 7] + assert info["validity_mask"].all() + done = rb[:]["next", "done"].squeeze(-1)[index] + assert done.tolist() == [False, True, False, False] + + def test_wraparound_seam(self): + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10), batch_size=4) + size = 14 + rb.extend( + TensorDict( + { + "obs": torch.arange(size, dtype=torch.float32), + ("next", "done"): torch.zeros(size, 1, dtype=torch.bool), + }, + batch_size=[size], + ) + ) + index, info = self._expand( + rb, Sequence(length=4, episode_boundary="include_reset"), [8] + ) + assert index.tolist() == [8, 9, 0, 1] + obs = rb[:]["obs"][index] + assert obs.tolist() == [8.0, 9.0, 10.0, 11.0] + assert info["validity_mask"].all() + + def test_multiple_anchors_sequence_ids(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand(rb, Sequence(length=3), [0, 6]) + assert index.tolist() == [0, 1, 2, 6, 7, 8] + assert info["sequence_id"].tolist() == [0, 0, 0, 1, 1, 1] + assert info["step_in_sequence"].tolist() == [0, 1, 2, 0, 1, 2] + + def test_metadata_flows_into_tensordict_sample(self): + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer( + storage=LazyTensorStorage(20), + batch_size=2, + sample_unit=Sequence(length=4), + ) + rb.extend( + TensorDict( + { + "obs": torch.arange(20, dtype=torch.float32), + ("next", "done"): torch.zeros(20, 1, dtype=torch.bool), + }, + batch_size=[20], + ) + ) + sample = rb.sample() + assert sample.batch_size[0] == 8 + for key in ("sequence_id", "step_in_sequence", "validity_mask"): + assert key in sample.keys() + valid = sample["validity_mask"] + obs = sample["obs"] + step = sample["step_in_sequence"].float() + starts = obs - step + assert (starts[0:4] == starts[0]).all() + assert (starts[4:8] == starts[4]).all() + assert valid.all() + + def test_prioritized_weights_expanded_to_records(self): + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer( + storage=LazyTensorStorage(20), + sampler=PrioritizedSampler(20, alpha=0.7, beta=0.9), + batch_size=2, + sample_unit=Sequence(length=4), + ) + rb.extend( + TensorDict( + { + "obs": torch.arange(20, dtype=torch.float32), + ("next", "done"): torch.zeros(20, 1, dtype=torch.bool), + }, + batch_size=[20], + ) + ) + sample, info = rb.sample(return_info=True) + assert sample.batch_size[0] == 8 + for value in info.values(): + assert torch.as_tensor(value).reshape(-1).shape[0] in (8,) + + def test_invalid_length_raises(self): + Sequence = self._sequence_cls() + with pytest.raises(ValueError): + Sequence(length=0) + with pytest.raises(ValueError): + Sequence(length=4, episode_boundary="teleport") + + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown) From 5f81acc2ca1a6bbcdee172efff0e55df3f42949b Mon Sep 17 00:00:00 2001 From: coder-jayp Date: Sat, 25 Jul 2026 18:32:05 +0530 Subject: [PATCH 2/3] feat: implement Sequence sample unit with exact boundary policies (#4039) --- torchrl/data/replay_buffers/sample_units.py | 109 +++++++++++++++++++- 1 file changed, 108 insertions(+), 1 deletion(-) diff --git a/torchrl/data/replay_buffers/sample_units.py b/torchrl/data/replay_buffers/sample_units.py index 7e4fa101f58..55cd3a3d3e8 100644 --- a/torchrl/data/replay_buffers/sample_units.py +++ b/torchrl/data/replay_buffers/sample_units.py @@ -12,7 +12,7 @@ if TYPE_CHECKING: from torchrl.data.replay_buffers.storages import Storage -__all__ = ["SampleUnit", "Transition"] +__all__ = ["SampleUnit", "Transition", "Sequence"] class SampleUnit(abc.ABC): @@ -99,3 +99,110 @@ def expand( storage: Storage, ) -> tuple[torch.Tensor | tuple, dict[str, Any]]: return index, info + +class Sequence(SampleUnit): + """Expands anchors into a fixed-length sequence of records. + + Args: + length (int): the length of the sequences. + episode_boundary (str, optional): boundary policy. One of: + - `"pad"`: repeat the last valid state if a boundary is reached, marking + padded steps as invalid. + - `"stop"`: shift the anchor backward so the sequence ends exactly + at the boundary, falling back to pad if the episode is shorter + than `length`. + - `"include_reset"`: cross boundaries blindly. + Defaults to `"pad"`. + done_key (str or tuple, optional): the key for the end-of-episode flag. + Defaults to `("next", "done")`. + """ + + def __init__( + self, + length: int, + episode_boundary: str = "pad", + done_key: str | tuple[str, ...] | None = ("next", "done"), + ): + if length <= 0: + raise ValueError(f"length must be strictly positive, got {length}.") + if episode_boundary not in ("pad", "stop", "include_reset"): + raise ValueError(f"Unknown episode_boundary {episode_boundary}") + self.length = length + self.episode_boundary = episode_boundary + self.done_key = done_key + + def expand( + self, + index: torch.Tensor | tuple, + info: dict[str, Any], + storage: Storage, + ) -> tuple[torch.Tensor | tuple, dict[str, Any]]: + if isinstance(index, tuple): + raise NotImplementedError("Multidimensional storage not yet supported by Sequence.") + + anchor = index.clone() + B = anchor.shape[0] + device = anchor.device + + expanded_info = {} + for k, v in info.items(): + val = torch.as_tensor(v) + expanded_info[k] = val.repeat_interleave(self.length, dim=0) + + seq_id = torch.arange(B, device=device).repeat_interleave(self.length) + step_idx = torch.arange(self.length, device=device).repeat(B) + + expanded_info["sequence_id"] = seq_id + expanded_info["step_in_sequence"] = step_idx + + offset = torch.arange(self.length, device=device, dtype=torch.long).unsqueeze(0).expand(B, self.length) + validity = torch.ones((B, self.length), device=device, dtype=torch.bool) + + if self.episode_boundary in ("pad", "stop"): + from torchrl.data.replay_buffers.utils import _derive_end_flags, _end_to_start_stop + done = storage.get(self.done_key) if self.done_key is not None else None + end, max_len = _derive_end_flags( + end=done, + at_capacity=storage._is_full, + cursor=storage._last_cursor, + ) + start, stop, _ = _end_to_start_stop(end=end, length=max_len, device=device) + start = start[:, 0] + stop = stop[:, 0] + + start_exp = start.unsqueeze(0) + stop_exp = stop.unsqueeze(0) + a_exp = anchor.unsqueeze(1) + + cond1 = (start_exp <= stop_exp) & (start_exp <= a_exp) & (a_exp <= stop_exp) + cond2 = (start_exp > stop_exp) & ((a_exp >= start_exp) | (a_exp <= stop_exp)) + mask = cond1 | cond2 + + traj_idx = mask.float().argmax(dim=1) + a_start = start[traj_idx] + a_stop = stop[traj_idx] + + dist_to_stop = ((a_stop - anchor) % max_len).to(torch.long) + dist_from_start = ((anchor - a_start) % max_len).to(torch.long) + + if self.episode_boundary == "pad": + clamped_offset = torch.min(offset, dist_to_stop.unsqueeze(1)) + indices = anchor.unsqueeze(1) + clamped_offset + validity = offset <= dist_to_stop.unsqueeze(1) + elif self.episode_boundary == "stop": + shortfall = (self.length - 1) - dist_to_stop + shift = torch.clamp(shortfall, min=torch.zeros_like(shortfall), max=dist_from_start) + new_anchor = anchor - shift + new_dist_to_stop = dist_to_stop + shift + + clamped_offset = torch.min(offset, new_dist_to_stop.unsqueeze(1)) + indices = new_anchor.unsqueeze(1) + clamped_offset + validity = offset <= new_dist_to_stop.unsqueeze(1) + else: + indices = anchor.unsqueeze(1) + offset + max_len = storage.max_size + + indices = indices % max_len + expanded_info["validity_mask"] = validity.flatten() + + return indices.flatten(), expanded_info From d15fe7c4027f10273f379a85400fe6bbea8f620f Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Wed, 5 Aug 2026 17:18:55 +0100 Subject: [PATCH 3/3] [BugFix] Harden Sequence sample unit: partial-fill and write-cursor safety, device handling, storage validation, public API - include_reset now computes against the written length instead of max_size: anchors near the write head of a partially filled buffer no longer produce out-of-range indices, and on a full ring buffer the window clamps at the write cursor instead of splicing the newest data with the oldest under an all-True validity mask. - Keep all index bookkeeping on the sampler's index device and move the end flags there before _end_to_start_stop, so non-CPU storages (CUDA/MPS) no longer trip cross-device comparisons; returned indices live on the same device Transition returns. - Raise an informative TypeError when the storage is not a TensorDict-backed TensorStorage (ListStorage, plain-tensor storages), and guard the _last_cursor/_is_full attribute accesses. - Leave 0-dim info entries untouched instead of crashing on repeat_interleave. - Export Sequence from torchrl.data / torchrl.data.replay_buffers, add it to the docs autosummary, and use the public import path in tests. - Move the utils imports to module top; Literal/NestedKey type hints; runnable Examples block; normalize sequence-form done_key to tuple. - Add TransitionConfig and SequenceConfig Hydra companions with registration and cross-references (config/class parity). - Tests: partial-fill and full-ring include_reset, ListStorage/plain tensor errors, scalar info entries, custom nested done_key, non-CPU storage device, sample-unit config instantiation. Co-Authored-By: Claude Fable 5 --- docs/source/reference/data_replaybuffers.rst | 7 +- test/rb/test_rb_core.py | 131 ++++++++++++- test/test_configs.py | 22 +++ torchrl/data/__init__.py | 2 + torchrl/data/replay_buffers/__init__.py | 3 +- torchrl/data/replay_buffers/sample_units.py | 177 ++++++++++++++---- .../trainers/algorithms/configs/__init__.py | 9 + torchrl/trainers/algorithms/configs/data.py | 38 ++++ 8 files changed, 349 insertions(+), 40 deletions(-) diff --git a/docs/source/reference/data_replaybuffers.rst b/docs/source/reference/data_replaybuffers.rst index d0068dcfefa..9aa531c68e9 100644 --- a/docs/source/reference/data_replaybuffers.rst +++ b/docs/source/reference/data_replaybuffers.rst @@ -57,8 +57,10 @@ Replay sampling combines two orthogonal decisions: which anchors are selected A :class:`~torchrl.data.replay_buffers.SampleUnit` passed through the ``sample_unit`` argument owns the second decision. The default behavior, equivalent to :class:`~torchrl.data.replay_buffers.Transition`, keeps every -anchor as a single transition; future units expand anchors into fixed-length -sequences or complete trajectories with explicit boundary policies. +anchor as a single transition; +:class:`~torchrl.data.replay_buffers.Sequence` expands each anchor into a +fixed-length sequence of records with explicit episode-boundary policies +(``"pad"``, ``"stop"`` or ``"include_reset"``). .. code-block:: python @@ -76,6 +78,7 @@ sequences or complete trajectories with explicit boundary policies. :template: rl_template.rst SampleUnit + Sequence Transition Offline-to-online helpers diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py index 83027b0117b..899e2526e5f 100644 --- a/test/rb/test_rb_core.py +++ b/test/rb/test_rb_core.py @@ -19,6 +19,7 @@ from torchrl.data import ( PrioritizedReplayBuffer, ReplayBuffer, + Sequence, TensorDictPrioritizedReplayBuffer, TensorDictReplayBuffer, ) @@ -1322,9 +1323,8 @@ class TestSequenceUnit: """ def _sequence_cls(self): - from torchrl.data.replay_buffers import sample_units - - return sample_units.Sequence + # Sequence is part of the public API: use the public import path. + return Sequence def _make_storage(self, capacity=10, done_at=(5, 9)): rb = TensorDictReplayBuffer(storage=LazyTensorStorage(capacity), batch_size=4) @@ -1478,6 +1478,131 @@ def test_invalid_length_raises(self): with pytest.raises(ValueError): Sequence(length=4, episode_boundary="teleport") + def test_include_reset_partially_filled_buffer(self): + # Anchors close to the write head of a partially filled buffer must + # not produce indices past the written region (repro: capacity 100, + # 10 written, anchor 8, length 4 used to raise an IndexError). + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer(storage=LazyTensorStorage(100), batch_size=2) + rb.extend( + TensorDict( + { + "obs": torch.arange(10, dtype=torch.float32), + ("next", "done"): torch.zeros(10, 1, dtype=torch.bool), + }, + batch_size=[10], + ) + ) + unit = Sequence(length=4, episode_boundary="include_reset") + index, info = self._expand(rb, unit, [8]) + assert index.tolist() == [8, 9, 9, 9] + assert info["validity_mask"].tolist() == [True, True, False, False] + # reading the storage with the produced indices must not raise + rb._storage.get(index) + + def test_include_reset_does_not_splice_across_write_cursor(self): + # On a full ring buffer the record after the newest one is the oldest + # record: include_reset must clamp at the write cursor instead of + # splicing new and old data with an all-True validity mask. + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10), batch_size=2) + size = 14 + rb.extend( + TensorDict( + { + "obs": torch.arange(size, dtype=torch.float32), + ("next", "done"): torch.zeros(size, 1, dtype=torch.bool), + }, + batch_size=[size], + ) + ) + # physical slots 0..3 hold obs 10..13 (newest at slot 3), + # slots 4..9 hold obs 4..9 (oldest at slot 4) + unit = Sequence(length=4, episode_boundary="include_reset") + index, info = self._expand(rb, unit, [2]) + assert index.tolist() == [2, 3, 3, 3] + assert info["validity_mask"].tolist() == [True, True, False, False] + obs = rb[:]["obs"][index] + assert obs.tolist() == [12.0, 13.0, 13.0, 13.0] + + def test_requires_tensordict_storage(self): + Sequence = self._sequence_cls() + unit = Sequence(length=3) + rb = ReplayBuffer(storage=ListStorage(10), batch_size=2) + with pytest.raises(TypeError, match="TensorStorage"): + unit.expand(torch.tensor([0, 1]), {}, rb._storage) + # plain-tensor TensorStorage is rejected as well + rb = ReplayBuffer(storage=LazyTensorStorage(10), batch_size=2) + rb.extend(torch.arange(5)) + with pytest.raises(TypeError, match="TensorDict"): + unit.expand(torch.tensor([0, 1]), {}, rb._storage) + + def test_scalar_info_entries_pass_through(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + unit = Sequence(length=3) + _, info = unit.expand( + torch.tensor([0, 6]), + {"scalar_meta": 3.0, "per_anchor": torch.tensor([1.0, 2.0])}, + rb._storage, + ) + assert info["scalar_meta"] == 3.0 + assert info["per_anchor"].tolist() == [1.0, 1.0, 1.0, 2.0, 2.0, 2.0] + + def test_custom_nested_done_key(self): + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer(storage=LazyTensorStorage(10), batch_size=2) + done = torch.zeros(10, 1, dtype=torch.bool) + done[5] = done[9] = True + rb.extend( + TensorDict( + { + "obs": torch.arange(10, dtype=torch.float32), + ("stats", "episode_end"): done, + }, + batch_size=[10], + ) + ) + unit = Sequence(length=4, done_key=("stats", "episode_end")) + index, info = self._expand(rb, unit, [4]) + assert index.tolist() == [4, 5, 5, 5] + assert info["validity_mask"].tolist() == [True, True, False, False] + + @pytest.mark.gpu + @pytest.mark.skipif( + not torch.cuda.is_available() and not torch.backends.mps.is_available(), + reason="needs a non-CPU device (CUDA or MPS)", + ) + def test_non_cpu_storage(self): + # Storage on an accelerator with anchors on CPU (as samplers produce + # them): expansion must not mix devices, and the returned indices + # live on the anchor device like Transition's. + Sequence = self._sequence_cls() + device = "cuda" if torch.cuda.is_available() else "mps" + done = torch.zeros(10, 1, dtype=torch.bool) + done[5] = done[9] = True + data = TensorDict( + { + "obs": torch.arange(10, dtype=torch.float32), + ("next", "done"): done, + }, + batch_size=[10], + ).to(device) + rb = TensorDictReplayBuffer( + storage=LazyTensorStorage(10, device=device), + batch_size=2, + sample_unit=Sequence(length=3), + ) + rb.extend(data) + anchors = torch.tensor([4], dtype=torch.long) + for boundary in ("pad", "stop", "include_reset"): + unit = Sequence(length=4, episode_boundary=boundary) + index, info = unit.expand(anchors, {}, rb._storage) + assert index.device == anchors.device + assert info["validity_mask"].device == anchors.device + sample = rb.sample() + assert sample["obs"].shape[0] == 6 + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() diff --git a/test/test_configs.py b/test/test_configs.py index 3d3023766d2..5189ef16e07 100644 --- a/test/test_configs.py +++ b/test/test_configs.py @@ -470,6 +470,28 @@ def test_random_sampler_config(self): sampler = instantiate(cfg) assert isinstance(sampler, RandomSampler) + @pytest.mark.skipif(not _has_hydra, reason="Hydra is not installed") + def test_sample_unit_configs(self): + """Test TransitionConfig and SequenceConfig.""" + from hydra.utils import instantiate + from torchrl.data.replay_buffers import Sequence, Transition + from torchrl.trainers.algorithms.configs.data import ( + SequenceConfig, + TransitionConfig, + ) + + cfg = TransitionConfig() + assert cfg._target_ == "torchrl.data.replay_buffers.Transition" + assert isinstance(instantiate(cfg), Transition) + + cfg = SequenceConfig(length=8, episode_boundary="stop") + assert cfg._target_ == "torchrl.data.replay_buffers.Sequence" + unit = instantiate(cfg) + assert isinstance(unit, Sequence) + assert unit.length == 8 + assert unit.episode_boundary == "stop" + assert unit.done_key == ("next", "done") + @pytest.mark.skipif(not _has_hydra, reason="Hydra is not installed") def test_tensor_storage_config(self): """Test TensorStorageConfig.""" diff --git a/torchrl/data/__init__.py b/torchrl/data/__init__.py index f9b4d40510d..7d06c3a9fdd 100644 --- a/torchrl/data/__init__.py +++ b/torchrl/data/__init__.py @@ -69,6 +69,7 @@ SamplerEnsemble, SamplerWithoutReplacement, SampleUnit, + Sequence, SliceSampler, SliceSamplerWithoutReplacement, StalenessAwareSampler, @@ -186,6 +187,7 @@ "RoundRobinWriter", "SampleUnit", "SamplerEnsemble", + "Sequence", "SamplerWithoutReplacement", "SipHash", "SliceSampler", diff --git a/torchrl/data/replay_buffers/__init__.py b/torchrl/data/replay_buffers/__init__.py index 20d9702dcba..0a115dcc3ee 100644 --- a/torchrl/data/replay_buffers/__init__.py +++ b/torchrl/data/replay_buffers/__init__.py @@ -31,7 +31,7 @@ TensorDictPrioritizedReplayBuffer, TensorDictReplayBuffer, ) -from .sample_units import SampleUnit, Transition +from .sample_units import SampleUnit, Sequence, Transition from .samplers import ( ConsumingSampler, PrioritizedSampler, @@ -97,6 +97,7 @@ "RemoteTensorDictReplayBuffer", "ReplayBuffer", "SampleUnit", + "Sequence", "Transition", "ReplayBufferEnsemble", "TensorDictPrioritizedReplayBuffer", diff --git a/torchrl/data/replay_buffers/sample_units.py b/torchrl/data/replay_buffers/sample_units.py index 55cd3a3d3e8..68d9f8e7387 100644 --- a/torchrl/data/replay_buffers/sample_units.py +++ b/torchrl/data/replay_buffers/sample_units.py @@ -5,9 +5,14 @@ from __future__ import annotations import abc -from typing import Any, TYPE_CHECKING +from typing import Any, Literal, TYPE_CHECKING import torch +from tensordict import is_tensor_collection +from tensordict.utils import NestedKey + +from torchrl.data.replay_buffers.storages import TensorStorage +from torchrl.data.replay_buffers.utils import _derive_end_flags, _end_to_start_stop if TYPE_CHECKING: from torchrl.data.replay_buffers.storages import Storage @@ -76,6 +81,9 @@ class Transition(SampleUnit): anchors selected by the sampler are the records of the batch, and the info dictionary is returned untouched. + .. seealso:: :class:`~torchrl.trainers.algorithms.configs.data.TransitionConfig` + for the Hydra configuration companion. + Examples: >>> import torch >>> from torchrl.data import LazyTensorStorage, ReplayBuffer @@ -100,37 +108,113 @@ def expand( ) -> tuple[torch.Tensor | tuple, dict[str, Any]]: return index, info + class Sequence(SampleUnit): """Expands anchors into a fixed-length sequence of records. + This unit requires a :class:`~torchrl.data.replay_buffers.TensorStorage` + backed by a TensorDict (e.g. :class:`~torchrl.data.LazyTensorStorage` + filled with TensorDict data), since episode boundaries are read from the + stored ``done_key`` entry. + Args: length (int): the length of the sequences. episode_boundary (str, optional): boundary policy. One of: - - `"pad"`: repeat the last valid state if a boundary is reached, marking - padded steps as invalid. - - `"stop"`: shift the anchor backward so the sequence ends exactly - at the boundary, falling back to pad if the episode is shorter - than `length`. - - `"include_reset"`: cross boundaries blindly. - Defaults to `"pad"`. - done_key (str or tuple, optional): the key for the end-of-episode flag. - Defaults to `("next", "done")`. + + - ``"pad"``: repeat the last valid state if a boundary is reached, + marking padded steps as invalid in the ``"validity_mask"`` info + entry. + - ``"stop"``: shift the anchor backward so the sequence ends + exactly at the boundary, falling back to pad if the episode is + shorter than ``length``. + - ``"include_reset"``: cross episode boundaries. The write seam + (the boundary between the newest and the oldest record of the + ring buffer) and unwritten slots are never crossed: records + beyond it are clamped and marked invalid. + + Defaults to ``"pad"``. + done_key (NestedKey, optional): the key for the end-of-episode flag. + Defaults to ``("next", "done")``. + + .. seealso:: :class:`~torchrl.trainers.algorithms.configs.data.SequenceConfig` + for the Hydra configuration companion. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from torchrl.data import LazyTensorStorage, ReplayBuffer, Sequence + >>> rb = ReplayBuffer( + ... storage=LazyTensorStorage(10), + ... batch_size=2, + ... sample_unit=Sequence(length=3), + ... ) + >>> done = torch.zeros(10, 1, dtype=torch.bool) + >>> done[4] = done[9] = True + >>> rb.extend(TensorDict( + ... { + ... "obs": torch.arange(10, dtype=torch.float32), + ... ("next", "done"): done, + ... }, + ... batch_size=[10], + ... )) + tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> sample, info = rb.sample(return_info=True) + >>> sample["obs"].shape # 2 anchors x 3 records each + torch.Size([6]) + >>> sorted(info.keys()) + ['index', 'sequence_id', 'step_in_sequence', 'validity_mask'] """ def __init__( self, length: int, - episode_boundary: str = "pad", - done_key: str | tuple[str, ...] | None = ("next", "done"), + episode_boundary: Literal["pad", "stop", "include_reset"] = "pad", + done_key: NestedKey | None = ("next", "done"), ): if length <= 0: raise ValueError(f"length must be strictly positive, got {length}.") if episode_boundary not in ("pad", "stop", "include_reset"): raise ValueError(f"Unknown episode_boundary {episode_boundary}") + if done_key is not None and not isinstance(done_key, str): + # normalize sequence-form nested keys (e.g. lists or omegaconf + # containers coming from Hydra configs) to plain tuples + done_key = tuple(done_key) self.length = length self.episode_boundary = episode_boundary self.done_key = done_key + @staticmethod + def _newest_index(storage: Storage, written: int) -> int: + """Physical index of the most recently written record.""" + cursor = getattr(storage, "_last_cursor", None) + if isinstance(cursor, torch.Tensor): + cursor = cursor.reshape(-1) + if cursor.numel(): + return int(cursor[-1].item()) % written + elif isinstance(cursor, range): + if len(cursor): + return int(cursor[-1]) % written + elif isinstance(cursor, int): + return cursor % written + return written - 1 + + def _check_storage(self, storage: Storage) -> None: + if not isinstance(storage, TensorStorage): + raise TypeError( + f"{type(self).__name__} requires a TensorDict-backed TensorStorage " + f"(e.g. LazyTensorStorage or LazyMemmapStorage written with " + f"TensorDict data) to recover episode boundaries and the write " + f"cursor; got {type(storage).__name__}." + ) + contents = getattr(storage, "_storage", None) + if contents is not None and not is_tensor_collection(contents): + raise TypeError( + f"{type(self).__name__} requires the TensorStorage to hold a " + f"TensorDict (or other tensor collection) so that the " + f"'{self.done_key}' entry can be read; the storage holds " + f"{type(contents).__name__} instead." + ) + def expand( self, index: torch.Tensor | tuple, @@ -138,16 +222,26 @@ def expand( storage: Storage, ) -> tuple[torch.Tensor | tuple, dict[str, Any]]: if isinstance(index, tuple): - raise NotImplementedError("Multidimensional storage not yet supported by Sequence.") + raise NotImplementedError( + "Multidimensional storage not yet supported by Sequence." + ) + self._check_storage(storage) anchor = index.clone() B = anchor.shape[0] + # All bookkeeping happens on the sampler's index device so that the + # returned indices live on the same device as the ones Transition + # (identity) would return. device = anchor.device expanded_info = {} for k, v in info.items(): val = torch.as_tensor(v) - expanded_info[k] = val.repeat_interleave(self.length, dim=0) + if val.ndim == 0: + # scalar metadata is not per-anchor: leave it untouched + expanded_info[k] = v + else: + expanded_info[k] = val.repeat_interleave(self.length, dim=0) seq_id = torch.arange(B, device=device).repeat_interleave(self.length) step_idx = torch.arange(self.length, device=device).repeat(B) @@ -155,17 +249,23 @@ def expand( expanded_info["sequence_id"] = seq_id expanded_info["step_in_sequence"] = step_idx - offset = torch.arange(self.length, device=device, dtype=torch.long).unsqueeze(0).expand(B, self.length) - validity = torch.ones((B, self.length), device=device, dtype=torch.bool) + offset = ( + torch.arange(self.length, device=device, dtype=torch.long) + .unsqueeze(0) + .expand(B, self.length) + ) if self.episode_boundary in ("pad", "stop"): - from torchrl.data.replay_buffers.utils import _derive_end_flags, _end_to_start_stop done = storage.get(self.done_key) if self.done_key is not None else None end, max_len = _derive_end_flags( end=done, at_capacity=storage._is_full, - cursor=storage._last_cursor, + cursor=getattr(storage, "_last_cursor", None), ) + # _end_to_start_stop returns its indices on the device of ``end`` + # (the storage device): move the flags first so start/stop live + # on the anchor device. + end = end.to(device) start, stop, _ = _end_to_start_stop(end=end, length=max_len, device=device) start = start[:, 0] stop = stop[:, 0] @@ -175,7 +275,9 @@ def expand( a_exp = anchor.unsqueeze(1) cond1 = (start_exp <= stop_exp) & (start_exp <= a_exp) & (a_exp <= stop_exp) - cond2 = (start_exp > stop_exp) & ((a_exp >= start_exp) | (a_exp <= stop_exp)) + cond2 = (start_exp > stop_exp) & ( + (a_exp >= start_exp) | (a_exp <= stop_exp) + ) mask = cond1 | cond2 traj_idx = mask.float().argmax(dim=1) @@ -185,24 +287,31 @@ def expand( dist_to_stop = ((a_stop - anchor) % max_len).to(torch.long) dist_from_start = ((anchor - a_start) % max_len).to(torch.long) - if self.episode_boundary == "pad": - clamped_offset = torch.min(offset, dist_to_stop.unsqueeze(1)) - indices = anchor.unsqueeze(1) + clamped_offset - validity = offset <= dist_to_stop.unsqueeze(1) - elif self.episode_boundary == "stop": + if self.episode_boundary == "stop": shortfall = (self.length - 1) - dist_to_stop - shift = torch.clamp(shortfall, min=torch.zeros_like(shortfall), max=dist_from_start) - new_anchor = anchor - shift - new_dist_to_stop = dist_to_stop + shift - - clamped_offset = torch.min(offset, new_dist_to_stop.unsqueeze(1)) - indices = new_anchor.unsqueeze(1) + clamped_offset - validity = offset <= new_dist_to_stop.unsqueeze(1) + shift = torch.clamp( + shortfall, min=torch.zeros_like(shortfall), max=dist_from_start + ) + anchor = anchor - shift + dist_to_stop = dist_to_stop + shift + + clamped_offset = torch.minimum(offset, dist_to_stop.unsqueeze(1)) + validity = offset <= dist_to_stop.unsqueeze(1) + indices = (anchor.unsqueeze(1) + clamped_offset) % max_len else: - indices = anchor.unsqueeze(1) + offset - max_len = storage.max_size + # "include_reset": cross episode boundaries, but never cross the + # write seam (between the newest and the oldest record of the + # ring buffer) nor read slots that were never written. + written = len(storage) + newest = self._newest_index(storage, written) + dist_to_newest = torch.remainder( + torch.as_tensor(newest, device=device, dtype=torch.long) - anchor, + written, + ) + clamped_offset = torch.minimum(offset, dist_to_newest.unsqueeze(1)) + validity = offset <= dist_to_newest.unsqueeze(1) + indices = torch.remainder(anchor.unsqueeze(1) + clamped_offset, written) - indices = indices % max_len expanded_info["validity_mask"] = validity.flatten() return indices.flatten(), expanded_info diff --git a/torchrl/trainers/algorithms/configs/__init__.py b/torchrl/trainers/algorithms/configs/__init__.py index 897ea52ebfa..86e5743a4cc 100644 --- a/torchrl/trainers/algorithms/configs/__init__.py +++ b/torchrl/trainers/algorithms/configs/__init__.py @@ -41,12 +41,15 @@ ReplayBufferConfig, RoundRobinWriterConfig, SamplerWithoutReplacementConfig, + SampleUnitConfig, + SequenceConfig, SliceSamplerConfig, SliceSamplerWithoutReplacementConfig, StorageEnsembleConfig, StorageEnsembleWriterConfig, TensorDictReplayBufferConfig, TensorStorageConfig, + TransitionConfig, ) from torchrl.trainers.algorithms.configs.envs import ( BatchedEnvConfig, @@ -383,6 +386,10 @@ "StorageEnsembleWriterConfig", "TensorDictReplayBufferConfig", "TensorStorageConfig", + # Sample units + "SampleUnitConfig", + "SequenceConfig", + "TransitionConfig", # Samplers "ConsumingSamplerConfig", "PrioritizedSamplerConfig", @@ -672,6 +679,8 @@ def _register_configs(): cs.store(group="storage", name="lazy_tensor", node=LazyTensorStorageConfig) cs.store(group="storage", name="lazy_memmap", node=LazyMemmapStorageConfig) cs.store(group="writer", name="round_robin", node=RoundRobinWriterConfig) + cs.store(group="sample_unit", name="transition", node=TransitionConfig) + cs.store(group="sample_unit", name="sequence", node=SequenceConfig) # ============================================================================= # Collector Configurations diff --git a/torchrl/trainers/algorithms/configs/data.py b/torchrl/trainers/algorithms/configs/data.py index 8fc153a58f7..7c3d2e7a3d6 100644 --- a/torchrl/trainers/algorithms/configs/data.py +++ b/torchrl/trainers/algorithms/configs/data.py @@ -202,6 +202,44 @@ class SamplerWithoutReplacementConfig(SamplerConfig): shuffle: bool = True +@dataclass +class SampleUnitConfig(ConfigBase): + """Base configuration class for replay buffer sample units. + + See also :class:`~torchrl.data.replay_buffers.SampleUnit`. + """ + + _target_: str = "torchrl.data.replay_buffers.SampleUnit" + + def __post_init__(self) -> None: + """Post-initialization hook for sample unit configurations.""" + + +@dataclass +class TransitionConfig(SampleUnitConfig): + """Hydra configuration for :class:`~torchrl.data.replay_buffers.Transition`. + + ``Transition.__init__`` takes no arguments, so this config only carries + the instantiation target. + """ + + _target_: str = "torchrl.data.replay_buffers.Transition" + + +@dataclass +class SequenceConfig(SampleUnitConfig): + """Hydra configuration for :class:`~torchrl.data.replay_buffers.Sequence`. + + Every kwarg accepted by ``Sequence.__init__`` is exposed as a field here + with the same default. + """ + + _target_: str = "torchrl.data.replay_buffers.Sequence" + length: int = MISSING + episode_boundary: str = "pad" + done_key: Any = ("next", "done") + + @dataclass class StorageConfig(ConfigBase): """Base configuration class for replay buffer storage."""