[Feature] Add generation-stamped replay buffer slots

[theap06]

Description

First step of the conditional replay-update RFC (#4040), implementing the dependency described in #4041: replay slots get a reuse counter, and samples report which version of a slot they came from.

Round-robin writers now stamp each storage slot with an int64 generation: the first write of a slot has generation 0 and every round-robin reuse increments it, so a captured (index, generation) pair identifies one specific record and becomes detectably stale once the slot is recycled or the buffer is emptied. Nothing acts on the stamps yet; the conditional write-back API that consumes them is submitted separately.

• writer.generations_of(index) is the one new public symbol (int64, dim-0 reduced for multidimensional storages). Writers that do not track reuse (ImmutableDatasetWriter, TensorDictMaxValueWriter, ensembles) report -1 and are exempt from the sampling surface via a tracks_generations flag.
• sample(return_info=True) gains an "index_generation" info entry aligned with "index", read under the same lock acquisition that produced the sample so the handle is coherent; TensorDict buffers surface it as a sample key through the existing info-copy path.
• empty() invalidates all outstanding handles (generations bump, never reset, and stay monotonic across the buffer lifetime).
• Persistence: generations travel through state_dict/load_state_dict and dumps/loads (memmap sidecar next to metadata.json); checkpoints from before this change still load.
• Lazily-initialized storages report a sentinel capacity (2**63-1); the generation tensor detects this and grows dynamically from the highest written slot instead of allocating by capacity.

Deliberate semantics, stated rather than implied:

• write_at and raw __setitem__/set_at_ patch records in place and do not advance generations (positional patching preserves identity; a follow-up could stamp __setitem__ if desired).
• A single extend that wraps the storage and writes a slot twice advances its generation once.
• Cross-process visibility of the stamps for share()-d buffers is deferred to a follow-up.

Testing

TestSlotGenerations (writer level: counting, add/extend wraparound, empty semantics, state_dict/dumps round-trips, legacy checkpoints) and TestSampleGenerationInfo (info entry, TensorDict sample key, staleness detectable after wraparound) — written spec-first before the implementation. Four existing tests that assert over all sample keys were adapted to exclude the new key. Full test/rb/ suite passes (4205 tests).

A wraparound-heavy extend benchmark (test_rb_extend_generation_stamping) covers the hot-path cost: the stamp is a single vectorized int64 increment per write batch.

Addresses #4041.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.