[Feature] Add generation-stamped replay slots for round-robin writer by harryfrzz · Pull Request #4046 · pytorch/rl

harryfrzz · 2026-07-24T09:09:34Z

Summary

A replay buffer index is a physical slot number, not a handle on a piece of
data. A round-robin writer reuses slots, so an index sampled at one point may
name completely different data a moment later — and nothing in the buffer tells
a consumer that happened.

This adds opt-in generation stamps: one counter per storage slot, advanced on
every write to that slot. A consumer that captured a stamp when it sampled can
compare it against the current stamp to know whether the slot still holds its
data.

rb = ReplayBuffer(
    storage=LazyTensorStorage(8),
    writer=RoundRobinWriter(track_generations=True),
)
rb.extend(torch.arange(8))
_, info = rb.sample(4, return_info=True)
index, generation = info["index"], info["index_generation"]

rb.extend(torch.arange(8, 11))            # overwrites slots 0, 1, 2
stale = rb.writer.generations_of(index) != generation
# index[stale] no longer holds the sampled data

This is the primitive that conditional writes ("update this record only if it is
still the one I read") need; it is useful on its own for async training and for
prioritized replay, where priorities are computed after the forward pass and
written back at an index that may since have been reused.

Design

Opt-in. The default is track_generations=False. Enabling it allocates one
int64 per storage slot and adds an "index_generation" entry to the sampler
output — which for tensordict buffers is a new sample key — so it is not
imposed on buffers that do not need it. sample() output is byte-identical to
before for every existing user.

The counters live on the storage, not on the writer. Two buffers sharing one
storage overwrite each other's slots, so a per-writer counter would let one
buffer's handles read as live after the other overwrote them — exactly the
staleness the feature exists to detect. The buffer is attached to the storage
object, and a writer registered against a storage that already has one adopts it
rather than replacing it.

Semantics. One stamp per write, not per extend call: an extend that wraps
the storage advances a reused slot once per write it receives. -1 means "no
usable stamp" (never written, out of range, or tracking disabled), not
"generation zero". empty() advances stamps rather than resetting them, so
handles taken before it correctly read as stale. Stamps are independent
per-slot counters and say nothing about write order between different slots.

Allocation. Storages small enough to allocate up front get one allocation, so
the buffer's shape never changes and the torch.compile extend/sample path does
not recompile. Larger and unbounded storages (ListStorage with no max_size
reports torch.iinfo(torch.int64).max) grow geometrically on demand.

Process-local, deliberately. The counters are not shared across processes. The
buffer is replaced rather than mutated when it grows, so a shared mapping would
silently stop tracking after the first growth. Cross-process staleness detection
needs a storage-owned fixed-size mapping and is a follow-up, not a half-working
share_memory_() call here.

Docs

docs/source/reference/data_replaybuffers.rst gains a "Detecting overwritten
slots: generation stamps" section (ref_buffers_generations) with the motivation,
a runnable example, the semantics above and the implementation notes, cross-
referenced from the Writers section. Writer.generations_of,
Writer.tracks_generations and RoundRobinWriter carry full docstrings.

Not in this PR

Cross-process visibility. Needs a storage-owned, fixed-size mapping; see
above.
ReplayBuffer(track_generations=True) passthrough. Enabling it means
passing your own writer today, as with samplers and storages. A buffer-level
kwarg would need matching fields on ReplayBufferConfig /
TensorDictReplayBufferConfig (CLAUDE.md [Feature Request] Trainer saving and loading utils #14) and is better added alongside
the first consumer.
Conditional writes (update_if_present with a version check) build on this
and are proposed separately.

Tested

test/rb: 4201 passed, 1029 skipped. Pre-PR baseline on the same merge base:
4175 passed, 1025 skipped — so +26 new tests, +4 CUDA-gated skips, no regressions.
The four pre-existing tests this PR previously had to patch pass unmodified again.
Doctest and docs examples verified to produce the output they claim.

pytorch-bot · 2026-07-24T09:09:38Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

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

meta-cla · 2026-07-24T09:09:40Z

Hi @harryfrzz!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

meta-cla · 2026-07-24T10:07:29Z

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

theap06

Thanks! Left some feedback.

theap06

Left some comments. Otherwise LGTM! cc @vmoens

harryfrzz · 2026-07-26T07:04:59Z

Left some comments. Otherwise LGTM! cc @vmoens

sure! i’ve pushed the changes you mentioned @theap06.

harryfrzz · 2026-07-29T15:42:31Z

@theap06 @vmoens any changes that i need to do here?

vmoens

With 84c1146 I made it optional to enable this - we can think about enabling it by default 2 releases after the next but the plan must be announced via deprecation warnings since this would be bc breaking.
I also added some doc.
Happy with this version if you are. Deprecation plan can follow in a follow up PR.

harryfrzz · 2026-08-05T07:46:56Z

With 84c1146 I made it optional to enable this - we can think about enabling it by default 2 releases after the next but the plan must be announced via deprecation warnings since this would be bc breaking. I also added some doc. Happy with this version if you are. Deprecation plan can follow in a follow up PR.

Happy with it, thanks. Opt-in is the right call and the storage owned counters fix the shared storage case. one thing that i want to point out is that the four CUDA tests I added in cfaa0ab still build default writers and assert on _writer._generation, which is None now that tracking is gated. They skip on CPU and are -m gpu-selected on the GPU job. full test/rb is green here (4332 passed, 898 skipped). Should I add track_generations=True to those four?
deprecation plan in a follow-up works for me!

…ytorch#4046) Drop this PR's own generation-stamp scaffold (writers.py machinery, always-on tracking, sample-info insertions and their test patches) in favor of the storage-owned, opt-in implementation from pytorch#4046: - update_if_present keeps its contract but now requires a writer constructed with track_generations=True and raises otherwise; its generation comparison and patch masks are device-aware so CUDA/MPS storages and mixed-device handles work. - TestUpdateIfPresent builds its buffers with tracking writers, gains a test for the non-tracking capability error; the superseded TestSlotGenerations and TestSampleGenerationInfo suites are removed (covered by pytorch#4046's TestWriterGeneration). - The Ray test threads a tracking writer factory to the remote buffer. - The wraparound benchmark enables tracking so it no longer measures a no-op; docs and docstrings describe the opt-in reality and link to the generation-stamp reference section. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ocument Addresses review feedback. - Opt-in. `RoundRobinWriter(track_generations=True)` now gates the feature; `tracks_generations` is an instance property rather than a class attribute set to True. Previously every replay buffer allocated an int64 per slot and every `sample()` grew an `"index_generation"` entry -- which for tensordict buffers is a new sample key. The four pre-existing tests this PR had to patch to `exclude("index", "index_generation")` are reverted, which is the check that the BC break is gone. - The counters live on the storage, not on the writer. Two buffers sharing one storage overwrite each other's slots, so a per-writer counter let buffer A's handles read as live after buffer B overwrote them -- defeating the one thing the feature is for. The buffer is attached to the storage; a writer registered against a storage that already has one adopts it. A storage-less writer (as `dumps`/`loads` and `load_state_dict` can be) keeps it locally until a storage is registered. - Dropped the `_GENERATION_UNBOUNDED = 2**40` sentinel. Rather than trying to classify a capacity as bounded or not, allocate eagerly below an explicit allocation limit (stable shape, no torch.compile recompile) and grow geometrically above it. Correct for genuinely unbounded storages (`torch.iinfo(torch.int64).max`) and for implausibly large bounded ones alike. - Dropped the `share_memory_()` calls. The buffer is replaced, not mutated, when it grows, so a shared mapping silently stops tracking after the first growth; cross-process visibility was never actually supported. Now documented as process-local instead of half-implemented. - `generations_of` no longer guesses that a 1-D tensor whose length matches `storage.ndim` is a coordinate vector: with `ndim == 2`, `generations_of(tensor([3, 7]))` silently returned one generation where the caller asked for two. Only `ndim >= 2` tensors are read as coordinate batches; pass a tuple for a single cell. - Documentation. A "Detecting overwritten slots: generation stamps" section in data_replaybuffers.rst covering why slot indices go stale, the semantics (one stamp per write, `-1` means no usable stamp, monotonic across `empty()`, not an ordering across slots) and the implementation notes (opt-in, storage- owned, allocation policy, process-local, multidim indexing, checkpointing), with a runnable example. Cross-referenced from the Writers section, and `Writer.generations_of` / `tracks_generations` / `RoundRobinWriter` carry full docstrings. - Tests for each: opt-in default leaves `sample()` untouched, shared-storage visibility, and the multidim 1-D index case. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Running it locally: LazyTensorStorage(4, ndim=2) with [4, 3] data gives a dim-0 capacity of 1, so the extend wrapped and the stamps were not all 0. Use LazyTensorStorage(12, ndim=2) (dim-0 capacity 3) with in-range indices, and also assert the [N, ndim] coordinate batch extend() returns is read as coordinates. test/rb: 4201 passed, 1029 skipped -- baseline is 4175/1025, so +26 new tests and no regressions. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

vmoens · 2026-08-07T14:25:19Z

Yes you are correct the tests had to be corrected which i did in 6d3af1a

…ytorch#4046) Drop this PR's own generation-stamp scaffold (writers.py machinery, always-on tracking, sample-info insertions and their test patches) in favor of the storage-owned, opt-in implementation from pytorch#4046: - update_if_present keeps its contract but now requires a writer constructed with track_generations=True and raises otherwise; its generation comparison and patch masks are device-aware so CUDA/MPS storages and mixed-device handles work. - TestUpdateIfPresent builds its buffers with tracking writers, gains a test for the non-tracking capability error; the superseded TestSlotGenerations and TestSampleGenerationInfo suites are removed (covered by pytorch#4046's TestWriterGeneration). - The Ray test threads a tracking writer factory to the remote buffer. - The wraparound benchmark enables tracking so it no longer measures a no-op; docs and docstrings describe the opt-in reality and link to the generation-stamp reference section. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

github-actions Bot added Feature New feature ReplayBuffers and removed Feature New feature labels Jul 24, 2026

harryfrzz mentioned this pull request Jul 24, 2026

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

Open

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 24, 2026

github-actions Bot added the Feature New feature label Jul 24, 2026

theap06 suggested changes Jul 25, 2026

View reviewed changes

Comment thread torchrl/data/replay_buffers/writers.py

Comment thread torchrl/data/replay_buffers/writers.py Outdated

harryfrzz requested a review from theap06 July 25, 2026 20:36

theap06 reviewed Jul 26, 2026

View reviewed changes

Comment thread torchrl/data/replay_buffers/writers.py

theap06 reviewed Jul 26, 2026

View reviewed changes

Comment thread torchrl/data/replay_buffers/writers.py

theap06 approved these changes Jul 26, 2026

View reviewed changes

vmoens marked this pull request as ready for review August 4, 2026 16:05

vmoens approved these changes Aug 4, 2026

View reviewed changes

github-actions Bot added the Documentation Improvements or additions to documentation label Aug 4, 2026

vmoens force-pushed the feature/rb-generation-tracking branch from 6f361a8 to 63471a8 Compare August 5, 2026 15:13

harryfrzz added 8 commits August 7, 2026 15:14

[Feature] Add opt-in generation tracking to RoundRobin writer

e1842b2

feat: Surface index_generation in replay sample info

fa9e722

feat: default-on slot generation tracking for round-robin writer

7e10c67

feat: align slot generation to 0-based reuse counter

8125e23

docs: drop verbose generation docstrings from RoundRobinWriter

bb3af3d

fix: allocate writer generation on the storage device

ff4fb4a

fix: keep unwritten generation slots at -1 on empty

164108d

docs: document generation bumps on writes

d4b4814

harryfrzz and others added 3 commits August 7, 2026 15:14

test: cover generation bump on wrapping extend

6b2d280

vmoens force-pushed the feature/rb-generation-tracking branch from 63471a8 to 7dc7c0e Compare August 7, 2026 14:14

[BugFix] Fix replay generation tracking edge cases

6d3af1a

github-actions Bot added the Trainers label Aug 7, 2026

[CI] Fix generation tracking lint

1fd72be

vmoens merged commit 5e2af20 into pytorch:main Aug 7, 2026
32 checks passed

theap06 mentioned this pull request Aug 9, 2026

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

Open

Navigation Menu

Uh oh!

[Feature] Add generation-stamped replay slots for round-robin writer - #4046

[Feature] Add generation-stamped replay slots for round-robin writer#4046
vmoens merged 13 commits into
pytorch:mainfrom
harryfrzz:feature/rb-generation-tracking

harryfrzz commented Jul 24, 2026 •

edited by vmoens

Loading

Uh oh!

pytorch-bot Bot commented Jul 24, 2026 •

edited

Loading

Uh oh!

meta-cla Bot commented Jul 24, 2026

Uh oh!

meta-cla Bot commented Jul 24, 2026

Uh oh!

theap06 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

theap06 left a comment

Uh oh!

harryfrzz commented Jul 26, 2026

Uh oh!

harryfrzz commented Jul 29, 2026

Uh oh!

vmoens left a comment

Uh oh!

harryfrzz commented Aug 5, 2026

Uh oh!

vmoens commented Aug 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

harryfrzz commented Jul 24, 2026 • edited by vmoens Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Docs

Not in this PR

Tested

Uh oh!

pytorch-bot Bot commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

Uh oh!

meta-cla Bot commented Jul 24, 2026

Action Required

Process

Uh oh!

meta-cla Bot commented Jul 24, 2026

Uh oh!

theap06 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

theap06 left a comment

Choose a reason for hiding this comment

Uh oh!

harryfrzz commented Jul 26, 2026

Uh oh!

harryfrzz commented Jul 29, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

harryfrzz commented Aug 5, 2026

Uh oh!

vmoens commented Aug 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

harryfrzz commented Jul 24, 2026 •

edited by vmoens

Loading

pytorch-bot Bot commented Jul 24, 2026 •

edited

Loading