From 64c23de494e11b01d1616df87b3600c256848144 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Tue, 20 Oct 2026 16:41:30 -0700 Subject: [PATCH 1/3] [Algorithm] Transitive RL training script on OGBench sota-implementations/transitive_rl trains TransitiveQLoss on OGBench goal-conditioned datasets with SubgoalTriplet sampling, following the reference implementation's hyperparameters (512x3 MLPs, batch 1024, discount 0.99, expectile 0.7, oracle distillation, rpg actor). dataset.name=null runs a small Pendulum smoke mode used by the SOTA CI entry; compile.compile and compile.cudagraphs enable the compiled update. Logging can be disabled with logger.backend=null. Adds the sota-check job and the test_sota smoke command. --- .../unittest/linux_sota/scripts/test_sota.py | 17 +- docs/source/reference/objectives_offline.rst | 3 +- sota-check/README.md | 2 +- sota-check/run_transitive_rl.sh | 26 ++ sota-check/submitit-release-check.sh | 1 + .../transitive_rl/config.yaml | 55 ++++ .../transitive_rl/transitive_rl.py | 103 +++++++ sota-implementations/transitive_rl/utils.py | 279 ++++++++++++++++++ 8 files changed, 481 insertions(+), 5 deletions(-) create mode 100755 sota-check/run_transitive_rl.sh create mode 100644 sota-implementations/transitive_rl/config.yaml create mode 100644 sota-implementations/transitive_rl/transitive_rl.py create mode 100644 sota-implementations/transitive_rl/utils.py diff --git a/.github/unittest/linux_sota/scripts/test_sota.py b/.github/unittest/linux_sota/scripts/test_sota.py index ae414406e9e..0aea2e74dfa 100644 --- a/.github/unittest/linux_sota/scripts/test_sota.py +++ b/.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( diff --git a/docs/source/reference/objectives_offline.rst b/docs/source/reference/objectives_offline.rst index 5dc20a40e3a..8861fdcc10e 100644 --- a/docs/source/reference/objectives_offline.rst +++ b/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. diff --git a/sota-check/README.md b/sota-check/README.md index 72573b810e9..dc4d1c9e7c7 100644 --- a/sota-check/README.md +++ b/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 diff --git a/sota-check/run_transitive_rl.sh b/sota-check/run_transitive_rl.sh new file mode 100755 index 00000000000..0e563920182 --- /dev/null +++ b/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 diff --git a/sota-check/submitit-release-check.sh b/sota-check/submitit-release-check.sh index a4e91a2c8a0..17b0f082f69 100755 --- a/sota-check/submitit-release-check.sh +++ b/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 diff --git a/sota-implementations/transitive_rl/config.yaml b/sota-implementations/transitive_rl/config.yaml new file mode 100644 index 00000000000..7f0619e2d00 --- /dev/null +++ b/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} diff --git a/sota-implementations/transitive_rl/transitive_rl.py b/sota-implementations/transitive_rl/transitive_rl.py new file mode 100644 index 00000000000..6c1c6341917 --- /dev/null +++ b/sota-implementations/transitive_rl/transitive_rl.py @@ -0,0 +1,103 @@ +# 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() + 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() diff --git a/sota-implementations/transitive_rl/utils.py b/sota-implementations/transitive_rl/utils.py new file mode 100644 index 00000000000..e7262525080 --- /dev/null +++ b/sota-implementations/transitive_rl/utils.py @@ -0,0 +1,279 @@ +# 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. +"""Building blocks of the Transitive RL recipe. + +``make_replay_buffer`` loads an OGBench dataset into a replay buffer sampling +:class:`~torchrl.data.replay_buffers.SubgoalTriplet` records, ``make_agent`` +builds the networks and :class:`~torchrl.objectives.TransitiveQLoss`, and +``evaluate`` runs the OGBench evaluation goals. +""" +from __future__ import annotations + +import importlib.util +import math + +import gymnasium as gym +import numpy as np +import torch +from omegaconf import DictConfig +from tensordict import TensorDict +from tensordict.nn import TensorDictModule +from torch import nn + +from torchrl.data import Bounded, LazyTensorStorage, TensorDictReplayBuffer +from torchrl.data.replay_buffers import SubgoalTriplet +from torchrl.envs import ( + default_info_dict_reader, + DoubleToFloat, + GymWrapper, + TransformedEnv, +) +from torchrl.envs.utils import ExplorationType, set_exploration_type +from torchrl.modules import IndependentNormal, ProbabilisticActor +from torchrl.objectives import SoftUpdate, TransitiveQLoss + +_has_ogbench = importlib.util.find_spec("ogbench") is not None +if _has_ogbench: + import ogbench + + +class SuccessReader(default_info_dict_reader): + """Copy OGBench's ``info["success"]`` flag into the tensordict for evaluation.""" + + def __call__(self, info, tensordict): + return tensordict.set( + "success", torch.tensor(info.get("success", 0.0), dtype=torch.float32) + ) + + +class GoalObservation(gym.Wrapper): + """Expose the OGBench evaluation goal, reported in the reset info, next to a float32 observation. + + Envs that report no goal use their reset observation, which only serves the + synthetic smoke-test setup. + """ + + def __init__(self, env: gym.Env, goal_dim: int): + super().__init__(env) + self.observation_space = gym.spaces.Dict( + { + "observation": gym.spaces.Box( + -np.inf, np.inf, env.observation_space.shape, np.float32 + ), + "goal": gym.spaces.Box(-np.inf, np.inf, (goal_dim,), np.float32), + } + ) + self._goal = np.zeros(goal_dim, np.float32) + + def reset(self, **kwargs): + obs, info = self.env.reset(**kwargs) + self._goal = np.asarray(info.get("goal", obs), np.float32) + return {"observation": obs.astype(np.float32), "goal": self._goal}, info + + def step(self, action): + obs, reward, terminated, truncated, info = self.env.step(action) + return ( + {"observation": obs.astype(np.float32), "goal": self._goal}, + reward, + terminated, + truncated, + info, + ) + + +def load_dataset(cfg: DictConfig) -> tuple[TensorDict, gym.Env]: + """Return the offline data and the raw evaluation env. + + Without ``cfg.dataset.name``, a random rollout of ``cfg.dataset.random_frames`` + steps on ``cfg.env.name`` stands in for the dataset and its observations for + the goals, for smoke tests only. + """ + if cfg.dataset.name: + if not _has_ogbench: + raise ImportError("Install ogbench to load an OGBench dataset.") + raw_env, dataset, _ = ogbench.make_env_and_datasets( + cfg.dataset.name, compact_dataset=True + ) + observation = torch.as_tensor(dataset["observations"], dtype=torch.float32) + data = TensorDict( + { + "observation": observation, + "oracle_goal": torch.as_tensor( + dataset.get("oracle_reps", dataset["observations"]), + dtype=torch.float32, + ), + "action": torch.as_tensor(dataset["actions"], dtype=torch.float32), + ("next", "done"): torch.as_tensor( + dataset["terminals"], dtype=torch.bool + ).unsqueeze(-1), + }, + batch_size=[observation.shape[0]], + ) + else: + raw_env = gym.make(cfg.env.name, max_episode_steps=cfg.env.max_episode_steps) + source = TransformedEnv( + GymWrapper( + gym.make(cfg.env.name, max_episode_steps=cfg.env.max_episode_steps) + ), + DoubleToFloat(), + ) + source.set_seed(cfg.seed) + rollout = source.rollout(cfg.dataset.random_frames, break_when_any_done=False) + data = rollout.select("observation", "action", ("next", "done")) + data["oracle_goal"] = data["observation"] + data["next", "done"][-1] = True + source.close() + data["action"] = data["action"].clamp(-1 + 1e-5, 1 - 1e-5) + return data, raw_env + + +def make_replay_buffer(cfg: DictConfig, device: torch.device): + """Return the replay buffer, the evaluation env and the policy-goal dimension.""" + data, raw_env = load_dataset(cfg) + replay = TensorDictReplayBuffer( + storage=LazyTensorStorage(data.numel(), device=device), + batch_size=cfg.optim.batch_size, + sample_unit=SubgoalTriplet( + goal_discount=cfg.loss.discount if cfg.sampling.geometric_goal else None, + policy_goal_discount=( + cfg.loss.discount if cfg.sampling.geometric_policy_goal else None + ), + ), + ) + replay.extend(data) + goal_dim = data["oracle_goal"].shape[-1] + env = GymWrapper(GoalObservation(raw_env, goal_dim), device="cpu") + env.set_info_dict_reader(SuccessReader(["success"])) + return replay, env, goal_dim + + +def mlp( + in_features: int, + hidden: list[int], + out_features: int, + *, + device: torch.device, + final_scale: float = 1.0, +) -> nn.Sequential: + """Dense -> GELU -> LayerNorm blocks with fan-average uniform init, as in the reference.""" + layers = [] + for width in hidden: + linear = nn.Linear(in_features, width, device=device) + nn.init.xavier_uniform_(linear.weight) + nn.init.zeros_(linear.bias) + layers += [ + linear, + nn.GELU(approximate="tanh"), + nn.LayerNorm(width, eps=1e-6, device=device), + ] + in_features = width + linear = nn.Linear(in_features, out_features, device=device) + nn.init.xavier_uniform_(linear.weight, gain=math.sqrt(final_scale)) + nn.init.zeros_(linear.bias) + layers.append(linear) + return nn.Sequential(*layers) + + +class Critic(nn.Module): + """Goal-conditioned critic logit ``Q(s, a, g)`` on concatenated inputs.""" + + def __init__(self, in_features: int, hidden: list[int], device: torch.device): + super().__init__() + self.net = mlp(in_features, hidden, 1, device=device) + + def forward(self, observation, goal, action): + return self.net(torch.cat([observation, goal, action], -1)) + + +class Actor(nn.Module): + """Gaussian policy head with a constant unit standard deviation.""" + + def __init__( + self, in_features: int, hidden: list[int], action_dim: int, device: torch.device + ): + super().__init__() + self.net = mlp(in_features, hidden, action_dim, device=device, final_scale=1e-2) + + def forward(self, observation, goal): + loc = self.net(torch.cat([observation, goal], -1)) + return loc, torch.ones_like(loc) + + +def make_agent(cfg: DictConfig, replay, goal_dim: int, device: torch.device): + """Return the policy, the TRL loss and its target updater.""" + sample = replay.storage[:2] + obs_dim = sample["observation"].shape[-1] + action_dim = sample["action"].shape[-1] + hidden = list(cfg.network.hidden) + oracle = cfg.loss.oracle_distill + if not oracle and goal_dim != obs_dim: + raise ValueError( + "loss.oracle_distill=false needs goals in the observation space, but " + f"{cfg.dataset.name} provides {goal_dim}-dimensional oracle goals; " + "set loss.oracle_distill=true." + ) + critic_keys = ["observation", "goal", "action"] + + def critics(goal_features): + return [ + TensorDictModule( + Critic(obs_dim + goal_features + action_dim, hidden, device), + in_keys=critic_keys, + out_keys=["state_action_value"], + ) + for _ in range(2) + ] + + actor = ProbabilisticActor( + TensorDictModule( + Actor(obs_dim + goal_dim, hidden, action_dim, device), + in_keys=["observation", "goal"], + out_keys=["loc", "scale"], + ), + in_keys=["loc", "scale"], + spec=Bounded(-1.0, 1.0, (action_dim,), device=device), + safe=True, + distribution_class=IndependentNormal, + default_interaction_type=ExplorationType.DETERMINISTIC, + ) + loss = TransitiveQLoss( + actor, + critics(obs_dim), + oracle_qvalue_network=critics(goal_dim) if oracle else None, + discount=cfg.loss.discount, + expectile=cfg.loss.expectile, + distance_weight=cfg.loss.distance_weight, + alpha=cfg.loss.alpha, + ) + loss.set_keys(policy_goal="oracle_goal") + updater = SoftUpdate(loss, tau=cfg.loss.tau) + return actor, loss, updater + + +@torch.no_grad() +def evaluate(policy, env, cfg: DictConfig) -> TensorDict: + """Success rate on each OGBench evaluation task and over all tasks. + + Each task runs ``cfg.evaluation.episodes`` deterministic episodes; an + episode succeeds if the env reports success on its last step. + """ + success = TensorDict(batch_size=[cfg.evaluation.episodes]) + with set_exploration_type(ExplorationType.DETERMINISTIC): + for task_id in range(1, cfg.evaluation.num_tasks + 1): + success[f"evaluation/task{task_id}"] = torch.stack( + [ + env.rollout( + cfg.evaluation.max_steps, + policy, + auto_reset=False, + tensordict=env.reset(options={"task_id": task_id}), + auto_cast_to_device=True, + )["next", "success"][-1].float() + for _ in range(cfg.evaluation.episodes) + ] + ) + success = success.mean(0) + success["evaluation/overall"] = torch.stack(list(success.values())).mean() + return success From 2d6fb672231fd147d37df984836c805a0407f4aa Mon Sep 17 00:00:00 2001 From: Achintya P Date: Wed, 21 Oct 2026 00:35:44 -0700 Subject: [PATCH 2/3] [Algorithm] Transitive RL recipe: compiled defaults, OGBench evaluation budget, per-dataset settings - compile and CUDA graphs are on by default, the setup the reported results and sota-check use. CUDA graphs (and the capturable Adam they need) apply only on a CUDA device, so device=cpu runs compiled without them; the CPU smoke test turns both off. - evaluation.num_tasks and evaluation.max_steps default to null and are read from the env (its OGBench evaluation tasks and time limit), as the reference evaluation runs each episode until the env ends it. The former 1000-step cap truncated tasks with longer limits, such as humanoidmaze-medium (2000); scene (750) and cube-single (200) are unaffected. - The module docstring lists the reference's per-dataset distance_weight, alpha and discount (aoberai/trl hyperparameters.sh): cube-single uses distance_weight 0.7, not the scene default of 1.0. --- .../unittest/linux_sota/scripts/test_sota.py | 1 + .../transitive_rl/config.yaml | 8 ++--- .../transitive_rl/transitive_rl.py | 30 +++++++++++++++++-- sota-implementations/transitive_rl/utils.py | 12 +++++++- 4 files changed, 43 insertions(+), 8 deletions(-) diff --git a/.github/unittest/linux_sota/scripts/test_sota.py b/.github/unittest/linux_sota/scripts/test_sota.py index 0aea2e74dfa..2cd0b801862 100644 --- a/.github/unittest/linux_sota/scripts/test_sota.py +++ b/.github/unittest/linux_sota/scripts/test_sota.py @@ -36,6 +36,7 @@ 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 \ diff --git a/sota-implementations/transitive_rl/config.yaml b/sota-implementations/transitive_rl/config.yaml index 7f0619e2d00..17729520e04 100644 --- a/sota-implementations/transitive_rl/config.yaml +++ b/sota-implementations/transitive_rl/config.yaml @@ -38,15 +38,15 @@ optim: lr: 0.0003 compile: - compile: false + compile: true compile_mode: - cudagraphs: false + cudagraphs: true evaluation: interval: 100000 - num_tasks: 5 + num_tasks: episodes: 15 - max_steps: 1000 + max_steps: hydra: job: diff --git a/sota-implementations/transitive_rl/transitive_rl.py b/sota-implementations/transitive_rl/transitive_rl.py index 6c1c6341917..dd15695b753 100644 --- a/sota-implementations/transitive_rl/transitive_rl.py +++ b/sota-implementations/transitive_rl/transitive_rl.py @@ -5,6 +5,29 @@ """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 @@ -31,8 +54,9 @@ def main(cfg: DictConfig) -> None: 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=cfg.compile.cudagraphs + loss_module.parameters(), lr=cfg.optim.lr, capturable=cudagraphs ) def update(batch): @@ -46,10 +70,10 @@ def update(batch): if cfg.compile.compile: compile_mode = cfg.compile.compile_mode or ( - "default" if cfg.compile.cudagraphs else "reduce-overhead" + "default" if cudagraphs else "reduce-overhead" ) update = compile_with_warmup(update, mode=compile_mode, warmup=1) - if cfg.compile.cudagraphs: + if cudagraphs: update = CudaGraphModule(update, in_keys=[], out_keys=[], warmup=5) output_dir = Path(HydraConfig.get().runtime.output_dir) logger = get_logger( diff --git a/sota-implementations/transitive_rl/utils.py b/sota-implementations/transitive_rl/utils.py index e7262525080..3f9de8e00b8 100644 --- a/sota-implementations/transitive_rl/utils.py +++ b/sota-implementations/transitive_rl/utils.py @@ -130,8 +130,18 @@ def load_dataset(cfg: DictConfig) -> tuple[TensorDict, gym.Env]: def make_replay_buffer(cfg: DictConfig, device: torch.device): - """Return the replay buffer, the evaluation env and the policy-goal dimension.""" + """Return the replay buffer, the evaluation env and the policy-goal dimension. + + A null ``cfg.evaluation.num_tasks`` or ``cfg.evaluation.max_steps`` is set + from the env: its OGBench evaluation tasks (one task for envs without any) + and its time limit, as the reference evaluation runs each episode until the + env ends it. + """ data, raw_env = load_dataset(cfg) + if cfg.evaluation.num_tasks is None: + cfg.evaluation.num_tasks = len(getattr(raw_env.unwrapped, "task_infos", [None])) + if cfg.evaluation.max_steps is None: + cfg.evaluation.max_steps = raw_env.spec.max_episode_steps replay = TensorDictReplayBuffer( storage=LazyTensorStorage(data.numel(), device=device), batch_size=cfg.optim.batch_size, From 2fdccb02741b5736b43191ed09d29f602c4a0df5 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Wed, 21 Oct 2026 01:59:05 -0700 Subject: [PATCH 3/3] [Doc] Transitive RL recipe: explain why GoalObservation wraps the gym env OGBench reports an evaluation episode's goal only in the reset info, and the policy reads it at every step, so the wrapper keeps it as an observation for the whole episode. --- sota-implementations/transitive_rl/utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/sota-implementations/transitive_rl/utils.py b/sota-implementations/transitive_rl/utils.py index 3f9de8e00b8..0890f9e6417 100644 --- a/sota-implementations/transitive_rl/utils.py +++ b/sota-implementations/transitive_rl/utils.py @@ -51,8 +51,11 @@ def __call__(self, info, tensordict): class GoalObservation(gym.Wrapper): """Expose the OGBench evaluation goal, reported in the reset info, next to a float32 observation. - Envs that report no goal use their reset observation, which only serves the - synthetic smoke-test setup. + OGBench reports the goal of an evaluation episode only in the info returned + by ``reset``, while the policy reads it at every step, so this wrapper keeps + it as a constant observation entry for the whole episode. Envs that report no + goal use their reset observation, which only serves the synthetic smoke-test + setup. """ def __init__(self, env: gym.Env, goal_dim: int):