[Feature Request] Add generation-safe conditional updates for mutable replay fields
@vmoens is already working on this.
Since Jul 22, 2026.
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 484
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 207
Description
Motivation
Some replay algorithms update stored data after sampling, not only sampler priorities. Examples
include:
- refreshing recurrent hidden states;
- caching embeddings or target-model outputs;
- attaching asynchronous labels or post-processing results;
- updating other explicitly mutable replay fields.
Existing physical-index mutation can target the wrong record if a round-robin writer reuses the
slot between sampling and update. Holding sampled records pinned throughout a learner step would
avoid the race but can create severe replay backpressure.
Dependency
This should build on generation-stamped replay indices. Priority updates remain sampler metadata
and should use their own conditional API; this issue covers stored replay fields.
Proposal
Add a best-effort conditional update API for storage implementations that support generation
validation:
result = rb.update_if_present(
index=sample["index"],
generation=sample["index_generation"],
patch={"recurrent_state": refreshed_state},
)
The operation validates all requested keys, shapes and dtypes before mutating storage, then updates
only records whose (index, generation) is still live. It reports updated and stale elements.
An optional compare-version facility could prevent an older asynchronous computation from
overwriting a newer refresh:
result = rb.update_if_present(
index=index,
generation=generation,
patch=patch,
version_key="state_model_version",
version=current_model_version,
require_newer=True,
)
The version feature may be deferred, but the base API should not preclude it.
Schema and storage contract
- Storages may restrict conditional mutation to fields designated mutable at initialization.
- Structural changes such as adding keys or changing shapes are out of scope.
- Generation validation and each record's patch write are atomic with respect to slot reuse.
- The API documents whether atomicity is per record or across the whole batch. Per-record atomicity
with a returned mask is likely the useful default. - Validation failure before writing must not leave a partially applied malformed patch.
- Multiple keys in one record's patch must not become observably torn to concurrent readers.
Backward compatibility
- Existing
set_at_,set_andupdate_APIs remain unchanged. - Conditional mutation is opt-in and supported first by compatible tensor storages/writers.
- Unsupported storage backends raise a capability error instead of performing an unsafe raw-index
update.
Acceptance criteria
- Live stamped indices update successfully.
- Reused slots are skipped and never modified.
- Mixed live/stale batches return an accurate mask and counts.
- Shape, dtype and key validation happen before mutation.
- Multi-key patches are atomic per record.
- Concurrent sample/write/update stress tests do not expose torn records.
- Optional version comparison, if implemented, rejects older updates deterministically.
- Multidimensional storage is covered.
- Dump/load behavior for mutable fields, slot generations and optional model versions is tested.
Open questions
- Should mutable keys be declared by storage, replay buffer, or a TensorDict schema object?
- Is per-record atomicity sufficient, or do any target algorithms need all-or-nothing batch updates?
- Should the returned result be a TensorDict, dataclass or simple boolean mask?
- Should compare-version semantics be generic compare-and-swap rather than a dedicated version
convention?
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.