[Feature] Extend Sequence with burn-in, bootstrap, dilation and per-anchor priority semantics - #4051
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4051
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (4 Unrelated Failures)As of commit 600c60d with merge base 678520e ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following jobs failed but was 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. |
34e7999 to
b719832
Compare
|
Thanks for the contribution — I routed priority updates through anchors (max over valid records), kept burn-in inside the ring seam, and documented index/anchor_index semantics in b719832. |
b719832 to
454ad75
Compare
Executable contract for piece 3 of the pytorch#4039 split: Sequence gains burn_in, bootstrap and stride keyword arguments. burn_in records precede the anchor and never shift it (entries before the episode start are invalid and clamped to it); the learning region of length records starts at the anchor; bootstrap records follow it and obey the episode_boundary policy at episode ends; stride spaces the whole window uniformly. A per-record learning_mask info entry is True exactly on the learning region, step_in_sequence spans the window, and defaults reproduce base Sequence behavior. Negative counts or a non-positive stride raise ValueError. Tests are expected to fail until the implementation lands.
…variance Executable contract for the first half of piece 4 of the pytorch#4039 split. Priorities live per anchor: the Sequence unit adds a per-record anchor_index info entry so priorities of sampled sequences can be updated through the ordinary update_priority path, and per-anchor sampler entries such as importance weights are expanded block-constant across each window. Range expansion must not change the anchor selection distribution: seeded statistical tests compare uniform and prioritized anchor frequencies with and without the unit, and an update through anchor_index must steer subsequent sampling. The SliceSampler compatibility adapters and the structured [B, T] layout remain the second half of piece 4 and get their own spec once this lands. Tests are expected to fail until the implementation lands.
…hor priority semantics Pieces 3 and 4 (first half) of the pytorch#4039 split, on top of the Sequence unit from pytorch#4050. The window around each anchor becomes burn_in records before the anchor, the learning region of length records starting at it, and bootstrap records after it, with stride spacing the whole window uniformly. A per-record learning_mask info entry is True exactly on the learning region. Burn-in never shifts the anchor: entries before the episode start are invalid and clamp to it; bootstrap entries obey the episode_boundary policy at episode ends; defaults reproduce the previous behavior exactly. Priorities live per anchor: a per-record anchor_index info entry reports the storage index of each record's sampled anchor (the original anchor, not the stop-shifted one, since that is what the sampler's distribution selected), so priorities of sampled sequences update through the ordinary update_priority path, and per-anchor sampler entries such as importance weights expand block-constant across the window. Seeded distribution tests pin that range expansion does not bias anchor selection for uniform or prioritized sampling. Part of pytorch#4039.
include_reset wrapped indices modulo storage.max_size, so on a storage that has not wrapped yet a burn-in window walking backward from an early anchor (or a forward window crossing the write cursor) landed on never-written slots and read uninitialized data. Crossing the physical seam is only meaningful once the ring has actually wrapped: indices now wrap modulo the written length when the storage is full, and clamp with validity_mask=False when it is not. Regression test covers both directions on a half-full storage.
…n inside the ring - update_tensordict_priority (sync and async paths) now detects a sample-unit expansion through the sample's anchor_index entry: the per-record priorities are reduced with a max over each anchor's valid records and written to the anchor slots only. Padded or bootstrap records no longer pollute the priorities of unrelated storage slots, and duplicate anchors are reduced before writing instead of last-write-wins. - include_reset burn-in now clamps at the oldest record of a full ring buffer instead of wrapping backward across the write seam into the newest data; the forward direction keeps the write-cursor clamp introduced at the Sequence level. - Sequence docstring: document info['index'] (expanded record indices) vs anchor_index semantics, the pre-shift anchor reported under episode_boundary='stop' with bootstrap/stride, and add Examples covering burn_in, bootstrap and stride; keep Literal/NestedKey hints on the extended signature. - SequenceConfig gains burn_in/bootstrap/stride fields (config/class parity). - Tests: anchor-routed update_tensordict_priority (pad pollution and duplicate-anchor reduction, plus an end-to-end sampled batch), full-ring burn-in seam clamp, public import path for the new test classes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
454ad75 to
d254bb3
Compare
|
Thanks @theap06! I think that performance and multidimensional-storage support can be handled in follow-up PRs rather than blocking this one, so I'll leave it as TODOs for you (please let me know if you'd rather have someone else do it). I pushed the missing documentation, including a recurrent Q-learning use case and an explanation of how learning_mask and validity_mask should be consumed. I also renamed stride to dilation: the parameter spaces records inside a window and does not control spacing or overlap between independently sampled windows. The PR title/body should be updated accordingly. For follow-ups, I suggest:
The boundary-detection work should be holistic rather than a local optimization of Sequence (we have that problem in many places in torchrl). Compile can help, it worked nicely with SliceSampler back in the days IIRC. That's a high difficulty task so I'm happy to take it over. |
vmoens
left a comment
There was a problem hiding this comment.
Thanks @theap06!
I think that performance and multidimensional-storage support can be handled in follow-up PRs rather than blocking this one, so I'll leave it as TODOs for you (please let me know if you'd rather have someone else do it).
I pushed the missing documentation, including a recurrent Q-learning use case and an explanation of how learning_mask and validity_mask should be consumed. I also renamed stride to dilation: the parameter spaces records inside a window and does not control spacing or overlap between independently sampled windows. The PR title/body should be updated accordingly.
For follow-ups, I suggest:
- Design shared, efficient replay-boundary detection used by Sequence, SliceSampler, and other trajectory-aware components, including caching/invalidation, ring wraparound, device placement, and benchmarks.
- Add ndim > 1 support by expanding the time coordinate while preserving lane coordinates, including multidimensional anchor-priority reduction.
The boundary-detection work should be holistic rather than a local optimization of Sequence (we have that problem in many places in torchrl). Compile can help, it worked nicely with SliceSampler back in the days IIRC.
IMO that's a high difficulty task.
Description
Extends the
Sequencesample unit from #4050 with recurrent-training context and anchor-based priority semantics.Burn-in, bootstrap, and dilation
Each sampled anchor expands into
burn_in + length + bootstraprecords. The learning region starts at the sampled anchor; burn-in records precede it and bootstrap records follow it.dilationspaces records within the window and does not control spacing or overlap between independently sampled windows.The returned metadata includes:
learning_mask: true only for the learning region;validity_mask: false for padded records;anchor_index: the original sampled anchor repeated across its expanded records;sequence_idandstep_in_sequence: flat-batch grouping metadata.A recurrent learner should replay the full window to reconstruct state and compute loss under
learning_mask & validity_mask.episode_boundary="pad"clamps records at trajectory boundaries and marks padding invalid."stop"shifts the effective window backward when possible."include_reset"may cross episode boundaries but cannot cross the replay write seam or enter unwritten slots. Defaults preserve the previousSequence(length=...)behavior.Priority semantics
Priorities remain attached to sampled anchors.
TensorDictReplayBuffer.update_tensordict_prioritygroups expanded per-record priorities byanchor_index, excludes records masked byvalidity_mask, max-reduces duplicate anchors, and updates only anchor slots. Per-anchor sampler metadata such as importance weights expands blockwise without changing the anchor sampling distribution.Documentation and tests
The reference documentation includes a recurrent Q-learning example and explains how
learning_maskandvalidity_maskare consumed. Tests cover burn-in, bootstrap, dilation, all boundary policies, partial and full rings, wraparound, duplicate anchors, priority routing, distribution invariance, nested done keys, and config parity.The focused Sequence suite passes on CPU and NVIDIA GB10 (
36 passed, 447 deselected). The original full replay-buffer run passed 4,212 tests. CI failures shown on the merged PR were classified as unrelated flaky or broken-trunk failures; GPU, wheel, docs, lint, and targeted jobs passed.Post-merge review and follow-ups
Post-merge profiling used recurrent Q-learning replay with one million records, average episode length 256, 256 anchors,
burn_in=40,length=64,bootstrap=5, anddilation=1.On NVIDIA GB10:
Sequence.expand()The compressed cache used 62,528 bytes versus 16,000,000 bytes for the dense map. The compressed design is the preferred default; a direct sorted-stop path is preferable for one-dimensional storage, with lane-encoded search reserved for multidimensional storage.
The follow-up work is intentionally split into reviewable PRs:
[Refactor] Centralize replay trajectory-boundary queries: introduce one internal provider shared bySequence,SliceSampler, and other trajectory-aware consumers; establish storage mutation revisions, cache lifecycle, ring-wrap behavior, device placement, and same-process/shared-process invalidation.[Performance] Cache and compile replay-boundary queries: add the compressed cache, the specialized one-dimensional stop lookup, the pure compile-friendly tensor kernel, and cold/hot/write-invalidated benchmarks forSequenceandSliceSampler.[Feature] Support multidimensional storage in Sequence sample units: expand only the time coordinate, preserve lane coordinates, return multidimensionalanchor_index, and max-reduce priorities by full anchor coordinate after validity masking.Part of #4039.