[Algorithm] Add Transitive RL

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

@@ -30,6 +30,13 @@
   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 \
+  hydra.run.dir=$SOTA_LOG_DIR/transitive_rl
 """,
     "tdmpc2": """python sota-implementations/tdmpc2_trainer/train.py \
   collector.total_frames=48 \
@@ -564,10 +571,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 +585,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/config.rst

@@ -478,6 +478,7 @@ Replay Buffer and Storage Configurations
     SliceSamplerConfig
     SliceSamplerWithoutReplacementConfig
     StreamingSliceSamplerConfig
+    SubgoalTripletConfig
     ListStorageConfig
     TensorStorageConfig
     LazyTensorStorageConfig
@@ -549,6 +550,7 @@ Trainer Hook Configurations
     CQLLossConfig
     TD3LossConfig
     TdMpc2LossConfig
+    TransitiveQLossConfig
     GAEConfig
     TargetNetUpdaterConfig
     SoftUpdateConfig

docs/source/reference/data_replaybuffers.rst

@@ -204,6 +204,16 @@ With ``B`` sampled anchors, the flat output contains
     sample, info = rb.sample(return_info=True)
     loss_mask = info["learning_mask"] & info["validity_mask"]
 
+:class:`~torchrl.data.replay_buffers.SubgoalTriplet` expands each anchor ``i``
+into four records of its own episode: the anchor, a subgoal ``k`` in
+``[i, j)``, a goal ``j > i`` and an independent policy goal. Goal offsets are
+geometric (or uniform) and truncated at the episode end, and the
+``"step_offset"`` info entry carries each record's distance to its anchor, so
+that a divide-and-conquer critic such as
+:class:`~torchrl.objectives.TransitiveQLoss` can use exact
+:math:`\gamma^{k - i}` values for one-step horizons. With ``B`` sampled
+anchors, the flat output contains ``4 * B`` records.
+
 Conditional record updates
 --------------------------
 
@@ -313,6 +323,7 @@ counting only generation-stale handles.
 
     SampleUnit
     Sequence
+    SubgoalTriplet
     Transition
     ConditionalUpdateResult
 

docs/source/reference/objectives_offline.rst

@@ -15,6 +15,7 @@ Loss modules for offline reinforcement learning.
     DiscreteIQLLoss
     TD3BCLoss
     FQLLoss
+    TransitiveQLoss
 
 Flow Q-learning
 ---------------
@@ -42,3 +43,12 @@ encoder receives gradients only from its own objective.
 
 Inference passes the one-step policy directly to TorchRL collectors.
 Exploration remains stochastic at evaluation, as in the reference implementation.
+
+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`;
+``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: false
+  compile_mode:
+  cudagraphs: false
+
+evaluation:
+  interval: 100000
+  num_tasks: 5
+  episodes: 15
+  max_steps: 1000
+
+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,101 @@
+# 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
+"""
+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)
+    optimizer = torch.optim.Adam(
+        loss_module.parameters(), lr=cfg.optim.lr, capturable=cfg.compile.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 cfg.compile.cudagraphs else "reduce-overhead"
+        )
+        update = compile_with_warmup(update, mode=compile_mode, warmup=1)
+    if cfg.compile.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()
+            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)
+            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()