[theap06]

Description

Part 3 of 3 of the Transitive RL stack, which replaces #4579. Based on part 2; review only the top commit.

Stack (merge in order):

1. [Feature] SubgoalTriplet sample unit for in-trajectory subgoal sampling #4596 SubgoalTriplet sample unit
2. [Feature] TransitiveQLoss for offline goal-conditioned RL #4597 TransitiveQLoss
3. [Algorithm] Transitive RL training script on OGBench #4598 OGBench training script and validation (this PR)

Adds sota-implementations/transitive_rl/, which trains TransitiveQLoss (part 2) on OGBench goal-conditioned datasets with SubgoalTriplet sampling (part 1). Transitive RL is from Park, Oberai, Atreya, Levine, "Transitive RL: Value Learning via Divide and Conquer", ICLR 2026, https://arxiv.org/abs/2510.22512.

• dataset.name=null runs a small Pendulum smoke mode, used by the new test_sota.py entry.
• The update is compiled with CUDA graphs by default, the setup behind the results below. CUDA graphs apply only on a CUDA device; the smoke test turns compile off.
• evaluation.num_tasks and evaluation.max_steps default to the env's OGBench evaluation tasks and time limit, as the reference evaluation runs each episode until the env ends it.
• The script's module docstring lists the reference's per-dataset distance_weight, alpha and discount (from its hyperparameters.sh); config.yaml holds the scene-play values.
• logger.backend=null disables logging.
• Adds sota-check/run_transitive_rl.sh, and ogbench to the sota-check setup.

Validation against the reference implementation

Reference: the authors' JAX code (github.com/aoberai/trl @ adfec6f249), with the same hyperparameters: 512x3 MLPs, batch 1024, discount 0.99, expectile 0.7, oracle distillation, rpg actor with alpha 1, and distance weight 1.0 (scene) or 0.7 (cube-single).

Numerical parity. On a fixed batch with identical weights, in float64, the critic, oracle and actor losses match to 1e-16 relative error, and the gradient norms to 2e-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. The score is the OGBench final score: per-seed mean of the 800k, 900k and 1M evaluations, each over 15 episodes x 5 tasks.

Task	TorchRL	Reference (JAX)	Paper
cube-single-play-oraclerep-v0	93.1 (91.1, 95.1)	88.7 (88.9, 88.4)	95
scene-play-oraclerep-v0	79.3 (71.1, 87.6)	75.3 (79.1, 71.6)	77

Both implementations land within seed spread of the paper, and of each other.

The success rates were measured on an earlier revision of this code (#4579 at 1f2ca7a). The later commits are review fixes. On the final code the losses still match the reference (2e-7) and a compiled OGBench run logs identical losses, but I did not retrain. The changes since do not affect these numbers: the recipe samples anchors uniformly (the priority-index fix only matters for prioritized buffers), q_min/q_max are logged metadata, scene and cube-single end episodes at 750 and 200 steps (below the former 1000-step evaluation cap), and the runs were already compiled with CUDA graphs.

Throughput (A100-SXM4-40GB, scene-play shapes, batch 1024, it/s of sample + update):

Implementation	it/s
Reference (JAX jit)	274
TorchRL eager	63
TorchRL compile + CudaGraphModule	630
TorchRL max-autotune + CudaGraphModule	686

At these model sizes, eager mode is bound by Python overhead and kernel launches.

Known differences and follow-ups

• An anchor sampled on the last record of its episode is moved one step back (documented). The reference samples anchors uniformly over non-final records, so the only difference is that penultimate records get double weight.
• Not included, and left as follow-ups:
  • flow-policy rejection-sampling extraction
  • mixing random and geometric policy goals, which the puzzle tasks need
  • Trainer integration
  • batched evaluation (scene evaluation takes about a third of wall-clock time)

Motivation and Context

Transitive RL is a recent offline goal-conditioned method with strong OGBench results. This PR adds the runnable recipe and the reference comparison for the building blocks added in parts 1 and 2.

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• New feature (non-breaking change which adds core functionality)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 2fdccb0 with merge base 53b10df ():
💚 Looks good so far! There are no failures yet. 💚

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