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

[vmoens]

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

[theap06]

@coder-jayp we can split this into 4 PRs:

1. Interfaces and wiring - I'll take this one. The sample_unit= arg on ReplayBuffer, a SampleUnit base class, and a Transition unit that keeps current behavior as-is. Settles the first open question too.
   Feel free to take up this scope - 2. Sequence Unit: Fixed-length expansion from anchors, episode boundary policies (stop/pad/include-reset), validity masks, wraparound handling. SliceSampler already has boundary recovery machinery you can reuse. I'll post pointers + spec tests once 1 is up, like the compare-version piece in [Feature Request] Add generation-safe conditional updates for mutable replay fields #4040. Start with whichever of the two you prefer.
2. Burn-in, bootstrap suffix, stride/overlap, builds on 2, we can figure out ownership later.
3. Priority semantics + the distribution tests from the acceptance criteria, then SliceSampler compat and the [B, T] layout.

[coder-jayp]

@theap06 Sounds perfect, I will take on Part 2 (Sequence Unit). I am ready to start the implementation as soon as you post the spec tests and pointers.

[theap06]

@coder-jayp sequence unit spec tests are on my fork, branch spec/sequence-unit; branch off #4045. full semantics are in the test docstrings.
General feedback: with this design sample(batch_size=B) returns B*length records, which differs from SliceSampler's batch_size convention.

[coder-jayp]

@theap06 I have opened the PR for the sequence unit and I would like to work on 3 and 4 next. Do you happen to have the spec tests ready for them?

[theap06]

@coder-jayp I'm taking 3 and 4 myself. let's land #4049 and #4050 first. Just opened #4051 for them, stacked on your #4050 branch so the stack stays consistent; it'll rebase as your PR evolves.

[coder-jayp]

@theap06 Sounds like a great plan. Standing by for any review feedback on my PRs to help get them merged.