[BugFix] Fix cudagraph_policy instability and add action std test by coder-jayp · Pull Request #4143 · pytorch/rl

coder-jayp · 2026-08-23T14:24:30Z

Description

This PR fixes a bug where the standard deviation of actions from a cudagraph_policy collapses to zero during training when using a stochastic policy compiled with torch.compile.

When a compiled module is captured into a CudaGraphModule by the SyncDataCollector, Inductor's internal Philox RNG state (seed and offset) gets frozen as static buffers. Consequently, the random samples drawn (e.g. from Beta or Normal distributions) are identical on every step replay.

Fix injects torch.compiler.cudagraph_mark_step_begin() immediately before the compiled policy is called during rollout. This signals PyTorch to emit an offset-increment kernel, ensuring fresh random numbers are generated on every execution of the CUDA graph. Additionally, a defensive .clone() is added to the policy outputs to prevent aliasing with the static CUDA graph buffers inside the TensorDict.

Motivation and Context

Fixes the learning instability and action standard deviation collapse observed when cudagraph_policy=True is used with a compiled stochastic policy.

close #3002

I have raised a bug report to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

Bug fix (non-breaking change which fixes a bug)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

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 · 2026-08-23T14:24:37Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 32b904a with merge base 3b6b5b9 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

Libs Tests on Linux / unittests-gym (3.10, 12.8) / linux-job (gh) (trunk failure)
test/libs/test_gym.py::TestGym::test_gym[True-True-1-HalfCheetah-v4]

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

coder-jayp · 2026-08-23T14:25:20Z

@torchrlbot reviewer @vmoens

github-actions · 2026-08-23T14:26:02Z

Requested review from @vmoens (requested by @coder-jayp).

vmoens

We're not really testing that we solved the bug

vmoens · 2026-08-23T15:28:39Z

+        from unittest.mock import patch
+
+        from tensordict.nn import TensorDictModule
+        from torchrl.envs.utils import ExplorationType, set_exploration_type


No local imports

vmoens · 2026-08-23T15:29:03Z

+
+        policy = TensorDictModule(
+            torch.nn.LazyLinear(n_actions), in_keys=["observation"], out_keys=["action"]
+        ).to("cuda:0")


Define device during construction

vmoens · 2026-08-23T15:30:08Z

+            assert (
+                mock_mark_step.called
+            ), "cudagraph_mark_step_begin was not called for cudagraphed_policy"


This doesn't test that the old behavior (the bug) existed and is solved, it just tests we called what we implemented which is circular thinking

coder-jayp · 2026-08-23T16:44:34Z

@vmoens Updated the test suite to address the feedback:

Added a behavioral variance check to verify the CUDA graph Philox RNG state advances across rollout steps (std > 0.1)
Moved the policy class to module scope, defined devices explicitly at construction, and added clean collector shutdown handling

Tested locally and ready for another look.

vmoens

Almost here
CI is failing and one minor comment on RandomPolicy

vmoens · 2026-08-23T18:03:20Z

        )


+class RandomPolicy(torch.nn.Module):


Maybe the one in the repo could be used?

I tried it, but the built-in RandomPolicy isn't an nn.Module, causing CudaGraphModule to crash when capturing its RNG allocations. The custom one subclasses nn.Module to fix the tracing.

coder-jayp · 2026-08-24T14:36:28Z

@vmoens Pushed a CI fix.

Unconditionally calling cudagraph_mark_step_begin() caused NotImplementedError on PT < 2.4 for native CUDA graphs. I updated _single.py to gracefully catch this for cudagraphed_policy, preserving backward compatibility.

… rsl_rl baselines (#13) * Add swarp.rl, a MAPPO trainer with the rollout and update CUDA-graph captured The torchrl trainer spends most of an iteration on framework overhead: TensorDict construction and clones per step, host syncs in the reset check and in every TanhNormal constructor, and an eager update with an .item() per minibatch. The simulator is a few percent of the iteration. swarp.rl keeps the algorithm and the recipes and removes the overhead. Per env step it replays an act graph, calls Environment.step and replays a store graph; per iteration it replays a prepare graph (critic, GAE, permutations) and one minibatch graph per minibatch, with one host sync for every logged statistic. Non-finite minibatches are skipped on device through fused Adam's found_inf. Torch only, no torchrl. End to end it is 2.4-2.5x faster than examples/train_mappo.py at the recipes' batch sizes, 3.1-3.4x with --compile --tf32, and the ratio holds up to 16384 envs. On giveway both trainers reach full difficulty at the same iteration (414 vs 415), in 83 s vs 208 s, and both final policies solve 1.000 of 2048 eval envs. Timeouts bootstrap from V(s_t), since the env auto-resets and has no terminal observation. The torchrl-free half of examples/mappo.py moves to swarp.rl.recipes and is re-exported. examples/train_mappo.py now times collection as well as the update, so its env_sps is end to end. eval_mappo.py reads both checkpoint formats. python -m swarp.benchmark.train_throughput measures both trainers through their CLIs. No change to the simulator step: swarp/core, swarp/scenarios and swarp/interop are untouched. * Make train_throughput's trainers a registry with per-config interpreters Each config is a Trainer entry (command, required module), and --python CONFIG=PATH runs a config under another venv, as compare_sims does, so baselines whose libraries cannot share the main venv plug in with one line. * Add an rsl_rl MAPPO baseline trainer for the training-throughput benchmark examples/train_rsl_rl.py drives rsl_rl's PPO and OnPolicyRunner.learn unmodified over a swarp Environment, one rsl_rl env per agent: the actor sees each agent's obs, the critic the env's joint obs repeated per agent, and per-env dones and time-outs are broadcast to the env's agents. The runner's Logger is replaced by device-side counters that write the shared metrics.csv (wall_s = rsl_rl's own collect_time + learn_time), run the TrainSpec curriculum gate and entropy schedule, and save rsl_rl/1 checkpoints that eval_mappo.py evaluates. Registered as the rsl_rl config of swarp.benchmark.train_throughput, with an rsl-rl optional extra. * Add a BenchMARL MAPPO baseline trainer examples/train_benchmarl.py runs MAPPO on a swarp scenario through BenchMARL's own Experiment: a VMAS-style task class over SwarpEnv with the per-agent entries nested under an "agents" group, MappoConfig/MlpConfig/ExperimentConfig set to the TrainSpec recipes as far as BenchMARL allows, and a callback that writes metrics.csv, saves benchmarl/1 checkpoints, runs the curriculum gate and the entropy schedule. Deviations from swarp.rl's MAPPO are listed in the module docstring. eval_mappo.py evaluates benchmarl/1 checkpoints (actor stored in swarp.rl's layout, no benchmarl import), train_throughput registers the benchmarl config, and the benchmarl extra is declared conflicting with torchrl in the uv lock since BenchMARL 1.5.2 pins torchrl < 0.12. * Make train_mappo's update, loss and CLI importable Split parse_args into build_parser + config_from_args, the loss construction into build_loss, and GAE plus the PPO epochs into an EagerUpdate class that train() takes as an argument, together with Collector overrides. A derived trainer can now swap the update and the collector settings while sharing the rest of the loop. Behaviour is unchanged: a seeded CPU run on navigation and giveway gives byte-identical metrics (wall time aside) and final weights before and after. * Add a TorchRL MAPPO trainer tuned with TorchRL's own speed features examples/train_mappo_tuned.py runs train_mappo's algorithm with a CUDA-graphed collector policy (cudagraph_policy), the update step in tensordict's CudaGraphModule, fused capturable Adam with an on-device non-finite skip through found_inf, one host read per iteration, and optional --compile/--tf32. Two TorchRL 0.13 issues are worked around: TanhNormal's float bounds are an H2D copy that cannot be captured, and cudagraph_policy stored the graph's static output buffers in every step of a batch, so the whole batch carried the last step's action and nothing learnt. Registered in train_throughput as torchrl-tuned (graphs, fp32) and torchrl-tuned-fast (+ --compile --tf32). Navigation, 512 envs, RTX 3070 Laptop: 113 and 94 ms/iter against 241 for train_mappo and 99.5/75 for swarp.rl graphs/fast. A GPU test checks the captured update against EagerUpdate on fresh batches through warm-up, capture and replay, and that a non-finite batch is skipped without touching the weights. * Sort test_mappo_tuned's imports as ruff resolves them from the main checkout * Document the training baselines and the throughput comparison against them docs/benchmarks.md's training-throughput tables now cover TorchRL as usually written, TorchRL with its own speed features, BenchMARL and rsl_rl at the recipes' sizes and over n_envs. Against tuned TorchRL swarp.rl is 1.1-1.4x faster rather than 2-3x, so the README and training page state both comparisons. * Make the step counter and the action-box check public Environment.step_count returns the live per-env counter that max_steps truncates on, so the rsl_rl adapter no longer reads env._step_count; swarp.rl.check_action_box replaces the private helper BenchMARL's script imported. A read-only property off the step path: no change to the step, so no throughput run. * Move to TorchRL 0.14.1 and record the giveway time-to-solve over three seeds The lockfile now resolves torchrl 0.14.1 and tensordict 0.14.3. TorchRL 0.14.0 fixed the cudagraph_policy aliasing bug (pytorch/rl#4143); the tuned trainer learns on 0.14.1 with or without FreshOutputs, at the same ms/iteration, so it stays for 0.13. Both TorchRL trainers time within 4% of their 0.13.3 numbers, as does swarp.rl alongside them. Giveway, full recipe, seeds 0-2: every trainer reaches full difficulty at iteration 413-417 and solves all 2048 evaluation envs, so time-to-solve follows iteration time; the curriculum, not learning, sets that iteration.

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 23, 2026

github-actions Bot added Collectors Integrations/torch_geometric Integrations BugFix labels Aug 23, 2026

github-actions Bot requested a review from vmoens August 23, 2026 14:26

vmoens reviewed Aug 23, 2026

View reviewed changes

vmoens mentioned this pull request Aug 23, 2026

Agents are bad at writing tests #4144

Open

[BugFix] Fix cudagraph_policy instability and add action std test

7b0a2af

coder-jayp force-pushed the fix/cudagraph-policy branch from d39f7a3 to 7b0a2af Compare August 23, 2026 16:37

vmoens reviewed Aug 23, 2026

View reviewed changes

Fix RandomPolicy shadowing

32b904a

vmoens approved these changes Aug 28, 2026

View reviewed changes

vmoens merged commit 60b1d01 into pytorch:main Aug 28, 2026
122 of 123 checks passed

Navigation Menu

Repository navigation

[BugFix] Fix cudagraph_policy instability and add action std test - #4143

[BugFix] Fix cudagraph_policy instability and add action std test#4143
vmoens merged 2 commits into
pytorch:mainfrom
coder-jayp:fix/cudagraph-policy

coder-jayp commented Aug 23, 2026

Uh oh!

pytorch-bot Bot commented Aug 23, 2026 •

edited

Loading

Uh oh!

coder-jayp commented Aug 23, 2026

Uh oh!

github-actions Bot commented Aug 23, 2026

Uh oh!

vmoens left a comment

Uh oh!

vmoens Aug 23, 2026

Uh oh!

vmoens Aug 23, 2026

Uh oh!

vmoens Aug 23, 2026

Uh oh!

coder-jayp commented Aug 23, 2026

Uh oh!

vmoens left a comment

Uh oh!

vmoens Aug 23, 2026

Uh oh!

coder-jayp Aug 24, 2026

Uh oh!

coder-jayp commented Aug 24, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

coder-jayp commented Aug 23, 2026

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

pytorch-bot Bot commented Aug 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

coder-jayp commented Aug 23, 2026

Uh oh!

github-actions Bot commented Aug 23, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

vmoens Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

vmoens Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

vmoens Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp commented Aug 23, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

vmoens Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp commented Aug 24, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Aug 23, 2026 •

edited

Loading