[Feature] Add ReplayBuffer.update_if_present for generation-safe conditional updates

[theap06]

Description

Second step of the conditional replay-update RFC: a best-effort conditional mutation API for stored replay fields, built on generation-stamped replay slots (#4041).

Some algorithms update stored data after sampling (refreshed recurrent states, cached embeddings, asynchronously computed labels). With a round-robin writer, the physical slot sampled from may hold a different record by write-back time, so an unconditional index write corrupts fresh data. update_if_present closes that race optimistically instead of pinning the buffer:

sample = buffer.sample()
result = buffer.update_if_present(
    index=sample["index"],
    generation=sample["index_generation"],
    patch={"recurrent_state": refreshed_state},
)
result.updated        # bool mask aligned with the input index order
result.updated_count  # live records patched
result.stale_count    # recycled/emptied records skipped, content untouched

Semantics, matching the RFC's acceptance criteria:

• Validate-then-write: the whole patch is checked against the storage (key existence -> KeyError, shape/dtype -> ValueError) before any mutation; a failed validation leaves storage byte-for-byte untouched, including multi-key patches where only one key is invalid.
• Per-record atomicity: the generation comparison and the patch write share one replay-lock acquisition, so a concurrent extend cannot interleave between validation of a record and its write; a stress test with a continuously extending writer thread verifies multi-key patches are never observed torn.
• Handles are not consumed: updating refreshes a record's content, not its identity; the same (index, generation) pair keeps working until the slot is rewritten.
• Capability-gated: tensor storages advertise supports_conditional_update; ListStorage and writers without generation tracking raise a clear capability error instead of performing an unsafe raw-index write.
• Nested keys and multidimensional (ndim>1) storages are supported.

Remote buffers: RayReplayBuffer.update_if_present delegates to the actor in a single RPC, with validation, comparison and write all executing inside the actor under its own lock; the distributed tensor transport raises a capability error pointing at transport="ray".

Deferred per the RFC, with the keyword-only signature left room to grow: the compare-version facility (version_key/require_newer), conditional priority updates (sampler metadata stays on its own API), and a declared-mutable-fields schema (any existing tensor field is currently patchable, validated at call time).

Testing

TestUpdateIfPresent, written spec-first: live updates, stale-skip with content verification, mask alignment under shuffled mixed batches, repeated updates on one handle, validation-before-write (including the mixed-validity patch case), nested keys, capability errors, empty() invalidation, sampled-handle round-trip, multidim storages, and the concurrent no-torn-records stress test. A Ray-gated test covers the one-RPC delegation with a wraparound. Full test/rb/ suite passes (4205 tests).

Note

Draft, deliberately. The generation-stamp implementation (step 1, #4041) is being built by @harryfrzz; the stamp commits currently present in this branch are reference scaffolding only and are NOT part of this PR's claimed scope. Once the #4041 implementation lands, this PR will be rebased onto it so its diff contains only the conditional-update layer: update_if_present, ConditionalUpdateResult, the storage capability/validation helpers, and the Ray delegation.

Part of #4040: the compare-version facility and checkpoint round-trip coverage remain open on the issue for other contributors.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4043

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit 11715f6 with merge base ae421b9 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.