Skip to content

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

Draft
harryfrzz wants to merge 9 commits into
pytorch:mainfrom
harryfrzz:feature/rb-generation-tracking
Draft

[Feature] Add generation-stamped replay slots for round-robin writer#4046
harryfrzz wants to merge 9 commits into
pytorch:mainfrom
harryfrzz:feature/rb-generation-tracking

Conversation

@harryfrzz

Copy link
Copy Markdown

Description

Adds generation-stamped slots to the round-robin replay writer (step 1 of the
work in #4041). Every physical storage slot carries an int64 generation counter:
the first write of a slot has generation 0 and every reuse increments it, so a
sampled (index, generation) pair identifies a specific record and becomes
detectably stale once its slot is reused.

  • Generation tracking is on by default for round-robin writers, exposed via
    writer.generations_of(index) and the tracks_generations capability flag.
    Writers that do not track reuse report -1.
  • sample(return_info=True) returns an "index_generation" entry aligned with
    "index", read under the sampling lock; TensorDict buffers surface it as a
    sample key.
  • empty() is monotonic: previously handed-out (index, generation) handles are
    never revived.
  • Generations survive state_dict/load_state_dict and dumps/loads; older
    checkpoints still load. Lazy (unbounded) storages grow the counter dynamically.

The conditional priority-update API that consumes these stamps is submitted
separately. Cross-process visibility of stamps for share()-d buffers is not
yet supported and is deferred to a follow-up.

Motivation and Context

With a round-robin writer a slot can be overwritten while a learner computes a
loss, so a delayed priority update may apply a TD-error to a different record
that now occupies the same slot. Stamping each index with a generation lets
callers detect that reuse. This PR adds the generation-tracking and sampling
side.

Part of #4041.

  • I have raised an issue to propose this change

Types of changes

  • New feature (non-breaking change which adds core functionality)

Checklist

  • I have read the CONTRIBUTION guide (required)
  • My change requires a change to the documentation.
  • I have updated the tests accordingly (required for a bug fix or a new feature).
  • I have updated the documentation accordingly.

@pytorch-bot

pytorch-bot Bot commented Jul 24, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 38 New Failures, 19 Unrelated Failures, 1 Unclassified Failure

As of commit 88ed9b4 with merge base ae421b9 (image):

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

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

@meta-cla

meta-cla Bot commented Jul 24, 2026

Copy link
Copy Markdown

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

meta-cla Bot commented Jul 24, 2026

Copy link
Copy Markdown

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

@meta-cla 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 github-actions Bot added the Feature New feature label Jul 24, 2026

@theap06 theap06 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Left some feedback.

Comment thread torchrl/data/replay_buffers/writers.py
Comment thread torchrl/data/replay_buffers/writers.py Outdated
@harryfrzz
harryfrzz requested a review from theap06 July 25, 2026 20:36
Comment thread torchrl/data/replay_buffers/writers.py
Comment thread torchrl/data/replay_buffers/writers.py

@theap06 theap06 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left some comments. Otherwise LGTM! cc @vmoens

@harryfrzz

Copy link
Copy Markdown
Author

Left some comments. Otherwise LGTM! cc @vmoens

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

@harryfrzz

Copy link
Copy Markdown
Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Feature New feature ReplayBuffers

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants