pytorch / pytorch/rl

[Feature Request] Add generation-safe conditional updates for mutable replay fields

Open
#4,040 10 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
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_ and update_ 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

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.