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

[theap06]

Description

Adds ReplayBuffer.update_if_present, a generation-safe conditional mutation API for replay fields, on top of the opt-in generation stamps merged in #4046.

rb = TensorDictReplayBuffer(
    storage=LazyTensorStorage(100),
    writer=TensorDictRoundRobinWriter(track_generations=True),
)
rb.extend(data)
sample = rb.sample(32)
result = rb.update_if_present(
    sample,
    TensorDict({"priority": new_priority}, batch_size=sample.batch_size),
)

The API:

• Validates the complete patch before making any writes, including keys, dtypes, and shapes.
• Compares generations and applies multi-key patches while holding the replay-buffer and writer locks.
• Skips entries whose slots were emptied or overwritten after sampling.
• Keeps a handle reusable until its slot is actually rewritten.
• Supports nested keys and replay buffers with ndim > 1.
• Raises a clear RuntimeError for unsupported storages and writers without generation tracking.
• Uses one actor RPC for Ray replay buffers; distributed replay buffers report the unsupported capability explicitly.
• Returns a ConditionalUpdateResult with the applied mask and update/skip counts.

Version-comparison support remains in #4049.

Tests and documentation

• Adds 14 focused TestUpdateIfPresent cases, including TensorDict patches, stale handles, validation, nested keys, and multidimensional storage.
• Adds Ray transport coverage and a replay-buffer benchmark.
• Documents the API and its generation-safe semantics in the replay-buffer reference.
• Locally: 497 core replay-buffer tests passed, the focused Ray test and doctests passed, and all pre-commit hooks passed.
• The broader replay-buffer run reached 4,376 passed and 891 skipped; the remaining failures require unavailable local dependencies or services (h5py, working Torch Inductor/OpenMP metadata, and a resolvable Gloo hostname).

Part of #4040. Builds on #4046.

[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 b8bf0a7 with merge base d7659c7 ():

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.

[theap06]

Rebased onto current main and #4046, resolved the obsolete generation-stamp history, and marked this ready for review. The focused core, Ray, and doctest coverage passes locally. I do not have permission to request reviewers or apply labels; @vmoens, could you review and add the ci/optdeps label, then rerun CI after the label is present? This PR touches the Ray optional-dependency integration.

[vmoens]

Reviewed rebased head b8bf0a7a0bc3ec5b995569aaf1414f86ac2553b7, including the latest generation-stamp reconciliation commit and the request for optional-dependency CI.

One blocking correctness issue is at torchrl/data/replay_buffers/replay_buffers.py in the live = current == generation comparison. Generation -1 explicitly means “no usable stamp”, but equality treats (-1, -1) as live. I reproduced this with a size-10 lazy storage containing only 5 records: update_if_present(index=tensor([7]), generation=tensor([-1]), patch={"x": tensor([99.])}) returns updated=True and writes slot 7 even though len(rb) == 5. Please exclude negative current/captured generations (and add the never-written-slot regression) so only genuine live handles can mutate storage.

Targeted result: all 14 TestUpdateIfPresent tests pass, which confirms the missing case. After the fix, the author-requested ci/optdeps label/rerun is still appropriate because the diff changes the Ray integration.