Repository navigation
Conversation
…iveQLoss Transitive RL (Park, Oberai, Atreya, Levine, ICLR 2026, https://arxiv.org/abs/2510.22512) learns a goal-conditioned critic Q(s, a, g) = gamma^d(s, g) with a divide-and-conquer update: Q(s_i, a_i, s_j) is regressed toward Q(s_i, a_i, s_k) Q(s_k, a_k, s_j) for an in-trajectory subgoal k in [i, j). - SubgoalTriplet sample unit: expands each anchor into its state, subgoal, goal and policy-goal records, with geometric or uniform goal offsets truncated at the episode end and a per-record "step_offset" info entry. - TransitiveQLoss: expectile-weighted BCE critic with exact gamma^offset base cases, optional distance re-weighting, optional oracle-representation critic distillation, and a reparameterized DDPG+BC actor; anchors flagged by validity_mask are excluded from every term. - SubgoalTripletConfig and TransitiveQLossConfig Hydra configs, docs entries. - sota-implementations/transitive_rl recipe on OGBench (compile + CUDA graphs optional), test_sota smoke entry on synthetic data, sota-check script.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4579
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 81bf4cb with merge base 264a3c2 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
github-actions
Bot
added
new algo
theap06
changed the title
…unit helpers - TransitiveQLoss reduces through LossModule._reduce_loss with loss_mask_key defaulting to "validity_mask", replacing the custom masked mean and the validity_mask accepted key; masked anchors are now selected out rather than multiplied by zero. Critic inputs are built with TensorDict.select instead of a hand-assembled TensorDict. - SampleUnit gains _check_storage, _repeat_info and _episode_boundaries, shared by Sequence and SubgoalTriplet instead of duplicated in both. - Shorten the objectives_offline section, which repeated the class docstring. Losses and gradients still match the JAX reference (worst relative error 1e-7), and a compiled + CUDA-graph OGBench run logs identical losses.
…recipe copy - TransitiveQLossConfig sets _convert_="object" so a list of critic configs reaches TransitiveQLoss as a list rather than a ListConfig. - TensorDictReplayBuffer reduces per-record priorities to anchors for SubgoalTriplet samples too, as it already did for Sequence; priorities were otherwise written onto subgoal and goal records. - Document that TransitiveQLoss assumes uniform anchor sampling. - The recipe extends the device storage from the host dataset directly instead of copying the dataset to the device first. - Add ogbench to the sota-check setup. Tests: the ordering test now also runs on a wrapped buffer whose newest episode is still open (it does not catch a write-seam bug otherwise), the loss config test instantiates a list of nested critic configs, and a new test checks that a priority update lands on the anchor.
…stronger mask test - SubgoalTriplet draws its offsets from the replay buffer's generator (set through generator= or set_rng) when it lives on the anchors' device, so seeding the buffer reproduces its batches like the other sample units. - The recipe rejects loss.oracle_distill=false on datasets whose goals are oracle representations, which previously failed with a shape mismatch. - The masking test now feeds a real flat SubgoalTriplet sample with single-record episodes and compares against the valid anchors indexed out independently; it replaces the flat-versus-unflattened test, which compared the loss's own reshape with itself. - q_mean and target_mean exclude invalid anchors, oracle parameters are cleared for compile like the others, the actor_network docstring type matches its hint, a stale replay-buffer comment is corrected, and the test_sota override is quoted.
theap06
marked this pull request as ready for review
October 16, 2026 08:52
This branch has not been deployed
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
1 participant
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
Adds Transitive RL (Park, Oberai, Atreya, Levine, "Transitive RL: Value Learning via Divide and Conquer", ICLR 2026, https://arxiv.org/abs/2510.22512) for offline goal-conditioned RL. The critic
Q(s, a, g)lives in(0, 1]and is read asgamma^d(s, g); instead of a one-step backup,Q(s_i, a_i, s_j)is regressed towardQ(s_i, a_i, s_k) Q(s_k, a_k, s_j)through a subgoalkthe dataset trajectory visited betweeniandj.Validation against the reference implementation
Reference: the authors' JAX code (github.com/aoberai/trl @ adfec6f249), same hyperparameters (512x3 MLPs, batch 1024, discount 0.99, expectile 0.7, oracle distillation, rpg actor with alpha 1).
Numerical parity. On a fixed batch with identical weights, in float64: critic, oracle and actor losses match to 1e-16 relative error, gradient norms to 1e-7. The compiled + CUDA-graph update matches eager after 60 steps (2e-6 on the loss, 2e-7 on parameters).
Success rates, 1M gradient steps, 2 seeds, OGBench final score (per-seed mean of the 800k, 900k and 1M evaluations, 15 episodes x 5 tasks each):
cube-single-play-oraclerep-v0scene-play-oraclerep-v0Both implementations land within seed spread of the paper, and of each other.
Throughput (A100-SXM4-40GB, scene-play shapes, batch 1024, it/s of sample + update):
compile+CudaGraphModulemax-autotune+CudaGraphModuleEager mode is Python and launch bound at these model sizes; the recipe exposes
compile.compile=true compile.cudagraphs=true.Known differences and follow-ups
Motivation and Context
Transitive RL is a recent offline goal-conditioned method with strong OGBench results whose two building blocks (in-trajectory subgoal sampling and a divide-and-conquer critic) are reusable beyond the paper's recipe.
Types of changes
Checklist