Repository navigation
[BugFix] Fix slice sampler, CQL Lagrange and multi-agent export correctness bugs - #4700
Merged
Merged
Conversation
PrioritizedSliceSampler accepted init_key and documented that it marks the first step of every slice, but its sample method built the truncated, done and terminated entries itself and never wrote init_key. Recurrent learners fed by a prioritized slice buffer therefore did not restart their hidden state at slice starts. The sampler now finalizes its index through SliceSampler._finalize_index, so it writes the same truncated, done, terminated and is_init entries as SliceSampler. The wrapped-buffer fuzz test now stores is_init and checks that every slice starts with it and has no other episode start.
With span=True or span=(..., True), SliceSampler drew relative starts in [start_point, length], so a start could land one step past the end of its trajectory. That slice had zero length: it was silently dropped, returning fewer slices than requested, or made the truncated marker index out of range. The exclusive upper bound is now the trajectory length, so the latest start is the last step and every slice keeps at least one step.
PrioritizedSliceSampler zeroes the priority of the last slice_len - 1 - span_right steps of every trajectory with pad[:, -num_forbidden:]. When num_forbidden is 0 (slice_len=1, or a right span of slice_len - 1) this slice became pad[:, 0:] and forbade every step. With strict_length=True sampling raised "non-positive p_sum"; with strict_length=False only the first step of each trajectory could be drawn. The slice now starts at width - num_forbidden, clamped at 0, so a non-positive count forbids nothing.
With with_lagrange=True, CQLLoss departed from the reference implementation (aviralkumar2907/CQL) in three ways: - loss_alpha_prime was built from the undetached conservative terms, so backpropagating it also pushed gradients into the Q-network parameters. - loss_cql was not scaled by alpha_prime. The reference trains the Q-networks on alpha_prime * (c - lagrange_thresh) with alpha_prime held fixed, and trains alpha_prime only through -alpha_prime * ((c_1 + c_2) / 2 - lagrange_thresh). - reduction="none" raised a batch-dimension mismatch because loss_alpha_prime was always reduced to a scalar. loss_cql is now alpha_prime.detach() * (c_1 + c_2 - 2 * lagrange_thresh), and loss_alpha_prime uses detached conservative terms and follows the loss reduction, so it is per-sample under reduction="none". Summing all losses into one backward pass, as the sota-implementations CQL scripts do, now gives each parameter group the reference gradient. This changes loss_cql values for with_lagrange=True users (alpha_prime scaling and the threshold offset) and the gradients both losses produce. Users without the Lagrange multiplier are not affected.
…strict export Exporting a MultiAgentMLP or MultiAgentConvNet built with share_params=True through torch.export.export(strict=False), the default, corrupted the eager network. In a non-strict trace tensordict's to_module takes its compile path, and its restore writes the meta placeholders of the template network into _parameters, while export puts the same placeholders back in the module __dict__. Afterwards the state_dict carries extra _empty_net.* meta entries, so a checkpoint saved after export no longer loads strictly into a fresh network, and eager forward reads the placeholders: MultiAgentMLP raises a meta device error and MultiAgentConvNet silently returns garbage, so eager and exported outputs disagree. During a non-strict export the shared branch now calls the template with torch.func.functional_call, with the placeholders taken out of the template for the call and put back afterwards. Eager, compiled and strict-export runs keep using to_module, and the parameter layout is unchanged.
theap06
added
the
ci/olddeps
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4700
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 2d52da0 with merge base c5a5e3a ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
PrioritizedSliceSampler with truncated_key=None now finishes through SliceSampler._finalize_index, whose no-truncation branch gathered every key of the sample just to read is_init. On 84x84 pixels that made rb.sample 2.7x slower than main. The marker now reads the single init_key entry from the underlying TensorDict and skips the read when init_key is None. Adds a PrioritizedSliceSampler case to the replay buffer sample benchmark.
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
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
Fixes silent wrong results in the replay samplers,
CQLLossand multi-agent export. Each fix has a regression test that fails on main.PrioritizedSliceSamplernow marks slice starts withis_init, likeSliceSampler, so recurrent policies get correct sequence boundaries. Adds a PSS case to the replay buffer benchmarks.SliceSamplerwith a right span no longer produces zero-length slices orIndexError.PrioritizedSliceSamplerwithslice_len=1can start on every step (it raised or only drew first steps).CQLLoss(with_lagrange=True)matches the reference CQL(H) Lagrange update:loss_alpha_primeno longer sends gradients into the Q networks,loss_cqlis scaled byalpha_prime, andreduction="none"works.MultiAgentMLP/MultiAgentConvNetwithshare_params=Trueare no longer corrupted by a non-stricttorch.export.Breaking changes
with_lagrange=True(used by both SOTA CQL configs),loss_cqlvalues and gradients change, so curves are not comparable with earlier runs.is_init, sampledPrioritizedSliceSamplerbatches now set it at every slice start.Motivation and Context
These produce wrong training inputs, wrong gradients or a broken model without an error.
Types of changes
Checklist