[Feature Request] Add generation-stamped replay indices for safe delayed priority updates

[vmoens]

Motivation

Replay samples currently expose physical storage indices. With a round-robin writer, an index can
be overwritten while the learner is computing a loss. A later priority update can therefore apply
the TD error of the sampled record to a different, newer record that now occupies the same slot.

The replay lock makes sampling and individual writes atomic, but it does not protect an index for
the full sample -> learner step -> priority update interval. Holding a storage lease for that whole
interval would prevent useful writer progress, so delayed updates need a way to detect slot reuse.

Proposal

Add an optional generation stamp to replay indices.

• A compatible writer increments the generation of every physical slot when that slot is reused.
• Sampling returns both the physical index and the generation that was observed for it.
• A conditional priority-update API compares both values while holding the appropriate replay/write
  lock and updates only entries that are still live.
• The operation returns which updates succeeded, allowing callers to report or ignore stale
  learner results.

For TensorDict replay buffers, the metadata could be represented as two tensor keys so it remains
easy to batch, reshape, move and serialize:

sample["index"]
sample["index_generation"]

An initial API could look like:

result = rb.update_priority_if_present(
    index=sample["index"],
    generation=sample["index_generation"],
    priority=sample["td_error"],
)

result.updated_mask
result.num_updated
result.num_stale

The exact public representation can differ, but the index and generation must travel together
through sampling, collation and multidimensional-storage reshaping.

Scope

Start with round-robin tensor storage and prioritized replay. The design should leave room for
other writers and storages to advertise generation support without requiring every storage backend
to implement it immediately.

Priority remains sampler metadata. This issue does not introduce general mutation of stored replay
fields; that can build on the same stamped-index mechanism separately.

Backward compatibility

• Existing raw index metadata and update_priority() remain supported.
• Existing calls retain their current behavior.
• Safe delayed updates use the new conditional method.
• Buffers whose writer/storage does not support generations should fail clearly at construction or
  when the conditional API is requested, rather than silently degrading to raw-index behavior.

Acceptance criteria

• Round-robin writes maintain a generation tensor with the same physical index space as storage.
• Samples return generation stamps matching their records.
• An update succeeds when a slot has not been reused.
• An update is skipped when the same physical index has been overwritten.
• Mixed batches report partial success correctly.
• The generation check and priority update are atomic with respect to writers.
• Multidimensional storage indices are supported.
• Generation state survives replay state_dict and dump/load round trips.
• Shared/multiprocess behavior is either supported and tested or rejected explicitly.
• Tests reproduce the delayed-learner wraparound race deterministically.

Open questions

• Should generation metadata be enabled by default for round-robin tensor storage or opt-in?
• Should the public object remain two tensors or become an opaque structured replay-index type?
• What generation dtype and overflow behavior should be guaranteed?
• Should all conditional update APIs share one result type?

cc @theap06

[harryfrzz]

Hi @vmoens, is this issue open for contribution or already being worked on? I'd like to pick it up if it's free.

[vmoens]

Maybe @theap06 is but I gave them a bunch of issues to work on so maybe this one is up for grabs?

[harryfrzz]

Thanks @vmoens! I'd love to pick this one up if @theap06 isn't already on it. If it's already taken, no worries at all. could you point me to another issue I could start on right now?

[theap06]

@harryfrzz Thanks for contributing! Feel free to take this up! A good first PR would be to implement generation tracking for the round-robin writer/storage.

[harryfrzz]

@harryfrzz Thanks for contributing! Feel free to take this up! A good first PR would be to implement generation tracking for the round-robin writer/storage.

Thanks @theap06, on it!

i'll start with the generation tracking on the round-robin side: tracking the reuse per slot, returning the stamp with the index when sampling and a test for the wraparound case.

quick question, should it be opt-in or on by default? I'll open a draft PR shortly either way. Also @vmoens could you assign this issue to me so others know it's being worked on?

[theap06]

@harryfrzz so sorry! I ended up building out the local round-robin writer/storage locally. I closed the PR. Feel free to continue working on this issue and I can provide feedback on your PR whenever you open it, even early. I did create a draft pr containing my design of the update_if_present API.
Also, regarding your question, I'd make it on by default, no flag: the counter is one int64 per slot and one increment per write.

The draft PR for update_if_present API is #4043

[harryfrzz]

@harryfrzz so sorry! I ended up building out the local round-robin writer/storage locally. I closed the PR. Feel free to continue working on this issue and I can provide feedback on your PR whenever you open it, even early. I did create a draft pr containing my design of the update_if_present API. Also, regarding your question, I'd make it on by default, no flag: the counter is one int64 per slot and one increment per write.

The draft PR for update_if_present API is #4043

Thanks @theap06, no worries at all!

I've opened a draft PR for this issue: #4046

On by default, no flag, as you suggested, each slot carries an int64 generation via writer.generations_of(index) / tracks_generations and sampling surfaces index_generation alongside index. It matches the spec your #4043 builds on, so update_if_present should rebase onto it cleanly.