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