[Feature] Extend Sequence with burn-in, bootstrap, stride and per-anchor priority semantics - #4051
Open
theap06 wants to merge 7 commits into
Open
[Feature] Extend Sequence with burn-in, bootstrap, stride and per-anchor priority semantics#4051theap06 wants to merge 7 commits into
theap06 wants to merge 7 commits into
Conversation
Introduces the interface layer of the sampler-decomposition RFC: sampling combines an anchor distribution (the sampler) with a range expansion (the sample unit), and this PR adds the composition point without changing any behavior. A SampleUnit receives the anchor index, the sampler info dict and the storage inside the buffer's sampling critical section, after the anchor sampler ran and before the storage read or index bookkeeping, and returns the expanded index plus (possibly augmented) info. Unit metadata added to info flows into sample(return_info=True) and becomes keys of TensorDict samples through the existing info-copy path. Transition is the identity unit and the implicit default: sample_unit=None and sample_unit=Transition() are behaviorally identical, which the tests pin under seeded generators. The sample_unit keyword is exposed on ReplayBuffer and threaded through the prioritized variants; Hydra config companions gain the matching field. This settles the RFC's first open question (units are buffer-owned) and gives the Sequence unit, boundary policies and priority-semantics follow-ups a stable target. Part of pytorch#4039.
Executable contract for piece 2 of the pytorch#4039 split: Sequence(length, episode_boundary, done_key) expands each anchor into the length records that follow it in stored-time order, wrapping ring indices across the storage seam. Boundary policies: pad keeps the anchor and marks the tail past an episode end invalid with indices clamped inside the episode; stop shifts the anchor backward to end exactly at the boundary, falling back to pad for episodes shorter than length; include_reset crosses the boundary with all entries valid. The unit adds per-record sequence_id, step_in_sequence and validity_mask info entries that surface as TensorDict sample keys, expands per-anchor sampler entries such as prioritized weights to the record count, and sample(batch_size=B) returns B*length records. Constructor validates length and the boundary policy. Tests are expected to fail until the implementation lands.
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.
🔗 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. ❌ 37 New Failures, 19 Unrelated Failures, 1 Unclassified FailureAs of commit 9dbc9bb with merge base ae421b9 ( 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
Bot
added
the
CLA Signed
github-actions
Bot
added
Feature
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Pieces 3 and 4 (first half) of the split proposed in #4039, built on top of the Sequence unit from #4050.
Burn-in, bootstrap, stride (piece 3): the window around each anchor becomes
burn_inrecords before the anchor, the learning region oflengthrecords starting at it, andbootstraprecords after it, withstridespacing the whole window uniformly. A per-recordlearning_maskinfo entry is True exactly on the learning region (recurrent burn-in states get replayed but masked out of the loss). Burn-in never shifts the anchor: entries before the episode start are invalid and clamp to it. Bootstrap entries obey theepisode_boundarypolicy at episode ends. Defaults reproduce the existing Sequence behavior exactly, pinned by a test comparing the two.Priority semantics + distribution invariance (piece 4, first half): priorities live per anchor. A per-record
anchor_indexinfo 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 ordinaryupdate_prioritypath. Per-anchor sampler entries such as importance weights expand block-constant across the window. Seeded statistical tests pin the RFC acceptance criterion that range expansion does not change the anchor selection distribution, for uniform and prioritized sampling, and that a priority update throughanchor_indexsteers subsequent sampling.The
SliceSamplercompatibility adapters and the structured[B, T]layout remain the second half of piece 4 and will follow separately.Testing
Spec-first:
TestSequenceBurnInBootstrap(8 tests) andTestSequencePrioritySemantics(5 tests) were written as the contract before the implementation, with deterministic direct-expandcalls for every boundary case and wide-band seeded statistical tests for the distribution properties. The piece-2 spec from #4050 stays green (10/10), and one of its tests was hardened against a legitimate sampling flake (anchors near the storage cursor produce valid-prefix sequences; the test now asserts the prefix property instead of full validity). Fulltest/rb/suite passes (4212 tests).Note
Stacked on #4050 (@coder-jayp's Sequence unit): its commits are included here and only the last three commits are new. Will rebase once #4050 lands.
Part of #4039.