pytorch / pytorch/rl

[Feature Request] Separate replay selection distributions from sampling units and boundary policies

Open
#4,039 6 comments 1 reaction 1 assignee View on GitHub

@vmoens is already working on this.

Since Jul 22, 2026.

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

  1. how an anchor is selected and weighted;
  2. how that anchor expands into a transition, sequence or complete trajectory;
  3. how episode and stored-segment boundaries are handled;
  4. 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, SliceSampler and
    PrioritizedSliceSampler APIs 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/truncated behavior 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 SliceSampler behavior 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.