From 1bf098f89fea3f9c4940ea601401b94ed281adec Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Fri, 15 May 2026 09:40:07 +0100 Subject: [PATCH 1/3] Update [ghstack-poisoned] --- .../collectors/isaaclab_rnn_ppo_memory.py | 316 ++++++++++++++++++ .../isaaclab_rnn_ppo_memory_utils.py | 137 ++++++++ knowledge_base/ISAACLAB.md | 173 ++++++++++ 3 files changed, 626 insertions(+) create mode 100644 examples/collectors/isaaclab_rnn_ppo_memory.py create mode 100644 examples/collectors/isaaclab_rnn_ppo_memory_utils.py diff --git a/examples/collectors/isaaclab_rnn_ppo_memory.py b/examples/collectors/isaaclab_rnn_ppo_memory.py new file mode 100644 index 00000000000..d210b322460 --- /dev/null +++ b/examples/collectors/isaaclab_rnn_ppo_memory.py @@ -0,0 +1,316 @@ +# 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. +"""Recurrent PPO on Isaac Lab with compact rollouts and shifted GAE. + +Demonstrates running large-vectorized Isaac Lab environments with an LSTM +policy through a :class:`~torchrl.collectors.MultiCollector`. Isaac Lab is +launched only in the worker subprocesses (the main process never imports +``isaaclab``) so the main process can stay light and own the trainer. + +Key TorchRL features exercised: + +- :class:`~torchrl.collectors.MultiCollector` with ``policy_factory`` (each + worker builds its own policy copy and receives weights via + :class:`~torchrl.weight_update.MultiProcessWeightSyncScheme`). +- ``compact_obs=True`` + ``final_obs=True`` to drop the redundant + ``("next", obs)`` and carry the boundary observation as an + ``UnbatchedTensor`` under ``("final", obs)``. +- :class:`~torchrl.objectives.value.GAE` with ``shifted=True``: bootstraps + from the boundary obs without re-stepping the env. +- :class:`~torchrl.modules.LSTMModule` with a configurable + ``recurrent_backend``: during collection (``set_recurrent_mode=False``) + the LSTM auto-uses cuDNN regardless of the backend; the configured + backend kicks in only during training (``set_recurrent_mode=True``). +- Optional ``torch.compile`` + ``CudaGraphModule`` around the update step. + +Run on a SLURM Isaac Lab container: + +.. code-block:: bash + + python isaaclab_rnn_ppo_memory.py --num-envs 16384 --num-collectors 2 \ + --rollout-steps 32 --rnn-backend scan --compile-update --cudagraph-update +""" +from __future__ import annotations + +import argparse +import logging +import math +import sys +from functools import partial + +import torch +import torch.optim +from tensordict.nn import CudaGraphModule +from torchrl._utils import logger as torchrl_logger, timeit +from torchrl.collectors import MultiCollector +from torchrl.data import LazyTensorStorage, ReplayBuffer +from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement +from torchrl.envs import ExplorationType, set_exploration_type +from torchrl.modules import set_recurrent_mode +from torchrl.objectives import ClipPPOLoss, group_optimizers +from torchrl.objectives.value.advantages import GAE +from torchrl.record import WandbLogger +from torchrl.weight_update import MultiProcessWeightSyncScheme + +from isaaclab_rnn_ppo_memory_utils import make_env, make_models + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + # Env / rollout + parser.add_argument("--task", default="Isaac-Ant-v0") + parser.add_argument("--num-envs", type=int, default=16_384) + parser.add_argument("--num-collectors", type=int, default=1) + parser.add_argument("--rollout-steps", type=int, default=32) + parser.add_argument("--max-episode-steps", type=int, default=500) + parser.add_argument("--iterations", type=int, default=20) + # Model + parser.add_argument("--hidden-size", type=int, default=256) + parser.add_argument("--obs-dim", type=int, default=60) + parser.add_argument("--action-dim", type=int, default=8) + parser.add_argument( + "--rnn-backend", + choices=["pad", "scan", "triton"], + default="scan", + help=( + "LSTM backend used during training (set_recurrent_mode=True). " + "Collection (set_recurrent_mode=False) always falls back to cuDNN." + ), + ) + # PPO + parser.add_argument("--ppo-epochs", type=int, default=3) + parser.add_argument("--mini-batch-steps", type=int, default=8_192) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--gamma", type=float, default=0.99) + parser.add_argument("--gae-lambda", type=float, default=0.95) + parser.add_argument("--clip-epsilon", type=float, default=0.2) + parser.add_argument("--entropy-coeff", type=float, default=0.0) + parser.add_argument("--critic-coeff", type=float, default=1.0) + parser.add_argument("--clip-grad-norm", type=float, default=1.0) + # Compile / cudagraph + parser.add_argument( + "--compile-update", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument( + "--compile-mode", + choices=["default", "reduce-overhead", "max-autotune"], + default="default", + ) + parser.add_argument( + "--cudagraph-update", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument("--cudagraph-warmup", type=int, default=8) + # Devices + parser.add_argument("--collector-device", default="cuda:0") + parser.add_argument("--train-device", default="cuda:1") + parser.add_argument( + "--sync-collector", + action=argparse.BooleanOptionalAction, + default=False, + help="Synchronous MultiCollector. Async (default) yields per-worker batches.", + ) + # Logging + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--wandb-project", default="torchrl-isaac-rnn-memory") + parser.add_argument("--wandb-group", default="default") + parser.add_argument("--wandb-name", default=None) + parser.add_argument("--wandb-mode", default="online") + parser.add_argument("--log-level", default="INFO") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + logging.basicConfig(level=getattr(logging, args.log_level)) + if args.num_envs % args.num_collectors: + raise ValueError("--num-envs must be divisible by --num-collectors.") + if args.compile_update or args.cudagraph_update: + torch._dynamo.config.capture_scalar_outputs = True + + torch.manual_seed(args.seed) + torch.set_float32_matmul_precision("high") + collector_device = torch.device(args.collector_device) + train_device = torch.device(args.train_device) + if train_device.type == "cuda": + torch.cuda.set_device(train_device) + torch.cuda.manual_seed_all(args.seed) + + per_collector_envs = args.num_envs // args.num_collectors + frames_per_batch = args.num_envs * args.rollout_steps + sample_num_slices = max(1, args.mini_batch_steps // args.rollout_steps) + + # ---- Training model (lives on train_device, recurrent_mode=True path) ---- + actor, critic, full_value = make_models( + obs_dim=args.obs_dim, + action_dim=args.action_dim, + hidden_size=args.hidden_size, + rnn_backend=args.rnn_backend, + device=train_device, + ) + + loss_module = ClipPPOLoss( + actor_network=actor, + critic_network=full_value, + clip_epsilon=args.clip_epsilon, + entropy_coeff=args.entropy_coeff, + critic_coeff=args.critic_coeff, + normalize_advantage=True, + ) + adv_module = GAE( + gamma=args.gamma, + lmbda=args.gae_lambda, + value_network=full_value, + average_gae=False, + shifted=True, + device=train_device, + ) + optim = group_optimizers( + torch.optim.Adam( + actor.parameters(), lr=args.lr, eps=1e-5, capturable=args.cudagraph_update + ), + torch.optim.Adam( + critic.parameters(), lr=args.lr, eps=1e-5, capturable=args.cudagraph_update + ), + ) + + def update(batch): + with set_recurrent_mode(True): + loss = loss_module(batch) + total = loss["loss_objective"] + loss["loss_entropy"] + loss["loss_critic"] + total.backward() + grad_norm = torch.nn.utils.clip_grad_norm_( + list(actor.parameters()) + list(critic.parameters()), + max_norm=args.clip_grad_norm, + ) + optim.step() + optim.zero_grad(set_to_none=True) + loss_out = loss.detach() + loss_out.set("loss_total", total.detach()) + loss_out.set("grad_norm", grad_norm.detach()) + return loss_out + + if args.compile_update: + update = torch.compile(update, mode=args.compile_mode) + if args.cudagraph_update: + update = CudaGraphModule( + update, in_keys=[], out_keys=[], warmup=args.cudagraph_warmup + ) + + # ---- Collector (Isaac lives in workers; main process never imports it) ---- + make_env_fn = partial( + make_env, + task=args.task, + num_envs=per_collector_envs, + max_episode_steps=args.max_episode_steps, + device=str(collector_device), + ) + # The worker's actor uses the same backend; during collection the LSTM + # runs with set_recurrent_mode=False and auto-dispatches to cuDNN. + collector_policy_factory = partial( + _make_collector_actor, + obs_dim=args.obs_dim, + action_dim=args.action_dim, + hidden_size=args.hidden_size, + rnn_backend=args.rnn_backend, + device=collector_device, + ) + collector = MultiCollector( + [make_env_fn] * args.num_collectors, + sync=args.sync_collector, + policy_factory=collector_policy_factory, + frames_per_batch=frames_per_batch, + total_frames=-1, + policy_device=collector_device, + storing_device="cpu", + no_cuda_sync=True, + trust_policy=True, + compact_obs=True, + final_obs=True, + auto_register_policy_transforms=True, + track_policy_version=True, + weight_sync_schemes={"policy": MultiProcessWeightSyncScheme()}, + ) + + train_buffer = ReplayBuffer( + storage=LazyTensorStorage(args.num_envs, device=train_device, ndim=1), + sampler=SamplerWithoutReplacement(drop_last=True), + batch_size=sample_num_slices, + ) + updates_per_epoch = math.ceil(args.num_envs / sample_num_slices) + + experiment_logger: WandbLogger | None = None + if args.wandb_mode != "disabled": + experiment_logger = WandbLogger( + exp_name=args.wandb_name or f"{args.task}-{args.rnn_backend}", + project=args.wandb_project, + offline=args.wandb_mode in ("offline", "dryrun"), + group=args.wandb_group, + ) + experiment_logger.log_hparams(vars(args)) + + # ---- Training loop ---- + try: + for iteration, collected_batch in enumerate(collector): + if iteration >= args.iterations: + break + timeit.erase() + with timeit("collector_policy_sync"): + collector.update_policy_weights_(actor) + with timeit("iteration"): + data = collected_batch.to(train_device) + loss_acc = None + loss_count = 0 + for _ in range(args.ppo_epochs): + with timeit("advantage"), torch.no_grad(), set_recurrent_mode(True): + epoch_data = adv_module(data) + train_buffer.empty() + train_buffer.extend(epoch_data) + for mini_batch in train_buffer: + with timeit("update"): + loss = update(mini_batch.to(train_device)) + loss_acc = loss if loss_acc is None else loss_acc + loss + loss_count += 1 + if experiment_logger is not None: + metrics = timeit.todict(percall=False, prefix="time") + if loss_acc is not None and loss_count > 0: + metrics.update( + {f"loss/{k}": float(v / loss_count) for k, v in loss_acc.items()} + ) + metrics.update( + { + "iteration": iteration, + "frames": (iteration + 1) * frames_per_batch, + "updates_per_epoch": updates_per_epoch, + } + ) + experiment_logger.log_metrics(metrics, step=iteration) + torchrl_logger.info({"phase": "iteration_done", "iteration": iteration}) + finally: + collector.shutdown() + if experiment_logger is not None: + experiment_logger.experiment.finish() + + +def _make_collector_actor( + obs_dim: int, + action_dim: int, + hidden_size: int, + rnn_backend: str, + device, +): + """Worker-side actor factory: one fresh ProbabilisticActor per worker.""" + actor, _, _ = make_models( + obs_dim=obs_dim, + action_dim=action_dim, + hidden_size=hidden_size, + rnn_backend=rnn_backend, + device=device, + ) + return actor + + +if __name__ == "__main__": + with set_exploration_type(ExplorationType.RANDOM): + main() diff --git a/examples/collectors/isaaclab_rnn_ppo_memory_utils.py b/examples/collectors/isaaclab_rnn_ppo_memory_utils.py new file mode 100644 index 00000000000..f0f671b57db --- /dev/null +++ b/examples/collectors/isaaclab_rnn_ppo_memory_utils.py @@ -0,0 +1,137 @@ +# 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. +"""Utilities for the Isaac Lab recurrent PPO example. + +Isaac Lab is only imported lazily inside :func:`make_env`, so the +top-level main script can run without ``isaaclab`` on the Python path +and only the worker subprocesses (where ``make_env`` is called) pay the +import cost. +""" +from __future__ import annotations + +from typing import Literal + +import torch +import torch.nn as nn +from tensordict.nn import ( + AddStateIndependentNormalScale, + TensorDictModule, + TensorDictSequential, +) +from torchrl.envs import ExplorationType +from torchrl.modules import ( + LSTMModule, + MLP, + ProbabilisticActor, + TanhNormal, + ValueOperator, +) + +RnnBackend = Literal["pad", "scan", "triton"] + + +def make_env(task: str, num_envs: int, max_episode_steps: int, device: str): + """Build an Isaac Lab env. Imports ``isaaclab`` lazily (worker-only).""" + import gymnasium as gym + import isaaclab_tasks # noqa: F401 + from isaaclab_tasks.manager_based.classic.ant.ant_env_cfg import AntEnvCfg + from torchrl.envs.libs.isaac_lab import IsaacLabWrapper + + if task == "Isaac-Ant-v0": + cfg = AntEnvCfg() + if hasattr(cfg, "scene") and hasattr(cfg.scene, "num_envs"): + cfg.scene.num_envs = num_envs + if hasattr(cfg, "sim") and hasattr(cfg.sim, "device"): + cfg.sim.device = device + if hasattr(cfg, "device"): + cfg.device = device + if ( + hasattr(cfg, "episode_length_s") + and hasattr(cfg, "sim") + and hasattr(cfg.sim, "dt") + ): + cfg.episode_length_s = max_episode_steps * cfg.sim.dt + env = gym.make(task, cfg=cfg) + else: + env = gym.make(task) + return IsaacLabWrapper(env, device=torch.device(device)) + + +def make_models( + *, + obs_dim: int, + action_dim: int, + hidden_size: int, + rnn_backend: RnnBackend, + device: torch.device, +) -> tuple[ProbabilisticActor, ValueOperator, TensorDictSequential]: + """Build the actor, critic head, and the full value module sharing the backbone. + + The actor is a :class:`~torchrl.modules.ProbabilisticActor` wrapping a + shared embed + LSTM backbone and a per-action Gaussian head. The value + module reuses the same backbone (so a single GAE pass amortises the + recurrent compute across the policy and the value head). + """ + embed = TensorDictModule( + nn.Linear(obs_dim, hidden_size, device=device), + in_keys=["policy"], + out_keys=["embed"], + ) + lstm = LSTMModule( + input_size=hidden_size, + hidden_size=hidden_size, + num_layers=1, + in_keys=["embed", "recurrent_state_h", "recurrent_state_c"], + out_keys=[ + "lstm_out", + ("next", "recurrent_state_h"), + ("next", "recurrent_state_c"), + ], + recurrent_backend=rnn_backend, + device=device, + ) + backbone = TensorDictSequential(embed, lstm) + + actor_head = TensorDictModule( + nn.Sequential( + MLP( + in_features=hidden_size, + out_features=action_dim, + num_cells=[], + activation_class=nn.Tanh, + device=device, + ), + AddStateIndependentNormalScale(action_dim, scale_lb=1e-4).to(device), + ), + in_keys=["lstm_out"], + out_keys=["loc", "scale"], + ) + actor = ProbabilisticActor( + module=TensorDictSequential(backbone, actor_head), + in_keys=["loc", "scale"], + out_keys=["action"], + distribution_class=TanhNormal, + distribution_kwargs={ + "low": -torch.ones(action_dim, device=device), + "high": torch.ones(action_dim, device=device), + "tanh_loc": False, + }, + return_log_prob=True, + default_interaction_type=ExplorationType.RANDOM, + ) + + # The value module re-uses the backbone (shared params, single LSTM call + # per GAE/update pass). An identity TDM caches the lstm_out under a + # distinct key so the critic head and the actor head don't fight over + # write semantics on a shared key. + value_feature = TensorDictModule( + nn.Identity(), in_keys=["lstm_out"], out_keys=["value_lstm_out"] + ) + critic = ValueOperator( + nn.Linear(hidden_size, 1, device=device), + in_keys=["value_lstm_out"], + ).to(device) + full_value = TensorDictSequential(backbone, value_feature, critic) + return actor, critic, full_value diff --git a/knowledge_base/ISAACLAB.md b/knowledge_base/ISAACLAB.md index b2f3ea2c8d9..08f8b9e5127 100644 --- a/knowledge_base/ISAACLAB.md +++ b/knowledge_base/ISAACLAB.md @@ -174,9 +174,158 @@ steve step -d "$JOBID" 'WANDB_MODE=online bash /root/setup-and-run.sh env.name=I | Variable | Value | Purpose | |----------|-------|---------| | `OMNI_KIT_ACCEPT_EULA` | `yes` | Accept IsaacLab EULA (required) | +| `ACCEPT_EULA` | `Y` | Accept NVIDIA container / Omniverse EULA in non-interactive jobs | +| `PRIVACY_CONSENT` | `Y` | Avoid privacy consent prompts in Isaac Sim / Omniverse startup | +| `OMNI_KIT_DISABLE_CUP` | `1` | Disable Customer Usage Profile prompts | +| `OMNI_KIT_ALLOW_ROOT` | `1` | Allow Kit to run as root in containerized jobs | | `PYTHONNOUSERSITE` | `1` | Avoid user-site packages conflicting | | `WANDB_MODE` | `online` | Weights & Biases logging mode | +## Pip-Based IsaacLab Flow on Cluster Jobs + +Prefer the IsaacLab container when possible. If the job starts from a generic +CUDA container, a pip-based flow can work, but the order matters. + +### Use `uv venv` when `ensurepip` is missing + +Some cluster images ship Python without `ensurepip`, so `python -m venv` creates +an unusable environment. Use `uv venv` directly: + +```bash +VENV_DIR=/root/.venv/rl-isaac +uv venv "$VENV_DIR" +source "$VENV_DIR/bin/activate" +``` + +Run installs from `/root` (or another directory without a restrictive +`pyproject.toml`) and clear inherited uv resolution pins: + +```bash +cd /root +unset UV_EXCLUDE_NEWER UV_EXCLUDE_NEWER_PACKAGE +``` + +### Install Isaac Sim before IsaacLab + +Install Isaac Sim from NVIDIA's package index, then install IsaacLab from source: + +```bash +uv pip install "isaacsim[all,extscache]==6.0.0" \ + --extra-index-url https://pypi.nvidia.com \ + --index-strategy unsafe-best-match \ + --prerelease=allow + +git clone https://github.com/isaac-sim/IsaacLab.git /root/IsaacLab +cd /root/IsaacLab +./isaaclab.sh --install || true +``` + +The `|| true` is intentional for non-interactive setup scripts: the package +installation can succeed, then the VSCode / Kit bootstrap step can still ask for +the EULA and exit with EOF. The runtime job must set the EULA variables listed +above. + +### Check package presence without importing Isaac + +Do not use `python -c "import isaacsim"` or `python -c "import isaaclab"` as an +installation check in setup scripts. Importing can initialize Kit or trigger EULA +prompts. Check package presence with `importlib.util.find_spec` instead: + +```bash +python -c "import importlib.util; raise SystemExit(importlib.util.find_spec('isaacsim') is None)" +python -c "import importlib.util; raise SystemExit(importlib.util.find_spec('isaaclab') is None)" +``` + +### Reinstall the target torch after IsaacLab + +IsaacLab's installer may install its preferred torch build, for example a cu128 +wheel. If the experiment needs a torch nightly, reinstall it after IsaacLab: + +```bash +uv pip install --upgrade --pre torch torchvision \ + --index-url https://download.pytorch.org/whl/nightly/cu130 +``` + +For cu130 nightlies, also reinstall the cu13 NCCL wheel after any IsaacLab or +cu12/cu128 torch install. `nvidia-nccl-cu12` and `nvidia-nccl-cu13` write to the +same `site-packages/nvidia/nccl/lib` path; the cu12 package can overwrite +`libnccl.so.2` and make `import torch` fail with: + +```text +ImportError: .../libtorch_cuda.so: undefined symbol: ncclCommResume +``` + +The fix is to force reinstall the matching cu13 NCCL wheel and make the venv's +NVIDIA libraries first on `LD_LIBRARY_PATH`: + +```bash +uv pip install --force-reinstall "nvidia-nccl-cu13==2.30.4" \ + --index-url https://download.pytorch.org/whl/nightly/cu130 + +SITE_PKGS="$(python -c 'import site; print(site.getsitepackages()[0])')" +NVIDIA_LIBS="$SITE_PKGS/nvidia/nccl/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/cublas/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/cuda_runtime/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/cudnn/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/cufft/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/curand/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/cusolver/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/cusparse/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/cuda_nvrtc/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/nvjitlink/lib" +NVIDIA_LIBS="$NVIDIA_LIBS:$SITE_PKGS/nvidia/nvtx/lib" +export LD_LIBRARY_PATH="$NVIDIA_LIBS:${LD_LIBRARY_PATH:-}" +python -c "import torch; print(torch.__version__, torch.version.cuda)" +``` + +### Install TorchRL and TensorDict without dependencies + +After torch is correct, install local TorchRL and TensorDict with `--no-deps` so +pip does not replace the torch runtime: + +```bash +uv pip install -e /root/tensordict --no-deps +uv pip install -e /root/rl-isaac --no-deps +``` + +If TensorDict and TorchRL are being tested across stacked PRs, make sure the +TensorDict branch contains the compatibility hooks expected by the TorchRL +branch. A stale TensorDict checkout can fail during `import torchrl` even when +torch itself imports correctly. + +## TorchRL Recurrent PPO Flow + +For IsaacLab PPO with an RNN policy, keep the dataflow explicitly on-policy: + +1. Build the collector policy with `policy_factory`, not by manually injecting + recurrent reset keys into the env. +2. On TorchRL branches where `auto_register_policy_transforms` still defaults + to `None`, pass `auto_register_policy_transforms=True` so the collector adds + `InitTracker` and recurrent state `TensorDictPrimer` transforms. +3. Use separate devices for large runs when possible: one GPU for Isaac + collection/inference, one GPU for learner updates. +4. Collect one rollout window, freeze it, and run GAE over the full window once + per PPO epoch. +5. For each epoch, empty/fill a training replay buffer from that processed + window, sample random slices for minibatches, then discard the buffer before + the next collection. +6. After training, sync learner weights back to the collector before collecting + the next window. + +In the current collector fallback path, the robust weight sync form is: + +```python +from tensordict import TensorDict + +collector.update_policy_weights_(weights=TensorDict.from_module(actor).data) +``` + +With compact rollout data, prefer `shifted=True` value estimation so the PPO +batch does not need `("next", "policy")` rehydration. If a backend requires +canonical strides, `td.contiguous()` and `td.clone()` may not be enough for +size-1 dimensions; `torch.empty_like(td).update_(td)` is the stronger +materialization pattern, and RNN backends should ideally handle this internally. + ## Gotchas and Pitfalls 1. **Import order**: AppLauncher MUST be initialized before `import torch`. This cannot be stressed enough. @@ -209,3 +358,27 @@ steve step -d "$JOBID" 'WANDB_MODE=online bash /root/setup-and-run.sh env.name=I 11. **Memory**: Pixel replay buffers are large. 500K frames at 64×64×3 float32 ≈ 24 GB on disk (memmap). 12. **Installing torchrl in Isaac container**: Use `--no-build-isolation --no-deps` to avoid conflicts with Isaac's pre-installed torch/numpy. + +13. **Pip IsaacLab can downgrade torch**: Always verify + `python -c "import torch; print(torch.__version__, torch.version.cuda)"` + after IsaacLab installation, then reinstall the desired torch build if + needed. + +14. **cu12 and cu13 NVIDIA wheels share runtime paths**: If a cu130 nightly + fails with `ncclCommResume`, force reinstall `nvidia-nccl-cu13` after all + cu12/cu128 installs and prepend the venv NVIDIA library paths to + `LD_LIBRARY_PATH`. + +15. **Non-interactive setup checks should not import Isaac**: Use + `importlib.util.find_spec` for install checks; save real Isaac imports for + the final runtime after EULA variables are exported. + +16. **RNN collector transforms**: For recurrent policies, use + `policy_factory` and collector auto-registration for `InitTracker` and + recurrent state primers. On branches where the default is still transitional, + passing `auto_register_policy_transforms=True` is required. + +17. **On-policy buffers**: Do not let a continuously filled replay buffer drive + PPO updates. Fill one training window, compute GAE over the whole window, + train for the configured epochs, empty the buffer, sync weights, then + collect again. From 35474f1d440b20b26162f31386072d1a0c5a6e4c Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Fri, 15 May 2026 13:59:46 +0100 Subject: [PATCH 2/3] Update [ghstack-poisoned] --- .gitignore | 7 + .../collectors/isaaclab_rnn_ppo_memory.py | 232 ++++++++++++++++-- .../isaaclab_rnn_ppo_memory_utils.py | 28 ++- 3 files changed, 238 insertions(+), 29 deletions(-) diff --git a/.gitignore b/.gitignore index c44d0f8f87b..9b4e51a80be 100644 --- a/.gitignore +++ b/.gitignore @@ -200,3 +200,10 @@ main/ # Additional cache directories .ruff_cache/ + +# Local agent / experiment artifacts +.codex/ +.cursorignore +.github/rulesets/ +TODO/ +uv.lock diff --git a/examples/collectors/isaaclab_rnn_ppo_memory.py b/examples/collectors/isaaclab_rnn_ppo_memory.py index d210b322460..fa9953093cb 100644 --- a/examples/collectors/isaaclab_rnn_ppo_memory.py +++ b/examples/collectors/isaaclab_rnn_ppo_memory.py @@ -14,11 +14,11 @@ - :class:`~torchrl.collectors.MultiCollector` with ``policy_factory`` (each worker builds its own policy copy and receives weights via :class:`~torchrl.weight_update.MultiProcessWeightSyncScheme`). -- ``compact_obs=True`` + ``final_obs=True`` to drop the redundant - ``("next", obs)`` and carry the boundary observation as an - ``UnbatchedTensor`` under ``("final", obs)``. -- :class:`~torchrl.objectives.value.GAE` with ``shifted=True``: bootstraps - from the boundary obs without re-stepping the env. +- ``compact_obs=True`` to drop the redundant ``("next", obs)``. Shifted + value estimation reconstructs next observations by shifting the root + observations when the next observations are absent. +- :class:`~torchrl.objectives.value.GAE` with ``shifted=True``: uses the + root observation shift when compact rollouts omit ``("next", obs)``. - :class:`~torchrl.modules.LSTMModule` with a configurable ``recurrent_backend``: during collection (``set_recurrent_mode=False``) the LSTM auto-uses cuDNN regardless of the backend; the configured @@ -37,11 +37,13 @@ import argparse import logging import math -import sys from functools import partial import torch import torch.optim + +from isaaclab_rnn_ppo_memory_utils import _init_isaac_app, make_env, make_models +from tensordict import TensorDictBase from tensordict.nn import CudaGraphModule from torchrl._utils import logger as torchrl_logger, timeit from torchrl.collectors import MultiCollector @@ -54,7 +56,148 @@ from torchrl.record import WandbLogger from torchrl.weight_update import MultiProcessWeightSyncScheme -from isaaclab_rnn_ppo_memory_utils import make_env, make_models + +_RECURRENT_STATE_KEYS = { + "recurrent_state_h", + "recurrent_state_c", + "('next', 'recurrent_state_h')", + "('next', 'recurrent_state_c')", +} + + +def _leaf_shape_summary(tensordict: TensorDictBase) -> dict[str, dict[str, str]]: + return { + str(key): { + "shape": str(tuple(value.shape)), + "dtype": str(value.dtype), + "device": str(value.device), + } + for key, value in tensordict.items(include_nested=True, leaves_only=True) + if hasattr(value, "shape") + } + + +def _metric_float(value) -> float: + if isinstance(value, torch.Tensor): + value = value.detach() + if value.numel() != 1: + value = value.float().mean() + return float(value.cpu()) + return float(value) + + +def _tensor_stats(prefix: str, value: torch.Tensor) -> dict[str, float]: + value = value.detach().float() + return { + f"{prefix}/mean": _metric_float(value.mean()), + f"{prefix}/std": _metric_float(value.std(unbiased=False)), + f"{prefix}/min": _metric_float(value.min()), + f"{prefix}/max": _metric_float(value.max()), + } + + +def _loss_metrics(loss_acc: TensorDictBase, loss_count: int) -> dict[str, float]: + metrics = {} + for key, value in loss_acc.items(): + value = value / loss_count + key = str(key) + if key.startswith("loss_"): + key = f"loss/{key.removeprefix('loss_')}" + elif key.startswith("grad_norm"): + key = key.replace("grad_norm", "grad_norm/") + metrics[f"training/{key}"] = _metric_float(value) + return metrics + + +def _inference_metrics( + data: TensorDictBase, + *, + frames: int, +) -> dict[str, float | int]: + metrics: dict[str, float | int] = { + "inference/frames": frames, + "inference/batch_numel": data.numel(), + "inference/batch_ndim": data.ndim, + } + reward = data.get(("next", "reward"), default=None) + if reward is not None: + metrics.update(_tensor_stats("inference/reward", reward)) + episode_reward = data.get(("next", "episode_reward"), default=None) + done = data.get(("next", "done"), default=None) + if episode_reward is not None and done is not None: + episode_reward = episode_reward.squeeze(-1) + done = done.squeeze(-1).to(torch.bool) + end_of_traj_reward = episode_reward[done] + if end_of_traj_reward.numel(): + metrics.update( + _tensor_stats( + "inference/end_of_traj_episode_reward", end_of_traj_reward + ) + ) + return metrics + + +def _cuda_metrics(prefix: str, device: torch.device) -> dict[str, float]: + if device.type != "cuda": + return {} + return { + f"telemetry/{prefix}/allocated_gb": torch.cuda.memory_allocated(device) / 1e9, + f"telemetry/{prefix}/reserved_gb": torch.cuda.memory_reserved(device) / 1e9, + f"telemetry/{prefix}/max_allocated_gb": torch.cuda.max_memory_allocated(device) + / 1e9, + f"telemetry/{prefix}/max_reserved_gb": torch.cuda.max_memory_reserved(device) + / 1e9, + } + + +def _assert_rollout_shapes( + tensordict: TensorDictBase, + *, + expected_shape: torch.Size, + hidden_size: int, + phase: str, +) -> None: + if tensordict.shape != expected_shape: + raise RuntimeError( + f"{phase}: expected TensorDict shape {expected_shape}, " + f"got {tensordict.shape}." + ) + expected_state_shape = (*expected_shape, 1, hidden_size) + for key, value in tensordict.items(include_nested=True, leaves_only=True): + if not hasattr(value, "shape"): + continue + if value.shape[: len(expected_shape)] != expected_shape: + raise RuntimeError( + f"{phase}: key {key} has shape {tuple(value.shape)}, " + f"which does not start with {tuple(expected_shape)}." + ) + if str(key) in _RECURRENT_STATE_KEYS and tuple(value.shape) != tuple( + expected_state_shape + ): + raise RuntimeError( + f"{phase}: key {key} has recurrent-state shape " + f"{tuple(value.shape)}, expected {tuple(expected_state_shape)}." + ) + + +def _normalize_rollout_batch( + tensordict: TensorDictBase, expected_shape: torch.Size +) -> TensorDictBase: + if tensordict.shape == expected_shape: + return tensordict + if tensordict.shape == torch.Size((1, *expected_shape)): + return tensordict.squeeze(0) + if tensordict.ndim < 2 or tensordict.shape[-1] != expected_shape[-1]: + raise RuntimeError( + f"Expected collected batch ending in time shape {tuple(expected_shape)}, " + f"got {tuple(tensordict.shape)}." + ) + if tensordict.shape[:-1].numel() != expected_shape[0]: + raise RuntimeError( + f"Expected collected batch with {expected_shape[0]} env elements before " + f"time, got shape {tuple(tensordict.shape)}." + ) + return tensordict.reshape(expected_shape) def parse_args() -> argparse.Namespace: @@ -72,10 +215,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--action-dim", type=int, default=8) parser.add_argument( "--rnn-backend", - choices=["pad", "scan", "triton"], + choices=["cudnn", "pad", "scan", "triton"], default="scan", help=( "LSTM backend used during training (set_recurrent_mode=True). " + "'cudnn' is an alias for the pad/nn.LSTM path. " "Collection (set_recurrent_mode=False) always falls back to cuDNN." ), ) @@ -140,6 +284,7 @@ def main() -> None: per_collector_envs = args.num_envs // args.num_collectors frames_per_batch = args.num_envs * args.rollout_steps sample_num_slices = max(1, args.mini_batch_steps // args.rollout_steps) + expected_rollout_shape = torch.Size((args.num_envs, args.rollout_steps)) # ---- Training model (lives on train_device, recurrent_mode=True path) ---- actor, critic, full_value = make_models( @@ -227,14 +372,13 @@ def update(batch): no_cuda_sync=True, trust_policy=True, compact_obs=True, - final_obs=True, + init_fn=partial(_init_isaac_app, device=str(collector_device)), auto_register_policy_transforms=True, - track_policy_version=True, weight_sync_schemes={"policy": MultiProcessWeightSyncScheme()}, ) train_buffer = ReplayBuffer( - storage=LazyTensorStorage(args.num_envs, device=train_device, ndim=1), + storage=LazyTensorStorage(args.num_envs, device="cpu", ndim=1), sampler=SamplerWithoutReplacement(drop_last=True), batch_size=sample_num_slices, ) @@ -251,22 +395,56 @@ def update(batch): experiment_logger.log_hparams(vars(args)) # ---- Training loop ---- + collector_iter = iter(collector) try: - for iteration, collected_batch in enumerate(collector): - if iteration >= args.iterations: - break + for iteration in range(args.iterations): timeit.erase() with timeit("collector_policy_sync"): collector.update_policy_weights_(actor) - with timeit("iteration"): - data = collected_batch.to(train_device) + with timeit("collector_next"): + collected_batch = next(collector_iter) + with timeit("training"): + data = _normalize_rollout_batch(collected_batch, expected_rollout_shape) loss_acc = None loss_count = 0 - for _ in range(args.ppo_epochs): + for epoch in range(args.ppo_epochs): + epoch_data = data.to(train_device) + _assert_rollout_shapes( + epoch_data, + expected_shape=expected_rollout_shape, + hidden_size=args.hidden_size, + phase="before_gae", + ) with timeit("advantage"), torch.no_grad(), set_recurrent_mode(True): - epoch_data = adv_module(data) + epoch_data = adv_module(epoch_data) + _assert_rollout_shapes( + epoch_data, + expected_shape=expected_rollout_shape, + hidden_size=args.hidden_size, + phase="after_gae", + ) + if iteration == 0 and epoch == 0: + torchrl_logger.info( + { + "phase": "pre_train_buffer_extend", + "epoch_data_shape": tuple(epoch_data.shape), + "epoch_data_leaf_shapes": _leaf_shape_summary( + epoch_data + ), + "epoch_data_0_shape": tuple(epoch_data[0].shape), + "epoch_data_0_leaf_shapes": _leaf_shape_summary( + epoch_data[0] + ), + } + ) train_buffer.empty() train_buffer.extend(epoch_data) + if train_buffer._storage._storage.shape != expected_rollout_shape: + raise RuntimeError( + "Expected train buffer storage shape " + f"{expected_rollout_shape}, got " + f"{train_buffer._storage._storage.shape}." + ) for mini_batch in train_buffer: with timeit("update"): loss = update(mini_batch.to(train_device)) @@ -275,16 +453,22 @@ def update(batch): if experiment_logger is not None: metrics = timeit.todict(percall=False, prefix="time") if loss_acc is not None and loss_count > 0: - metrics.update( - {f"loss/{k}": float(v / loss_count) for k, v in loss_acc.items()} - ) + metrics.update(_loss_metrics(loss_acc, loss_count)) metrics.update( { - "iteration": iteration, - "frames": (iteration + 1) * frames_per_batch, - "updates_per_epoch": updates_per_epoch, + "training/iteration": iteration, + "training/updates_per_epoch": updates_per_epoch, + "training/updates_total": loss_count, } ) + metrics.update( + _inference_metrics( + data, + frames=(iteration + 1) * frames_per_batch, + ) + ) + metrics.update(_cuda_metrics("collector_cuda", collector_device)) + metrics.update(_cuda_metrics("train_cuda", train_device)) experiment_logger.log_metrics(metrics, step=iteration) torchrl_logger.info({"phase": "iteration_done", "iteration": iteration}) finally: diff --git a/examples/collectors/isaaclab_rnn_ppo_memory_utils.py b/examples/collectors/isaaclab_rnn_ppo_memory_utils.py index f0f671b57db..1ff8196f16e 100644 --- a/examples/collectors/isaaclab_rnn_ppo_memory_utils.py +++ b/examples/collectors/isaaclab_rnn_ppo_memory_utils.py @@ -11,6 +11,7 @@ """ from __future__ import annotations +import argparse from typing import Literal import torch @@ -20,7 +21,7 @@ TensorDictModule, TensorDictSequential, ) -from torchrl.envs import ExplorationType +from torchrl.envs import ExplorationType, RewardSum, TransformedEnv from torchrl.modules import ( LSTMModule, MLP, @@ -29,7 +30,20 @@ ValueOperator, ) -RnnBackend = Literal["pad", "scan", "triton"] +RnnBackend = Literal["cudnn", "pad", "scan", "triton"] + + +def _init_isaac_app(device: str | None = None) -> None: + """Start Isaac Lab's AppLauncher in headless mode inside a worker.""" + from isaaclab.app import AppLauncher + + parser = argparse.ArgumentParser(description="TorchRL Isaac Lab env launcher.") + AppLauncher.add_app_launcher_args(parser) + launch_args = ["--headless"] + if device is not None: + launch_args.extend(["--device", device]) + args_cli, _ = parser.parse_known_args(launch_args) + AppLauncher(args_cli) def make_env(task: str, num_envs: int, max_episode_steps: int, device: str): @@ -56,7 +70,10 @@ def make_env(task: str, num_envs: int, max_episode_steps: int, device: str): env = gym.make(task, cfg=cfg) else: env = gym.make(task) - return IsaacLabWrapper(env, device=torch.device(device)) + return TransformedEnv( + IsaacLabWrapper(env, device=torch.device(device)), + RewardSum(), + ) def make_models( @@ -79,6 +96,7 @@ def make_models( in_keys=["policy"], out_keys=["embed"], ) + lstm_backend = "pad" if rnn_backend == "cudnn" else rnn_backend lstm = LSTMModule( input_size=hidden_size, hidden_size=hidden_size, @@ -89,7 +107,7 @@ def make_models( ("next", "recurrent_state_h"), ("next", "recurrent_state_c"), ], - recurrent_backend=rnn_backend, + recurrent_backend=lstm_backend, device=device, ) backbone = TensorDictSequential(embed, lstm) @@ -122,7 +140,7 @@ def make_models( default_interaction_type=ExplorationType.RANDOM, ) - # The value module re-uses the backbone (shared params, single LSTM call + # The value module reuses the backbone (shared params, single LSTM call # per GAE/update pass). An identity TDM caches the lstm_out under a # distinct key so the critic head and the actor head don't fight over # write semantics on a shared key. From d4e5ea893b3767ed88dcddc49bc8149fa4552986 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Fri, 15 May 2026 16:13:22 +0100 Subject: [PATCH 3/3] Update [ghstack-poisoned] --- examples/collectors/isaaclab_rnn_ppo_memory.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/examples/collectors/isaaclab_rnn_ppo_memory.py b/examples/collectors/isaaclab_rnn_ppo_memory.py index fa9953093cb..cd33e79d03b 100644 --- a/examples/collectors/isaaclab_rnn_ppo_memory.py +++ b/examples/collectors/isaaclab_rnn_ppo_memory.py @@ -17,8 +17,10 @@ - ``compact_obs=True`` to drop the redundant ``("next", obs)``. Shifted value estimation reconstructs next observations by shifting the root observations when the next observations are absent. -- :class:`~torchrl.objectives.value.GAE` with ``shifted=True``: uses the - root observation shift when compact rollouts omit ``("next", obs)``. +- :class:`~torchrl.objectives.value.GAE` with ``shifted="compact"``: a + constant-shape single value-network call along the time dim. No reads + of ``("next", obs)``, no Python branches on tensor values, no + ``.item()`` syncs — friendly to ``torch.compile`` + scan/triton LSTM. - :class:`~torchrl.modules.LSTMModule` with a configurable ``recurrent_backend``: during collection (``set_recurrent_mode=False``) the LSTM auto-uses cuDNN regardless of the backend; the configured @@ -308,7 +310,7 @@ def main() -> None: lmbda=args.gae_lambda, value_network=full_value, average_gae=False, - shifted=True, + shifted="compact", device=train_device, ) optim = group_optimizers(