[Algorithm] Transitive RL training script on OGBench

.github/unittest/linux_sota/scripts/test_sota.py

@@ -30,6 +30,14 @@
   network.width=32 network.depth=2 network.num_steps=3 \
   evaluation.interval=10 evaluation.episodes=1 log_interval=1 \
   hydra.run.dir=$SOTA_LOG_DIR/fql
+""",
+    "transitive_rl": """python sota-implementations/transitive_rl/transitive_rl.py \
+  device=cpu dataset.name=null dataset.random_frames=64 env.max_episode_steps=16 \
+  optim.steps=10 optim.batch_size=8 'network.hidden=[32,32]' \
+  evaluation.interval=10 evaluation.num_tasks=1 evaluation.episodes=1 \
+  evaluation.max_steps=5 log_interval=1 \
+  compile.compile=false compile.cudagraphs=false \
+  hydra.run.dir=$SOTA_LOG_DIR/transitive_rl
 """,
     "tdmpc2": """python sota-implementations/tdmpc2_trainer/train.py \
   collector.total_frames=48 \
@@ -564,10 +572,10 @@ def test_commands(algo, monkeypatch, tmp_path):
     if dataset_id is not None:
         monkeypatch.setenv("HOME", str(tmp_path))
         _write_synthetic_d4rl_dataset(tmp_path, dataset_id)
-    if algo in {"ppo_mujoco", "sac", "fql"}:
+    if algo in {"ppo_mujoco", "sac", "fql", "transitive_rl"}:
         monkeypatch.setenv("SOTA_LOG_DIR", str(tmp_path))
     run_command(commands[algo])
-    if algo in {"ppo_mujoco", "sac", "fql"}:
+    if algo in {"ppo_mujoco", "sac", "fql", "transitive_rl"}:
         scalar_roots = list(tmp_path.rglob("scalars"))
         assert len(scalar_roots) == 1
         scalar_names = {
@@ -578,9 +586,13 @@ def test_commands(algo, monkeypatch, tmp_path):
             "ppo_mujoco": "training/loss_objective.csv",
             "sac": "training/q_loss.csv",
             "fql": "training/loss_actor.csv",
+            "transitive_rl": "training/loss_qvalue.csv",
         }[algo]
         assert expected_training_metric in scalar_names
-        expected_evaluation_metric = "return" if algo == "fql" else "reward"
+        expected_evaluation_metric = {
+            "fql": "return",
+            "transitive_rl": "overall",
+        }.get(algo, "reward")
         assert f"evaluation/{expected_evaluation_metric}.csv" in scalar_names
         assert scalar_names
         assert all(

docs/source/reference/objectives_offline.rst

@@ -50,4 +50,5 @@ Transitive RL
 :class:`TransitiveQLoss` implements Transitive RL (Park et al., ICLR 2026,
 https://arxiv.org/abs/2510.22512), a divide-and-conquer critic for offline
 goal-conditioned RL. It reads the flat batches of a replay buffer sampling with
-:class:`~torchrl.data.replay_buffers.SubgoalTriplet`.
+:class:`~torchrl.data.replay_buffers.SubgoalTriplet`;
+``sota-implementations/transitive_rl/`` trains it on OGBench.

sota-check/README.md

@@ -26,7 +26,7 @@ export MUJOCO_GL=egl
 conda create -n rl-sota-bench python=3.10 -y 
 conda install anaconda::libglu -y
 pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
-pip3 install "gymnasium[atari,mujoco]" vmas tqdm wandb pygame "moviepy<2.0.0" imageio submitit "hydra-core>=1.3,<1.4" transformers datasets
+pip3 install "gymnasium[atari,mujoco]" vmas tqdm wandb pygame "moviepy<2.0.0" imageio submitit "hydra-core>=1.3,<1.4" transformers datasets ogbench
 
 cd /path/to/tensordict
 python setup.py develop

sota-check/run_transitive_rl.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+#SBATCH --job-name=transitive_rl
+#SBATCH --ntasks=32
+#SBATCH --cpus-per-task=1
+#SBATCH --gres=gpu:1
+#SBATCH --output=slurm_logs/transitive_rl_%j.txt
+#SBATCH --error=slurm_errors/transitive_rl_%j.txt
+
+current_commit=$(git rev-parse --short HEAD)
+project_name="torchrl-example-check-$current_commit"
+group_name="transitive_rl"
+export PYTHONPATH=$(dirname $(dirname $PWD))
+python $PYTHONPATH/sota-implementations/transitive_rl/transitive_rl.py \
+  logger.backend=wandb \
+  logger.project="$project_name" \
+  logger.group="$group_name" \
+  compile.compile=true \
+  compile.cudagraphs=true
+
+exit_status=$?
+if [ $exit_status -eq 0 ]; then
+  echo "${group_name}_${SLURM_JOB_ID}=success" >> report.log
+else
+  echo "${group_name}_${SLURM_JOB_ID}=error" >> report.log
+fi

sota-check/submitit-release-check.sh

@@ -70,6 +70,7 @@ scripts=(
     run_td3.sh
     run_td3bc.sh
     run_tdmpc2.sh
+    run_transitive_rl.sh
     run_dt.sh
     run_dt_online.sh
     run_vla_grpo.sh

sota-implementations/transitive_rl/config.yaml

@@ -0,0 +1,55 @@
+seed: 0
+device: cuda
+matmul_precision: high
+log_interval: 10000
+
+logger:
+  backend: csv
+  project: torchrl-transitive-rl
+  group: null
+  mode: online
+
+env:
+  name: Pendulum-v1
+  max_episode_steps: 200
+
+dataset:
+  name: scene-play-oraclerep-v0
+  random_frames: 1000
+
+sampling:
+  geometric_goal: true
+  geometric_policy_goal: false
+
+network:
+  hidden: [512, 512, 512]
+
+loss:
+  discount: 0.99
+  expectile: 0.7
+  distance_weight: 1.0
+  alpha: 1.0
+  tau: 0.005
+  oracle_distill: true
+
+optim:
+  steps: 1000000
+  batch_size: 1024
+  lr: 0.0003
+
+compile:
+  compile: true
+  compile_mode:
+  cudagraphs: true
+
+evaluation:
+  interval: 100000
+  num_tasks:
+  episodes: 15
+  max_steps:
+
+hydra:
+  job:
+    chdir: false
+  run:
+    dir: outputs/transitive_rl/${dataset.name}/seed${seed}_${now:%H-%M-%S}

sota-implementations/transitive_rl/transitive_rl.py

@@ -0,0 +1,127 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""Transitive RL (Park et al., 2026) on OGBench goal-conditioned datasets.
+
+Reference: https://arxiv.org/abs/2510.22512
+
+``config.yaml`` holds the scene-play settings. The reference implementation
+(github.com/aoberai/trl, ``hyperparameters.sh``) tunes three values per
+dataset; pass them as overrides, e.g. ``dataset.name=cube-single-play-oraclerep-v0
+loss.distance_weight=0.7``:
+
+=============================  =========================  ==============  =================
+dataset (``-oraclerep-v0``)    ``loss.distance_weight``   ``loss.alpha``  ``loss.discount``
+=============================  =========================  ==============  =================
+scene-play                     1.0                        1               0.99
+cube-single-play               0.7                        1               0.99
+cube-double-play               1.0                        10              0.99
+puzzle-3x3-play                0.5                        2               0.99
+puzzle-4x4-play                2.0                        2               0.99
+antmaze-large-navigate         0.0                        0.7             0.99
+antsoccer-arena-navigate       0.5                        0.3             0.99
+pointmaze-medium-navigate      0.7                        10              0.99
+humanoidmaze-medium-navigate   0.0                        0.1             0.995
+humanoidmaze-large-navigate    0.1                        0.1             0.995
+=============================  =========================  ==============  =================
+
+The reference runs puzzle-4x5, puzzle-4x6 and humanoidmaze-giant with larger
+networks and a different goal sampling, which this script does not set up.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import hydra
+import numpy as np
+import torch
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import DictConfig, OmegaConf
+from tensordict.nn import CudaGraphModule
+
+from torchrl._utils import compile_with_warmup, logger as torchrl_logger, timeit
+from torchrl.record.loggers import get_logger
+from utils import evaluate, make_agent, make_replay_buffer
+
+
+@hydra.main(config_path=".", config_name="config", version_base="1.3")
+def main(cfg: DictConfig) -> None:
+    torch.set_float32_matmul_precision(cfg.matmul_precision)
+    torch.manual_seed(cfg.seed)
+    np.random.seed(cfg.seed)
+    device = torch.device(cfg.device)
+    replay, eval_env, goal_dim = make_replay_buffer(cfg, device)
+    eval_env.set_seed(cfg.seed)
+    policy, loss_module, updater = make_agent(cfg, replay, goal_dim, device)
+    cudagraphs = cfg.compile.cudagraphs and device.type == "cuda"
+    optimizer = torch.optim.Adam(
+        loss_module.parameters(), lr=cfg.optim.lr, capturable=cudagraphs
+    )
+
+    def update(batch):
+        losses = loss_module(batch)
+        total = sum(v for k, v in losses.items() if k.startswith("loss_"))
+        optimizer.zero_grad(set_to_none=True)
+        total.backward()
+        optimizer.step()
+        updater.step()
+        return losses.detach()
+
+    if cfg.compile.compile:
+        compile_mode = cfg.compile.compile_mode or (
+            "default" if cudagraphs else "reduce-overhead"
+        )
+        update = compile_with_warmup(update, mode=compile_mode, warmup=1)
+    if cudagraphs:
+        update = CudaGraphModule(update, in_keys=[], out_keys=[], warmup=5)
+    output_dir = Path(HydraConfig.get().runtime.output_dir)
+    logger = get_logger(
+        logger_type=cfg.logger.backend,
+        logger_name=str(output_dir),
+        experiment_name="transitive_rl",
+        wandb_kwargs={
+            "project": cfg.logger.project,
+            "group": cfg.logger.group,
+            "mode": cfg.logger.mode,
+            "config": OmegaConf.to_container(cfg, resolve=True),
+        },
+    )
+
+    for step in range(1, cfg.optim.steps + 1):
+        with timeit("train/sample"):
+            batch = replay.sample()
+        with timeit("train/update"):
+            losses = update(batch)
+        if step % cfg.log_interval == 0:
+            metrics = losses.apply(lambda x: x.mean()).to_dict()
+            metrics = {f"training/{k}": float(v) for k, v in metrics.items()}
+            timings = timeit.todict(percall=True)
+            metrics.update(
+                {
+                    "timing/step_s": timings["train/sample"] + timings["train/update"],
+                    "timing/sample_s": timings["train/sample"],
+                    "timing/update_s": timings["train/update"],
+                }
+            )
+            timeit.erase()
+            if logger is not None:
+                logger.log_metrics(metrics, step=step)
+            torchrl_logger.info(
+                f"step {step} loss_qvalue {metrics['training/loss_qvalue']:.4f} "
+                f"q_mean {metrics['training/q_mean']:.3f} "
+                f"{1 / metrics['timing/step_s']:.1f} it/s"
+            )
+        if step % cfg.evaluation.interval == 0:
+            with timeit("eval"):
+                results = evaluate(policy, eval_env, cfg)
+            if logger is not None:
+                logger.log_metrics({k: float(v) for k, v in results.items()}, step=step)
+            torchrl_logger.info(
+                f"step {step} eval overall {float(results['evaluation/overall']):.3f}"
+            )
+    eval_env.close()
+
+
+if __name__ == "__main__":
+    main()