From 6795078b8912274535f52830f5ecd65605557ed5 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Fri, 24 Jul 2026 01:59:52 -0700 Subject: [PATCH 1/7] [Feature] Add SampleUnit composition point for replay sampling Introduces the interface layer of the sampler-decomposition RFC: sampling combines an anchor distribution (the sampler) with a range expansion (the sample unit), and this PR adds the composition point without changing any behavior. A SampleUnit receives the anchor index, the sampler info dict and the storage inside the buffer's sampling critical section, after the anchor sampler ran and before the storage read or index bookkeeping, and returns the expanded index plus (possibly augmented) info. Unit metadata added to info flows into sample(return_info=True) and becomes keys of TensorDict samples through the existing info-copy path. Transition is the identity unit and the implicit default: sample_unit=None and sample_unit=Transition() are behaviorally identical, which the tests pin under seeded generators. The sample_unit keyword is exposed on ReplayBuffer and threaded through the prioritized variants; Hydra config companions gain the matching field. This settles the RFC's first open question (units are buffer-owned) and gives the Sequence unit, boundary policies and priority-semantics follow-ups a stable target. Part of #4039. --- docs/source/reference/data_replaybuffers.rst | 29 +++++ test/rb/test_rb_core.py | 104 ++++++++++++++++++ torchrl/data/__init__.py | 4 + torchrl/data/replay_buffers/__init__.py | 3 + torchrl/data/replay_buffers/replay_buffers.py | 25 +++++ torchrl/data/replay_buffers/sample_units.py | 101 +++++++++++++++++ torchrl/trainers/algorithms/configs/data.py | 2 + 7 files changed, 268 insertions(+) create mode 100644 torchrl/data/replay_buffers/sample_units.py diff --git a/docs/source/reference/data_replaybuffers.rst b/docs/source/reference/data_replaybuffers.rst index 8ca46081ad5..d0068dcfefa 100644 --- a/docs/source/reference/data_replaybuffers.rst +++ b/docs/source/reference/data_replaybuffers.rst @@ -49,6 +49,35 @@ discovery and buffer lifecycle. RemoteTensorDictReplayBuffer +Sample units +------------ + +Replay sampling combines two orthogonal decisions: which anchors are selected +(the sampler's probability distribution) and what each anchor expands into. +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. + +.. code-block:: python + + from torchrl.data import LazyTensorStorage, ReplayBuffer + from torchrl.data.replay_buffers import Transition + + rb = ReplayBuffer( + storage=LazyTensorStorage(1000), + batch_size=32, + sample_unit=Transition(), + ) + +.. autosummary:: + :toctree: generated/ + :template: rl_template.rst + + SampleUnit + Transition + Offline-to-online helpers ------------------------- diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py index a5763e9aa1d..3786a93a83b 100644 --- a/test/rb/test_rb_core.py +++ b/test/rb/test_rb_core.py @@ -22,6 +22,7 @@ TensorDictPrioritizedReplayBuffer, TensorDictReplayBuffer, ) +from torchrl.data.replay_buffers.sample_units import SampleUnit, Transition from torchrl.data.replay_buffers.samplers import ( ConsumingSampler, PrioritizedSampler, @@ -1196,6 +1197,109 @@ def test_stats_with_non_counting_writer(self): assert stats["capacity"] == 10 +class _RepeatTwiceUnit(SampleUnit): + """Toy unit doubling every anchor and recording per-record provenance.""" + + def expand(self, index, info, storage): + index = torch.as_tensor(index).repeat_interleave(2) + info = dict(info) + info["unit_repeat"] = torch.arange(index.numel()) % 2 + return index, info + + +class TestSampleUnit: + """Executable spec for the SampleUnit composition point (#4039, PR 1). + + Contract pinned by this class: + + - ``sample_unit=None`` (default) and ``sample_unit=Transition()`` are + behaviorally identical: same sampled data under the same generator + state, same info entries. + - A unit's ``expand`` runs after the anchor sampler and before storage + read and index bookkeeping, so ``info["index"]`` reports the expanded + indices and the returned batch is built from them. + - Metadata a unit adds to ``info`` flows into ``sample(return_info=True)`` + and becomes keys of TensorDict samples. + - ``sample_unit`` must be a ``SampleUnit`` instance; anything else raises + ``TypeError`` at construction. + """ + + def test_default_and_transition_are_identical(self): + data = torch.arange(20) + samples = {} + for name, unit in (("default", None), ("transition", Transition())): + generator = torch.Generator() + generator.manual_seed(0) + rb = ReplayBuffer( + storage=LazyTensorStorage(20), + batch_size=4, + generator=generator, + sample_unit=unit, + ) + rb.extend(data) + samples[name] = rb.sample(return_info=True) + default_sample, default_info = samples["default"] + transition_sample, transition_info = samples["transition"] + torch.testing.assert_close(default_sample, transition_sample) + assert set(default_info) == set(transition_info) + torch.testing.assert_close( + torch.as_tensor(default_info["index"]), + torch.as_tensor(transition_info["index"]), + ) + + def test_transition_adds_no_info_keys(self): + rb = ReplayBuffer( + storage=LazyTensorStorage(10), batch_size=4, sample_unit=Transition() + ) + rb.extend(torch.arange(10)) + _, info = rb.sample(return_info=True) + assert "unit_repeat" not in info + + def test_custom_unit_expands_batch_and_metadata(self): + rb = ReplayBuffer( + storage=LazyTensorStorage(10), + batch_size=4, + sample_unit=_RepeatTwiceUnit(), + ) + rb.extend(torch.arange(10, dtype=torch.float32)) + sample, info = rb.sample(return_info=True) + assert sample.shape[0] == 8 + torch.testing.assert_close(sample[0::2], sample[1::2]) + index = torch.as_tensor(info["index"]) + assert index.numel() == 8 + assert (index[0::2] == index[1::2]).all() + assert info["unit_repeat"].numel() == 8 + + def test_tensordict_buffer_carries_unit_metadata(self): + rb = TensorDictReplayBuffer( + storage=LazyTensorStorage(10), + batch_size=4, + sample_unit=_RepeatTwiceUnit(), + ) + rb.extend(TensorDict({"obs": torch.randn(10, 3)}, batch_size=[10])) + sample = rb.sample() + assert sample.batch_size[0] == 8 + assert "unit_repeat" in sample.keys() + torch.testing.assert_close(sample["obs"][0::2], sample["obs"][1::2]) + + def test_invalid_sample_unit_raises(self): + with pytest.raises(TypeError, match="sample_unit"): + ReplayBuffer(storage=LazyTensorStorage(10), sample_unit=object()) + + def test_prioritized_buffer_with_transition_unit(self): + rb = TensorDictPrioritizedReplayBuffer( + alpha=0.7, + beta=0.9, + storage=LazyTensorStorage(10), + batch_size=4, + sample_unit=Transition(), + ) + rb.extend(TensorDict({"obs": torch.randn(10, 3)}, batch_size=[10])) + sample = rb.sample() + assert sample.batch_size[0] == 4 + rb.update_tensordict_priority(sample) + + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown) diff --git a/torchrl/data/__init__.py b/torchrl/data/__init__.py index 1fe24678119..f9b4d40510d 100644 --- a/torchrl/data/__init__.py +++ b/torchrl/data/__init__.py @@ -68,6 +68,7 @@ RoundRobinWriter, SamplerEnsemble, SamplerWithoutReplacement, + SampleUnit, SliceSampler, SliceSamplerWithoutReplacement, StalenessAwareSampler, @@ -87,6 +88,7 @@ traj, Trajectory, TrajectoryPredicate, + Transition, Writer, WriterEnsemble, ) @@ -182,6 +184,7 @@ "RobotDatasetMetadata", "RolloutFromModel", "RoundRobinWriter", + "SampleUnit", "SamplerEnsemble", "SamplerWithoutReplacement", "SipHash", @@ -200,6 +203,7 @@ "TensorDictMap", "TensorDictMaxValueWriter", "TensorDictPrioritizedReplayBuffer", + "Transition", "TensorDictReplayBuffer", "TensorDictRoundRobinWriter", "TensorDictTokenizer", diff --git a/torchrl/data/replay_buffers/__init__.py b/torchrl/data/replay_buffers/__init__.py index 1d258b8d78b..20d9702dcba 100644 --- a/torchrl/data/replay_buffers/__init__.py +++ b/torchrl/data/replay_buffers/__init__.py @@ -31,6 +31,7 @@ TensorDictPrioritizedReplayBuffer, TensorDictReplayBuffer, ) +from .sample_units import SampleUnit, Transition from .samplers import ( ConsumingSampler, PrioritizedSampler, @@ -95,6 +96,8 @@ "PrioritizedReplayBuffer", "RemoteTensorDictReplayBuffer", "ReplayBuffer", + "SampleUnit", + "Transition", "ReplayBufferEnsemble", "TensorDictPrioritizedReplayBuffer", "TensorDictReplayBuffer", diff --git a/torchrl/data/replay_buffers/replay_buffers.py b/torchrl/data/replay_buffers/replay_buffers.py index a5d781e71df..bacac4c02b0 100644 --- a/torchrl/data/replay_buffers/replay_buffers.py +++ b/torchrl/data/replay_buffers/replay_buffers.py @@ -58,6 +58,7 @@ def tree_leaves(data): # noqa: D103 rl_warnings, ) from torchrl.data.replay_buffers.query import _query_source, Trajectory +from torchrl.data.replay_buffers.sample_units import SampleUnit from torchrl.data.replay_buffers.samplers import ( ConsumingSampler, PrioritizedSampler, @@ -137,6 +138,12 @@ class ReplayBuffer(metaclass=_RayServiceMetaClass): If a callable is passed, it is used as constructor for the sampler. If none is provided, a default :class:`~torchrl.data.replay_buffers.RandomSampler` will be used. + sample_unit (SampleUnit, optional): expands the anchors selected by + the sampler into the records of the batch (see + :class:`~torchrl.data.replay_buffers.SampleUnit`). ``None`` + (default) is equivalent to + :class:`~torchrl.data.replay_buffers.Transition`: every anchor is + one transition and classic behavior is preserved. writer (Writer, Callable[[], Writer], optional): the writer to be used. If a callable is passed, it is used as constructor for the writer. If none is provided a default :class:`~torchrl.data.replay_buffers.RoundRobinWriter` @@ -339,6 +346,7 @@ def __init__( *, storage: Storage | Callable[[], Storage] | None = None, sampler: Sampler | Callable[[], Sampler] | None = None, + sample_unit: SampleUnit | None = None, writer: Writer | Callable[[], Writer] | None = None, collate_fn: Callable | None = None, pin_memory: bool = False, @@ -411,6 +419,11 @@ def __init__( # Update _delayed_init after auto-detection self._delayed_init = delayed_init + if sample_unit is not None and not isinstance(sample_unit, SampleUnit): + raise TypeError( + f"sample_unit must be a SampleUnit instance, got {type(sample_unit).__name__}." + ) + self._sample_unit = sample_unit self._pin_memory = pin_memory self._prefetch = bool(prefetch) self._prefetch_cap = prefetch or 0 @@ -1515,6 +1528,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]: nc = contextlib.nullcontext() 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) + if self._sample_unit is not None: + index, info = self._sample_unit.expand(index, info, self._storage) info["index"] = index data = self._storage.get(_storage_index(index, self._storage)) if not isinstance(index, INT_CLASSES): @@ -2040,6 +2055,7 @@ def __init__( dtype: torch.dtype = torch.float, storage: Storage | None = None, sampler: Sampler | None = None, + sample_unit: SampleUnit | None = None, sampler_device: DEVICE_TYPING | None = None, sync: bool = True, collate_fn: Callable | None = None, @@ -2078,6 +2094,7 @@ def __init__( super().__init__( storage=storage, sampler=sampler, + sample_unit=sample_unit, collate_fn=collate_fn, pin_memory=pin_memory, prefetch=prefetch, @@ -2134,6 +2151,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]: self._write_lock if not is_comp else nc, ): index, info = self.prioritized_sampler.sample(self._storage, batch_size) + if self._sample_unit is not None: + index, info = self._sample_unit.expand(index, info, self._storage) info["index"] = index data = self._storage.get(_storage_index(index, self._storage)) if not isinstance(index, INT_CLASSES): @@ -2563,6 +2582,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]: nc = contextlib.nullcontext() 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) + if self._sample_unit is not None: + index, info = self._sample_unit.expand(index, info, self._storage) info["index"] = index data = self._storage.get(_storage_index(index, self._storage)) if not isinstance(index, INT_CLASSES): @@ -2752,6 +2773,7 @@ def __init__( priority_key: NestedKey = "td_error", eps: float = 1e-8, storage: Storage | None = None, + sample_unit: SampleUnit | None = None, sampler_device: DEVICE_TYPING | None = None, sync: bool = True, collate_fn: Callable | None = None, @@ -2793,6 +2815,7 @@ def __init__( priority_key=priority_key, storage=storage, sampler=sampler, + sample_unit=sample_unit, collate_fn=collate_fn, pin_memory=pin_memory, prefetch=prefetch, @@ -2928,6 +2951,8 @@ def _sample(self, batch_size: int) -> tuple[Any, dict]: self._write_lock if not is_comp else nc, ): index, info = self.prioritized_sampler.sample(self._storage, batch_size) + if self._sample_unit is not None: + index, info = self._sample_unit.expand(index, info, self._storage) info["index"] = index data = self._storage.get(_storage_index(index, self._storage)) if not isinstance(index, INT_CLASSES): diff --git a/torchrl/data/replay_buffers/sample_units.py b/torchrl/data/replay_buffers/sample_units.py new file mode 100644 index 00000000000..7e4fa101f58 --- /dev/null +++ b/torchrl/data/replay_buffers/sample_units.py @@ -0,0 +1,101 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import abc +from typing import Any, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torchrl.data.replay_buffers.storages import Storage + +__all__ = ["SampleUnit", "Transition"] + + +class SampleUnit(abc.ABC): + """Expands sampled anchors into the records a batch is made of. + + Replay sampling combines two orthogonal decisions: which anchors are + selected (the sampler's probability distribution) and what each anchor + expands into (a single transition, a fixed-length sequence, a complete + trajectory). A ``SampleUnit`` owns the second decision. The buffer calls + :meth:`expand` inside its sampling critical section, after the anchor + sampler ran and before the storage is read or any index bookkeeping + happens, so the indices it returns are the ones the batch is built from + and the ones reported in the sample info. + + Contract for implementations: + + - ``expand`` receives the anchor index (a tensor, or a tuple of + coordinate tensors for multidimensional storages), the sampler's info + dictionary and the storage. It returns the expanded index and info, + which may be new objects; it must not mutate the storage. + - Entries of ``info`` that are aligned with the anchors (for example + priority weights) are the unit's responsibility: a unit that changes + the number of records must expand or reduce those entries so they stay + aligned with the index it returns. + - Metadata describing the expansion (validity masks, learning masks, + per-record anchor provenance) is communicated by adding entries to + ``info``; scalar-per-record tensors are surfaced as keys of + TensorDict samples automatically. + + .. seealso:: :class:`Transition`, the identity unit reproducing classic + one-anchor-one-transition sampling. + """ + + @abc.abstractmethod + def expand( + self, + index: torch.Tensor | tuple, + info: dict[str, Any], + storage: Storage, + ) -> tuple[torch.Tensor | tuple, dict[str, Any]]: + """Expands anchor indices into the final record indices of the batch. + + Args: + index (torch.Tensor or tuple of torch.Tensor): the anchor indices + selected by the sampler. + info (dict): the sampler's info dictionary. + storage (Storage): the storage the batch will be read from. + + Returns: + A tuple ``(index, info)`` with the expanded indices and the + (possibly augmented) info dictionary. + """ + ... + + +class Transition(SampleUnit): + """The identity sample unit: every anchor is one transition. + + This unit reproduces the classic replay-buffer behavior exactly and is + the implicit default when no ``sample_unit`` is passed to the buffer: + anchors selected by the sampler are the records of the batch, and the + info dictionary is returned untouched. + + Examples: + >>> import torch + >>> from torchrl.data import LazyTensorStorage, ReplayBuffer + >>> from torchrl.data.replay_buffers import Transition + >>> rb = ReplayBuffer( + ... storage=LazyTensorStorage(10), + ... batch_size=4, + ... sample_unit=Transition(), + ... ) + >>> rb.extend(torch.arange(10)) + tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> sample = rb.sample() + >>> sample.shape + torch.Size([4]) + """ + + def expand( + self, + index: torch.Tensor | tuple, + info: dict[str, Any], + storage: Storage, + ) -> tuple[torch.Tensor | tuple, dict[str, Any]]: + return index, info diff --git a/torchrl/trainers/algorithms/configs/data.py b/torchrl/trainers/algorithms/configs/data.py index 822287257c0..8fc153a58f7 100644 --- a/torchrl/trainers/algorithms/configs/data.py +++ b/torchrl/trainers/algorithms/configs/data.py @@ -326,6 +326,7 @@ class TensorDictReplayBufferConfig(ReplayBufferBaseConfig): _target_: str = "torchrl.data.replay_buffers.TensorDictReplayBuffer" priority_key: str = "td_error" sampler: Any = None + sample_unit: Any = None storage: Any = None writer: Any = None collate_fn: Any = None @@ -359,6 +360,7 @@ class ReplayBufferConfig(ReplayBufferBaseConfig): _target_: str = "torchrl.data.replay_buffers.ReplayBuffer" storage: Any = None sampler: Any = None + sample_unit: Any = None writer: Any = None collate_fn: Any = None pin_memory: bool = False From f2225a14492309a45bbadda9d4cf59a9420fef79 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Sat, 25 Jul 2026 01:04:46 -0700 Subject: [PATCH 2/7] [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 78b322514124406fc6fcf7d7958736925c0bc1ad Mon Sep 17 00:00:00 2001 From: coder-jayp Date: Sat, 25 Jul 2026 18:32:05 +0530 Subject: [PATCH 3/7] 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 b621923ae971cf75cbc2a93a82a495bda91cb22e Mon Sep 17 00:00:00 2001 From: Achintya P Date: Sat, 25 Jul 2026 12:15:44 -0700 Subject: [PATCH 4/7] [Test] Spec tests for Sequence burn-in, bootstrap and stride Executable contract for piece 3 of the #4039 split: Sequence gains burn_in, bootstrap and stride keyword arguments. burn_in records precede the anchor and never shift it (entries before the episode start are invalid and clamped to it); the learning region of length records starts at the anchor; bootstrap records follow it and obey the episode_boundary policy at episode ends; stride spaces the whole window uniformly. A per-record learning_mask info entry is True exactly on the learning region, step_in_sequence spans the window, and defaults reproduce base Sequence behavior. Negative counts or a non-positive stride raise ValueError. Tests are expected to fail until the implementation lands. --- test/rb/test_rb_core.py | 136 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 136 insertions(+) diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py index 83027b0117b..d7267ee95be 100644 --- a/test/rb/test_rb_core.py +++ b/test/rb/test_rb_core.py @@ -1479,6 +1479,142 @@ def test_invalid_length_raises(self): Sequence(length=4, episode_boundary="teleport") +class TestSequenceBurnInBootstrap: + """Executable spec for Sequence burn-in, bootstrap and stride (#4039, piece 3). + + Contract pinned by this class: + + - ``Sequence(..., burn_in=0, bootstrap=0, stride=1)`` extends the window + around each anchor: ``burn_in`` records before the anchor, the learning + region of ``length`` records starting at the anchor, then ``bootstrap`` + records after it. Total records per anchor: + ``burn_in + length + bootstrap``. + - A new per-record ``"learning_mask"`` info entry is True exactly on the + learning region; ``"validity_mask"`` keeps its meaning (real, in-episode + data). ``"step_in_sequence"`` runs 0..total-1 across the whole window. + - Burn-in never shifts the anchor: entries before the anchor's episode + start are invalid and clamped to the episode start, whatever the + boundary policy. Bootstrap entries follow ``episode_boundary`` at the + episode end like any tail entry. + - ``stride`` spaces the records of the window uniformly. + - Defaults reproduce the base Sequence behavior exactly. + - ``burn_in < 0``, ``bootstrap < 0`` or ``stride < 1`` raise ``ValueError``. + """ + + 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) + done = torch.zeros(capacity, 1, dtype=torch.bool) + for idx in done_at: + done[idx] = True + rb.extend( + TensorDict( + { + "obs": torch.arange(capacity, dtype=torch.float32), + ("next", "done"): done, + }, + batch_size=[capacity], + ) + ) + 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_burn_in_prepends_records(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand(rb, Sequence(length=2, burn_in=2), [8]) + assert index.tolist() == [6, 7, 8, 9] + assert info["learning_mask"].tolist() == [False, False, True, True] + assert info["validity_mask"].all() + assert info["step_in_sequence"].tolist() == [0, 1, 2, 3] + + def test_burn_in_clamped_at_episode_start(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand(rb, Sequence(length=2, burn_in=2), [6]) + assert index.tolist() == [6, 6, 6, 7] + assert info["validity_mask"].tolist() == [False, False, True, True] + assert info["learning_mask"].tolist() == [False, False, True, True] + + def test_bootstrap_appends_records(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand(rb, Sequence(length=2, bootstrap=1), [6]) + assert index.tolist() == [6, 7, 8] + assert info["learning_mask"].tolist() == [True, True, False] + assert info["validity_mask"].all() + + def test_bootstrap_masked_at_episode_end(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + index, info = self._expand( + rb, Sequence(length=2, bootstrap=1, episode_boundary="pad"), [8] + ) + assert index.tolist() == [8, 9, 9] + assert info["validity_mask"].tolist() == [True, True, False] + assert info["learning_mask"].tolist() == [True, True, False] + + def test_stride_spaces_the_window(self): + Sequence = self._sequence_cls() + rb = self._make_storage(done_at=(9,)) + index, info = self._expand(rb, Sequence(length=3, stride=2), [0]) + assert index.tolist() == [0, 2, 4] + assert info["step_in_sequence"].tolist() == [0, 1, 2] + assert info["validity_mask"].all() + + def test_defaults_match_base_sequence(self): + Sequence = self._sequence_cls() + rb = self._make_storage() + base_index, base_info = self._expand(rb, Sequence(length=3), [1, 6]) + ext_index, ext_info = self._expand( + rb, Sequence(length=3, burn_in=0, bootstrap=0, stride=1), [1, 6] + ) + assert base_index.tolist() == ext_index.tolist() + torch.testing.assert_close( + base_info["validity_mask"], ext_info["validity_mask"] + ) + assert ext_info["learning_mask"].all() + + def test_window_size_through_sampling(self): + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer( + storage=LazyTensorStorage(30), + batch_size=2, + sample_unit=Sequence(length=3, burn_in=2, bootstrap=1), + ) + rb.extend( + TensorDict( + { + "obs": torch.arange(30, dtype=torch.float32), + ("next", "done"): torch.zeros(30, 1, dtype=torch.bool), + }, + batch_size=[30], + ) + ) + sample = rb.sample() + assert sample.batch_size[0] == 2 * (2 + 3 + 1) + assert "learning_mask" in sample.keys() + assert sample["learning_mask"].sum() == 2 * 3 + + def test_validation(self): + Sequence = self._sequence_cls() + with pytest.raises(ValueError): + Sequence(length=3, burn_in=-1) + with pytest.raises(ValueError): + Sequence(length=3, bootstrap=-1) + with pytest.raises(ValueError): + Sequence(length=3, stride=0) + + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown) From ffc0c7f1335d62883d2cbd96a2196ba0d16ed699 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Sat, 25 Jul 2026 12:16:34 -0700 Subject: [PATCH 5/7] [Test] Spec tests for sequence priority semantics and distribution invariance Executable contract for the first half of piece 4 of the #4039 split. Priorities live per anchor: the Sequence unit adds a per-record anchor_index info entry so priorities of sampled sequences can be updated through the ordinary update_priority path, and per-anchor sampler entries such as importance weights are expanded block-constant across each window. Range expansion must not change the anchor selection distribution: seeded statistical tests compare uniform and prioritized anchor frequencies with and without the unit, and an update through anchor_index must steer subsequent sampling. The SliceSampler compatibility adapters and the structured [B, T] layout remain the second half of piece 4 and get their own spec once this lands. Tests are expected to fail until the implementation lands. --- test/rb/test_rb_core.py | 121 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py index d7267ee95be..f04386d322b 100644 --- a/test/rb/test_rb_core.py +++ b/test/rb/test_rb_core.py @@ -1615,6 +1615,127 @@ def test_validation(self): Sequence(length=3, stride=0) +class TestSequencePrioritySemantics: + """Executable spec for sequence priority semantics and distribution + invariance (#4039, piece 4, first half). + + Contract pinned by this class: + + - Priorities live per anchor. The Sequence unit adds a per-record + ``"anchor_index"`` info entry (the storage index of each record's + anchor) so priorities can be updated for the anchors of sampled + sequences through the ordinary ``update_priority`` path. + - Per-anchor sampler entries such as importance weights are expanded + block-constant: reshaped to ``[anchors, window]``, every row is + constant. + - Range expansion does not change the anchor selection distribution: + anchors drawn through a Sequence unit follow the same distribution as + anchors drawn through Transition, both for uniform and prioritized + sampling. + + The distribution tests are statistical: seeded, with wide tolerance + bands chosen to keep them deterministic in CI. + """ + + def _sequence_cls(self): + from torchrl.data.replay_buffers import sample_units + + return sample_units.Sequence + + def _make_rb(self, sampler=None, unit=None, capacity=20, batch_size=4): + rb = ReplayBuffer( + storage=LazyTensorStorage(capacity), + sampler=sampler if sampler is not None else RandomSampler(), + batch_size=batch_size, + sample_unit=unit, + generator=torch.Generator().manual_seed(0), + ) + rb.extend( + TensorDict( + { + "obs": torch.arange(capacity, dtype=torch.float32), + ("next", "done"): torch.zeros(capacity, 1, dtype=torch.bool), + }, + batch_size=[capacity], + ) + ) + return rb + + def _anchor_counts(self, rb, unit_length, draws=400, capacity=20): + counts = torch.zeros(capacity) + for _ in range(draws): + _, info = rb.sample(return_info=True) + if unit_length == 1: + anchors = torch.as_tensor(info["index"]).reshape(-1) + else: + anchors = torch.as_tensor(info["anchor_index"]).reshape(-1)[ + ::unit_length + ] + counts += torch.bincount(anchors, minlength=capacity) + return counts / counts.sum() + + def test_anchor_index_metadata(self): + Sequence = self._sequence_cls() + rb = self._make_rb() + index, info = Sequence(length=3).expand(torch.tensor([1, 6]), {}, rb._storage) + assert info["anchor_index"].tolist() == [1, 1, 1, 6, 6, 6] + + def test_weights_are_block_constant(self): + Sequence = self._sequence_cls() + length = 4 + rb = self._make_rb( + sampler=PrioritizedSampler(20, alpha=0.7, beta=0.9), + unit=Sequence(length=length), + ) + rb.update_priority(index=torch.arange(20), priority=torch.rand(20) + 0.1) + _, info = rb.sample(return_info=True) + records = 4 * length + for value in info.values(): + value = torch.as_tensor(value).reshape(-1) + if value.numel() != records or value.dtype == torch.bool: + continue + blocks = value.reshape(4, length).float() + torch.testing.assert_close(blocks, blocks[:, :1].expand_as(blocks)) + + def test_uniform_anchor_distribution_unchanged(self): + Sequence = self._sequence_cls() + expected = 1.0 / 20 + for unit, unit_length in ((None, 1), (Sequence(length=2), 2)): + rb = self._make_rb(unit=unit) + freqs = self._anchor_counts(rb, unit_length) + assert (freqs > 0.4 * expected).all() + assert (freqs < 1.9 * expected).all() + + def test_prioritized_anchor_distribution_unchanged(self): + Sequence = self._sequence_cls() + priorities = torch.ones(20) + priorities[10:] = 9.0 + for unit, unit_length in ((None, 1), (Sequence(length=2), 2)): + rb = self._make_rb( + sampler=PrioritizedSampler(20, alpha=1.0, beta=1.0), + unit=unit, + ) + rb.update_priority(index=torch.arange(20), priority=priorities) + freqs = self._anchor_counts(rb, unit_length) + high_share = freqs[10:].sum().item() + assert 0.8 < high_share < 0.98 + + def test_update_priority_via_anchor_index(self): + Sequence = self._sequence_cls() + rb = self._make_rb( + sampler=PrioritizedSampler(20, alpha=1.0, beta=1.0), + unit=Sequence(length=2), + ) + rb.update_priority(index=torch.arange(20), priority=torch.ones(20)) + _, info = rb.sample(return_info=True) + boosted = int(torch.as_tensor(info["anchor_index"]).reshape(-1)[0]) + rb.update_priority( + index=torch.tensor([boosted]), priority=torch.tensor([1000.0]) + ) + freqs = self._anchor_counts(rb, unit_length=2, draws=200) + assert freqs[boosted] > 0.5 + + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown) From a420ef4de7c0fefbfea60897bf0756f3b81aa796 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Sat, 25 Jul 2026 12:31:31 -0700 Subject: [PATCH 6/7] [Feature] Extend Sequence with burn-in, bootstrap, stride and per-anchor priority semantics Pieces 3 and 4 (first half) of the #4039 split, on top of the Sequence unit from #4050. The window around each anchor becomes burn_in records before the anchor, the learning region of length records starting at it, and bootstrap records after it, with stride spacing the whole window uniformly. A per-record learning_mask info entry is True exactly on the learning region. Burn-in never shifts the anchor: entries before the episode start are invalid and clamp to it; bootstrap entries obey the episode_boundary policy at episode ends; defaults reproduce the previous behavior exactly. Priorities live per anchor: a per-record anchor_index info entry reports the storage index of each record's sampled anchor (the original anchor, not the stop-shifted one, since that is what the sampler's distribution selected), so priorities of sampled sequences update through the ordinary update_priority path, and per-anchor sampler entries such as importance weights expand block-constant across the window. Seeded distribution tests pin that range expansion does not bias anchor selection for uniform or prioritized sampling. Part of #4039. --- test/rb/test_rb_core.py | 32 ++++++-- torchrl/data/replay_buffers/sample_units.py | 90 +++++++++++++++------ 2 files changed, 91 insertions(+), 31 deletions(-) diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py index f04386d322b..3fddc9e90cc 100644 --- a/test/rb/test_rb_core.py +++ b/test/rb/test_rb_core.py @@ -1441,13 +1441,15 @@ def test_metadata_flows_into_tensordict_sample(self): 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() + valid = sample["validity_mask"].reshape(2, 4) + obs = sample["obs"].reshape(2, 4) + step = sample["step_in_sequence"].reshape(2, 4).float() starts = obs - step - assert (starts[0:4] == starts[0]).all() - assert (starts[4:8] == starts[4]).all() - assert valid.all() + assert valid[:, 0].all() + assert (valid.int().diff(dim=1) <= 0).all() + for row in range(2): + row_valid = valid[row] + assert (starts[row][row_valid] == starts[row][0]).all() def test_prioritized_weights_expanded_to_records(self): Sequence = self._sequence_cls() @@ -1690,12 +1692,26 @@ def test_weights_are_block_constant(self): rb.update_priority(index=torch.arange(20), priority=torch.rand(20) + 0.1) _, info = rb.sample(return_info=True) records = 4 * length - for value in info.values(): + structural = { + "index", + "index_generation", + "sequence_id", + "step_in_sequence", + "anchor_index", + } + checked = 0 + for key, value in info.items(): value = torch.as_tensor(value).reshape(-1) - if value.numel() != records or value.dtype == torch.bool: + if ( + key in structural + or value.numel() != records + or value.dtype == torch.bool + ): continue blocks = value.reshape(4, length).float() torch.testing.assert_close(blocks, blocks[:, :1].expand_as(blocks)) + checked += 1 + assert checked > 0 def test_uniform_anchor_distribution_unchanged(self): Sequence = self._sequence_cls() diff --git a/torchrl/data/replay_buffers/sample_units.py b/torchrl/data/replay_buffers/sample_units.py index 55cd3a3d3e8..b29fb69def1 100644 --- a/torchrl/data/replay_buffers/sample_units.py +++ b/torchrl/data/replay_buffers/sample_units.py @@ -100,6 +100,7 @@ def expand( ) -> tuple[torch.Tensor | tuple, dict[str, Any]]: return index, info + class Sequence(SampleUnit): """Expands anchors into a fixed-length sequence of records. @@ -115,6 +116,19 @@ class Sequence(SampleUnit): Defaults to `"pad"`. done_key (str or tuple, optional): the key for the end-of-episode flag. Defaults to `("next", "done")`. + burn_in (int, optional): number of records preceding the anchor, + marked False in the ``"learning_mask"`` info entry. Burn-in never + shifts the anchor: entries before the anchor's episode start are + invalid and clamp to the episode start. Defaults to 0. + bootstrap (int, optional): number of records following the learning + region, marked False in ``"learning_mask"`` and subject to the + ``episode_boundary`` policy at episode ends. Defaults to 0. + stride (int, optional): spacing between the records of the window. + Defaults to 1. + + The unit also reports a per-record ``"anchor_index"`` info entry holding + the storage index of each record's sampled anchor, so priorities of + sampled sequences can be updated per anchor through ``update_priority``. """ def __init__( @@ -122,14 +136,26 @@ def __init__( length: int, episode_boundary: str = "pad", done_key: str | tuple[str, ...] | None = ("next", "done"), + burn_in: int = 0, + bootstrap: int = 0, + stride: int = 1, ): 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 burn_in < 0: + raise ValueError(f"burn_in must be non-negative, got {burn_in}.") + if bootstrap < 0: + raise ValueError(f"bootstrap must be non-negative, got {bootstrap}.") + if stride < 1: + raise ValueError(f"stride must be strictly positive, got {stride}.") self.length = length self.episode_boundary = episode_boundary self.done_key = done_key + self.burn_in = burn_in + self.bootstrap = bootstrap + self.stride = stride def expand( self, @@ -138,28 +164,39 @@ 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." + ) anchor = index.clone() B = anchor.shape[0] device = anchor.device + total = self.burn_in + self.length + self.bootstrap expanded_info = {} for k, v in info.items(): val = torch.as_tensor(v) - expanded_info[k] = val.repeat_interleave(self.length, dim=0) + expanded_info[k] = val.repeat_interleave(total, dim=0) - seq_id = torch.arange(B, device=device).repeat_interleave(self.length) - step_idx = torch.arange(self.length, device=device).repeat(B) + steps = torch.arange(total, device=device, dtype=torch.long) + learning = (steps >= self.burn_in) & (steps < self.burn_in + self.length) - expanded_info["sequence_id"] = seq_id - expanded_info["step_in_sequence"] = step_idx + expanded_info["sequence_id"] = torch.arange(B, device=device).repeat_interleave( + total + ) + expanded_info["step_in_sequence"] = steps.repeat(B) + expanded_info["learning_mask"] = learning.repeat(B) + expanded_info["anchor_index"] = anchor.repeat_interleave(total) - 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 = ((steps - self.burn_in) * self.stride).unsqueeze(0).expand(B, total) + validity = torch.ones((B, total), 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 + 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, @@ -175,7 +212,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,19 +224,24 @@ 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": - 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) + anchor_eff = anchor + if self.episode_boundary == "stop": + shortfall = self.stride * (self.length - 1) - dist_to_stop + shift = torch.clamp( + shortfall, min=torch.zeros_like(shortfall), max=dist_from_start + ) + anchor_eff = anchor - shift + dist_to_stop = dist_to_stop + shift + dist_from_start = dist_from_start - shift + + validity = (offset <= dist_to_stop.unsqueeze(1)) & ( + offset >= -dist_from_start.unsqueeze(1) + ) + clamped_offset = torch.minimum(offset, dist_to_stop.unsqueeze(1)) + clamped_offset = torch.maximum( + clamped_offset, -dist_from_start.unsqueeze(1) + ) + indices = anchor_eff.unsqueeze(1) + clamped_offset else: indices = anchor.unsqueeze(1) + offset max_len = storage.max_size From 9dbc9bb0d65ab0202eb1eb6d02dd082bf61f4985 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Sat, 25 Jul 2026 13:26:16 -0700 Subject: [PATCH 7/7] [BugFix] Keep include_reset sequences inside written storage include_reset wrapped indices modulo storage.max_size, so on a storage that has not wrapped yet a burn-in window walking backward from an early anchor (or a forward window crossing the write cursor) landed on never-written slots and read uninitialized data. Crossing the physical seam is only meaningful once the ring has actually wrapped: indices now wrap modulo the written length when the storage is full, and clamp with validity_mask=False when it is not. Regression test covers both directions on a half-full storage. --- test/rb/test_rb_core.py | 22 +++++++++++++++++++++ torchrl/data/replay_buffers/sample_units.py | 11 ++++++++--- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/test/rb/test_rb_core.py b/test/rb/test_rb_core.py index 3fddc9e90cc..44ab6e38f68 100644 --- a/test/rb/test_rb_core.py +++ b/test/rb/test_rb_core.py @@ -1607,6 +1607,28 @@ def test_window_size_through_sampling(self): assert "learning_mask" in sample.keys() assert sample["learning_mask"].sum() == 2 * 3 + def test_include_reset_does_not_read_unwritten_slots(self): + Sequence = self._sequence_cls() + rb = TensorDictReplayBuffer(storage=LazyTensorStorage(20), 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], + ) + ) + backward = Sequence(length=3, burn_in=2, episode_boundary="include_reset") + index, info = backward.expand(torch.tensor([0]), {}, rb._storage) + assert index.min() >= 0 + assert index.max() < 10 + assert info["validity_mask"].tolist() == [False, False, True, True, True] + forward = Sequence(length=4, episode_boundary="include_reset") + index, info = forward.expand(torch.tensor([8]), {}, rb._storage) + assert index.max() < 10 + assert info["validity_mask"].tolist() == [True, True, False, False] + def test_validation(self): Sequence = self._sequence_cls() with pytest.raises(ValueError): diff --git a/torchrl/data/replay_buffers/sample_units.py b/torchrl/data/replay_buffers/sample_units.py index b29fb69def1..6482ec8259e 100644 --- a/torchrl/data/replay_buffers/sample_units.py +++ b/torchrl/data/replay_buffers/sample_units.py @@ -242,11 +242,16 @@ def expand( clamped_offset, -dist_from_start.unsqueeze(1) ) indices = anchor_eff.unsqueeze(1) + clamped_offset + indices = indices % max_len else: - indices = anchor.unsqueeze(1) + offset - max_len = storage.max_size + written = len(storage) + raw = anchor.unsqueeze(1) + offset + if getattr(storage, "_is_full", False): + indices = raw % written + else: + validity = (raw >= 0) & (raw < written) + indices = raw.clamp(min=0, max=max(written - 1, 0)) - indices = indices % max_len expanded_info["validity_mask"] = validity.flatten() return indices.flatten(), expanded_info