feat: support staleness-window in ReplayBufferNew

[yuki-97]

Part of RL-727. Stacks on #2448.

Implements ReplayBufferNew, a temporary replacement for ReplayBuffer until TQReplayBuffer is ready.

Motivation: ReplayBuffer.sample() requires target_weight_version == current_weight_version, which stalls training when the exact-match trajectories haven't arrived yet (buffer starvation). ReplayBufferNew fixes this by allowing slightly older trajectories to be used, with an importance-sampling correction.

Changes:

• Add max_staleness config: trajectories with trainer_version - weight_version > max_staleness are evicted at the start of each sample() call.
• sample() selects from the staleness window [trainer_version - max_staleness, trainer_version], removing the strict target_weight_version == current_weight_version gate.
• Add sample_freshest_first flag (default True): when True, selects the highest-version trajectories first; when False, uses FIFO (insertion order).
• target_weight_versions is intentionally unused in ReplayBufferNew — it gates generation on specific trainer steps, causing generation pauses. Will be removed when cleaning up after TQReplayBuffer lands.
• Unit tests covering eviction, staleness-window sampling, freshest-first ordering, and FIFO ordering.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[mehraakash]

👍

[mehraakash]

Could we refactored into another function:

  def _remove_indices(self, indices: Iterable[int]) -> None:
      for idx in sorted(indices, reverse=True):
          self.trajectory_versions.pop(idx)
          self.target_weight_versions.pop(idx)
          self.trajectories.pop(idx)

can then use it in _evict and sample and provide different Iterables?

[yuki-97]

good idea! e743b7d

[mehraakash]

See comment below on adding it in a function.

[mehraakash]

I know we want to eventually get rid of target_weight_versions but since we inherited from ReplayBufferImpl that list will be created. So we either keep that state aligned or remove it?

[yuki-97]

also covered in e743b7d

[mehraakash]

I think we need a follow-up task here before wiring this in:

ReplayBufferNew removes exact target matching in sample() here, but the collector still enforces
target-version reservation and generation-limit pauses through last_target_weight_already_generated.

For end-to-end staleness-window sampling, the collector needs a mode that generates based on
current generation_weight_version and buffer/backpressure capacity, and not future target_weight_version slots.

We'll control generation using SingleController by:

• Buffer Capacity
• Inflight Semaphore
• Refit pause
• Any manual Pause
• Dataloader availability

[yuki-97]

For end-to-end staleness-window sampling, the collector needs a mode that generates based on
current generation_weight_version and buffer/backpressure capacity, and not future target_weight_version slots.

yes, that's what we will do in SingleController, and this PR is a foundational work of it. when implementing SingleController we can just continue based on this.

currently collector still have generation-limit, and if we want to update it, will need many changes. and given that we will eventually use SingleController to control this, I think there's no need for us to update the collector here.

I think we can get this PR in first (have unit test to guard the new implementation), and it won't be a dead code since we'll definitely use it later. I'm also fine with keep the PR until we implement SingleController. wdyt? @mehraakash @terrykong

[mehraakash]

Yeah, just meant that we should add a note that end-to-end staleness-window sampling will get completed when ...

[yuki-97]

/ok to test e743b7d

[mehraakash]

LGTM

[saumishr]

@yuki-97 Please add a disclaimer on the new interfaces that its a work in progress and is subject to change. Users should not be using these WIP interfaces.