[Feature] Add SampleUnit composition point for replay sampling

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
 -------------------------
 

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)

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",

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",

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):