[Feature Request] Separate replay selection distributions from sampling units and boundary policies
@vmoens is already working on this.
Since Jul 22, 2026.
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 487
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 207
Description
Motivation
Transition sampling and trajectory-slice sampling currently expose different sampler classes even
though they combine several orthogonal decisions:
- how an anchor is selected and weighted;
- how that anchor expands into a transition, sequence or complete trajectory;
- how episode and stored-segment boundaries are handled;
- how short/ragged output is represented.
This makes combinations such as prioritized sequence starts, recurrent burn-in, n-step bootstrap
suffixes and alternative padding policies harder to express consistently. It also makes individual
samplers responsible for both probability distributions and trajectory-range mechanics.
Proposal
Introduce an optional composition layer with three responsibilities:
- Anchor sampler: selects record/trajectory anchors and returns probabilities or weights.
- Sampling unit / range resolver: expands anchors into transitions, fixed sequences, n-step
transitions or complete trajectories. - Boundary/output policy: controls episode/rollout crossing, padding, masks and output layout.
An illustrative API is:
rb = ReplayBuffer(
storage=storage,
sampler=PrioritizedSampler(...),
sample_unit=Sequence(
learn_length=32,
burn_in=8,
bootstrap=1,
episode_boundary="stop",
rollout_boundary="cross",
short_sequence="pad",
layout="flat", # preserve current TorchRL-friendly default
),
)
The exact API need not use these names. The important property is that the distribution chooses
anchors while the sampling unit owns range expansion and boundary enforcement.
Required semantics
- Transition, fixed-sequence and complete-trajectory units.
- Recurrent burn-in distinct from the learning portion.
- Optional n-step/bootstrap suffix.
- Configurable stride/overlap where meaningful.
- Episode boundary policies such as stop, pad or include-reset.
- Collector-rollout boundary policies independent of episode boundaries.
- Validity and learning masks for padded/ragged data.
- Clearly defined priority/importance-weight semantics: per anchor, per sequence or expanded per
transition. - Both current flat concatenated output and an optional structured
[B, T]index/data layout.
Backward compatibility
- Existing
RandomSampler,PrioritizedSampler,SliceSamplerand
PrioritizedSliceSamplerAPIs remain supported. - Existing slice samplers can initially be implemented as adapters over the new decomposition, or
the new API can be introduced alongside them. - Flat output and current
is_init/truncatedbehavior remain the compatibility default.
Acceptance criteria
- Uniform and prioritized anchor samplers work with transition and sequence units.
- Boundary behavior is tested independently of selection distribution.
- Burn-in, learning region, bootstrap suffix and validity masks are unambiguous in returned
metadata. - Existing
SliceSamplerbehavior is covered by compatibility tests. - Episode and collector-rollout boundaries can be configured independently.
- Range expansion works with ring wraparound and partial trajectories.
- Probability and importance-weight tests demonstrate that range expansion does not silently
change the documented sampling distribution.
Open questions
- Should sampling units be owned by
ReplayBuffer, composed inside a sampler, or passed to
sample()? - Should structured
[B, T]data be produced directly or only structured indices plus a gather
step? - How should priorities be reduced or assigned when a sampled sequence contains many records?
- Can this be introduced without expanding the public sampler type hierarchy further?
cc @theap06
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.