[BugFix] Fix checkpointing of replay buffer prefetch state

[theap06]

Summary

Persist the replay buffer prefetch queue across state_dict / load_state_dict, dumps / loads, and Python pickle round-trips.

A restored buffer now returns the batches that had already been prefetched at checkpoint time before producing new samples. This keeps sample order, sampler state, and RNG state aligned, including for SamplerWithoutReplacement.

Root cause

Prefetch workers advance the sampler and RNG when they produce a batch, before that batch is consumed by ReplayBuffer.sample(). Previous checkpoints serialized the already-advanced sampler and RNG but discarded the completed or in-flight queue entries. Restoring such a checkpoint therefore skipped the prefetched batches and resumed from a later sampling state.

Implementation

• Wait for queued futures while holding the futures lock so queue consumption cannot race checkpoint capture.
• Acquire the replay lock only after futures settle, preserving the existing futures -> replay lock order and avoiding deadlock with workers sampling under the replay lock.
• Capture storage, sampler, writer, transform, RNG, and the ordered prefetch queue at one checkpoint boundary.
• Store a versioned prefetch payload containing the configured capacity and completed queue results.
• Recreate saved queue entries as completed Future objects so the existing sampling path can consume them without starting replacement worker pools or special-casing restored batches.
• Cancel and settle a target buffer's current futures before loading checkpoint state.
• Validate prefetch version, capacity, and queue shape before mutating a live target buffer. Invalid or capacity-mismatched state therefore leaves its current queue intact.
• Preserve backward compatibility with state dictionaries and checkpoint directories that predate the prefetch payload; loading them clears any target queue and resumes without restored entries.

Autograd and ownership

Queued samples can contain non-leaf tensors, either because list-backed samples are stacked or because a transform produced autograd-connected values. copy.deepcopy rejects these tensors.

The queue clone now walks the sample PyTree, detaches and clones tensor leaves, restores their requires_grad property, and deep-copies non-tensor leaves. Checkpoints preserve sample values and tensor metadata without retaining a live autograd graph.

Ownership is handled per checkpoint API:

• state_dict() and __getstate__() clone queued results so returned state cannot alias the live queue.
• dumps() serializes the queue directly while checkpoint locks are held, avoiding a redundant full queue-sized pre-copy.
• loads() and __setstate__() transfer ownership of freshly deserialized queue entries into completed futures instead of cloning them again.
• load_state_dict() still clones queue entries because the caller retains ownership of the supplied state dictionary.

The unavoidable checkpoint-size increase is bounded by the configured prefetch capacity: at most prefetch completed, post-transform batches are stored. No additional executor is created for state-dict or disk restoration, and the restored queue remains bounded by that capacity.

Validation

Added or extended coverage for:

• state-dict, disk, and pickle round-trips;
• in-flight checkpoint capture and loading over active worker futures;
• plain tensor and TensorDict samples containing non-leaf autograd tensors;
• queue ownership / non-aliasing;
• legacy checkpoints without prefetch state;
• capacity mismatch without partial queue mutation;
• exact continuation for sampling without replacement.

Local results:

• 14 focused prefetch/checkpoint tests passed with warnings treated as errors.
• Both deterministic in-flight concurrency tests passed 20 consecutive repetitions (40 runs).
• 385 replay-buffer tests passed on the rebased branch, with one CUDA-only skip; prioritized/priority tests requiring the unavailable local C++ segment-tree extension were excluded.
• A concurrent-write stress probe completed 4,010 consistent checkpoint snapshots.
• Repeated pickle/unpickle/use/discard probes returned prefetch worker threads to their baseline after in-flight work settled and garbage collection ran.
• Python 3.10 syntax parsing, ufmt, flake8, pydocstyle, codespell, and git diff --check passed for the changed files.

Current status

Rebased onto main at 7923c6da0. Conflict resolution preserves the current replay-buffer sample-unit and storage mutation-revision changes alongside prefetch checkpointing. The force-push to the contributor branch triggered fresh CI.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 8d23cd0 with merge base 7923c6d ():
💚 Looks good so far! There are no failures yet. 💚

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

[github-actions]

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Fix checkpointing of replay buffer prefetch state

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[Algorithm]	new algo	[Algorithm] Add new RL objective
[BE]	BE	[BE] Improve error messages
[Benchmark] or [Benchmarks]	Benchmarks	[Benchmark] Add collector benchmark
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Example] or [Examples]	Examples	[Example] Add training script
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Trainer] or [Trainers]	Trainers	[Trainer] Add trainer config
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[LLM]	llm/	[LLM] Add reward model integration
[Minor]	small change	[Minor] Fix typo in error message
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Algorithm] or [Algorithms]	new algo	[Algorithm] Add new objective
[Quality]	Quality	[Quality] Fix typos and add codespell
[Versioning]	versioning	[Versioning] Bump release version
[WIP]	WIP	[WIP] Draft implementation

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

[vmoens]

LGTM thanks!