Skip to content

[Feature] Add GRPO candidate selection and token advantages - #3875

Merged
vmoens merged 5 commits into
gh/vmoens/287/basefrom
gh/vmoens/287/head
Jun 18, 2026
Merged

[Feature] Add GRPO candidate selection and token advantages#3875
vmoens merged 5 commits into
gh/vmoens/287/basefrom
gh/vmoens/287/head

Conversation

@vmoens

@vmoens vmoens commented Jun 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

@vmoens
Update
6d0dd61 
[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jun 17, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 New Failure

As of commit a8215d4 with merge base 6364a19 (image):

NEW FAILURE - The following job has failed:

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

@meta-cla 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 Jun 17, 2026
@github-actions github-actions Bot added Documentation Improvements or additions to documentation Objectives llm/ LLM-related PR, triggers LLM CI tests Trainers Feature New feature labels Jun 17, 2026
This was referenced Jun 17, 2026
Merged
Merged
Merged
Merged
Merged
Merged
vmoens added 2 commits June 17, 2026 17:20
@vmoens
Update
5846cf7 
[ghstack-poisoned]
@vmoens
Update
95ba841 
[ghstack-poisoned]
@github-actions

github-actions Bot commented Jun 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Benchmark Results: PR a8215d49 vs main 2ab5b651

Benchmark run: https://github.com/pytorch/rl/actions/runs/27750668444

Higher ops/sec is better. Tables are sorted by largest absolute change.

CPU

Compared 192 benchmarks. Regressions over 5%: 9. Improvements over 5%: 15.

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 49.80 197.06 +295.73%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 859.39 3,061 +256.13%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 193.81 51.31 -73.53%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] 108.81 141.10 +29.68%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 2,826 3,654 +29.30%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 2,854 3,519 +23.33%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward] 239.64 274.36 +14.49%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 2,801 3,186 +13.74%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 868.18 750.30 -13.58%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 738.52 823.66 +11.53%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,688 2,381 -11.42%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 2,702 3,000 +11.03%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,235 2,028 -9.26%
benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False] 8,080 7,334 -9.23%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 1,037 953.48 -8.03%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 692.19 746.86 +7.90%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 528.19 489.33 -7.36%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None] 438.83 471.03 +7.34%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 517.03 479.66 -7.23%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 516.56 480.73 -6.94%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] 112.40 118.87 +5.76%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 639.77 675.62 +5.60%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False] 48,961 51,605 +5.40%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None] 265.78 279.90 +5.31%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None] 273.44 287.08 +4.99%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None] 270.66 284.10 +4.96%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 2,858 2,998 +4.88%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,139 2,038 -4.74%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant] 4,329 4,139 -4.39%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True] 23,935 22,998 -3.91%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle] 12,151 11,689 -3.80%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True] 21,088 20,373 -3.39%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward] 947.11 977.56 +3.21%
benchmarks/test_envs_benchmark.py::test_parallel 0.9677 0.9368 -3.19%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None] 113.33 116.90 +3.15%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-gru] 1.3055 1.2658 -3.04%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True] 17,826 18,367 +3.03%
benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False] 23.81 23.11 -2.94%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] 172.23 167.76 -2.60%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None] 257.27 263.74 +2.51%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] 305.34 312.91 +2.48%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-backward] 28.09 27.40 -2.46%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False] 32,125 31,363 -2.37%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari] 679.66 695.36 +2.31%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari] 639.34 654.07 +2.30%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward] 128.94 131.84 +2.25%
benchmarks/test_envs_benchmark.py::test_transformed 0.8887 0.9086 +2.24%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch] 324.85 332.10 +2.23%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch] 766.01 749.25 -2.19%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None] 82.28 84.03 +2.12%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward] 81.69 83.40 +2.09%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-gru] 1.3934 1.3660 -1.97%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None] 700.34 686.78 -1.94%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 1,060 1,080 +1.81%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None] 326.47 332.38 +1.81%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False] 75,783 77,136 +1.79%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cpu] 96.82 98.53 +1.77%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward] 78.14 76.78 -1.73%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True] 0.5204 0.5294 +1.73%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch] 14.97 15.23 +1.73%
benchmarks/test_objectives_benchmarks.py::test_redq_speed[True-None] 223.90 227.78 +1.73%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] 29,949 30,463 +1.72%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-lstm] 0.8323 0.8180 -1.71%
benchmarks/test_envs_benchmark.py::test_simple 1.7643 1.7944 +1.71%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] 52.97 53.87 +1.70%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None] 535.64 544.64 +1.68%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-None] 208.52 211.98 +1.66%
benchmarks/test_objectives_benchmarks.py::test_redq_speed[False-None] 95.46 93.88 -1.66%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None] 175.98 178.89 +1.65%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[10000000-cpu] 52.07 52.90 +1.58%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False] 51.87 52.69 +1.57%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 34,576 34,054 -1.51%
benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True] 55.71 54.89 -1.46%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 1,977 1,948 -1.46%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-None] 38.23 37.67 -1.45%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] 263.51 267.24 +1.41%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward] 112.95 114.53 +1.40%
benchmarks/test_collectors_benchmark.py::test_sync_preempt 16.34 16.57 +1.39%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True] 20,858 20,571 -1.38%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None] 348.40 343.73 -1.34%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[100-img_shape0-atari] 29.56 29.95 +1.32%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False] 27,636 27,997 +1.31%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same] 27.49 27.13 -1.28%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True] 18,689 18,449 -1.28%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False] 63,803 63,034 -1.21%
benchmarks/test_objectives_benchmarks.py::test_redq_speed[False-backward] 55.79 55.12 -1.19%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small] 7,060 7,144 +1.19%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True] 25.17 25.47 +1.19%
benchmarks/test_collectors_benchmark.py::test_async 17.30 17.51 +1.18%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,518 2,488 -1.17%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] 19,789 19,559 -1.16%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] 122.81 124.21 +1.14%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 28.94 28.61 -1.14%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False] 56,909 57,547 +1.12%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari] 25.84 26.12 +1.10%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 162.74 164.50 +1.08%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False] 54,394 54,972 +1.06%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward] 518.11 512.69 -1.05%
benchmarks/test_collectors_benchmark.py::test_single 8.8496 8.7579 -1.04%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True] 32,522 32,854 +1.02%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] 22,142 21,922 -1.00%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False] 31,582 31,895 +0.99%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant] 2,640 2,614 -0.98%
benchmarks/test_envs_benchmark.py::test_serial 0.5705 0.5761 +0.98%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] 47.80 48.27 +0.98%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch] 13.18 13.31 +0.96%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] 6,988 6,921 -0.95%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] 4,852 4,897 +0.94%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None] 82.80 83.57 +0.93%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari] 272.44 269.95 -0.91%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same] 19.21 19.03 -0.91%
benchmarks/test_collectors_benchmark.py::test_sync 16.79 16.64 -0.90%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-None] 122.33 121.22 -0.90%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None] 681.25 687.33 +0.89%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None] 703.92 697.65 -0.89%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward] 32.42 32.70 +0.87%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False] 35,019 34,718 -0.86%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] 8.1217 8.1908 +0.85%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False] 37,937 38,251 +0.83%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-lstm] 3.1055 3.1302 +0.80%
... ... ... Showing 120 of 192 comparisons, sorted by absolute change.

GPU

Compared 202 benchmarks. Regressions over 5%: 10. Improvements over 5%: 12.

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 1,025 2,752 +168.41%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None] 96.92 120.59 +24.43%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,734 2,848 -23.73%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,943 3,455 +17.42%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3,002 3,494 +16.38%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,235 3,737 +15.52%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 44.41 50.57 +13.87%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 843.07 733.18 -13.03%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 3,143 2,734 -13.03%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] 3,738 4,183 +11.91%
benchmarks/test_collectors_benchmark.py::test_single_with_rb_pixels 5.2913 4.7012 -11.15%
benchmarks/test_collectors_benchmark.py::test_sync_pixels 9.3012 10.31 +10.89%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,116 2,340 +10.59%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 787.79 710.21 -9.85%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3,385 3,698 +9.24%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,032 1,851 -8.94%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,100 1,915 -8.82%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 770.73 830.87 +7.80%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img] 574.78 538.47 -6.32%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512] 48.52 45.55 -6.13%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward] 312.02 329.33 +5.55%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 510.55 483.16 -5.37%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] 134.95 141.18 +4.61%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 1,336 1,279 -4.26%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 479.57 499.50 +4.15%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None] 217.76 226.32 +3.93%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru] 21.61 22.45 +3.90%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True] 1.3130 1.3635 +3.84%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 1,303 1,255 -3.73%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False] 49,860 51,659 +3.61%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 165.00 170.78 +3.50%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] 1,859 1,924 +3.50%
benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False] 12.10 11.69 -3.37%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None] 44.80 46.23 +3.20%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,046 1,983 -3.07%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None] 800.28 824.30 +3.00%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 970.68 999.59 +2.98%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward] 380.00 368.75 -2.96%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True] 28,639 27,802 -2.92%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True] 21.91 22.50 +2.69%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False] 41,706 42,829 +2.69%
benchmarks/test_collectors_benchmark.py::test_single 6.7184 6.5381 -2.68%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 2,869 2,946 +2.66%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm] 20.69 21.23 +2.62%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward] 77.17 79.18 +2.61%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward] 436.82 447.69 +2.49%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle] 12,406 12,101 -2.46%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None] 108.97 111.64 +2.45%
benchmarks/test_envs_benchmark.py::test_transformed 0.7082 0.6909 -2.45%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False] 38,263 37,331 -2.44%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 502.17 514.02 +2.36%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False] 55,430 54,126 -2.35%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 160.69 164.26 +2.23%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb_cuda[100-img_shape0-atari] 17.07 17.45 +2.22%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward] 123.37 126.10 +2.21%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] 801.25 818.55 +2.16%
benchmarks/test_envs_benchmark.py::test_simple 1.2102 1.1842 -2.15%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch] 707.18 722.29 +2.14%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True] 23.30 23.79 +2.07%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward] 70.01 71.46 +2.07%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img] 386.90 394.82 +2.05%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] 22,720 22,261 -2.02%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True] 20,702 20,328 -1.81%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True] 18,885 18,544 -1.80%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[200-img_shape1-large_batch] 8.1582 8.3033 +1.78%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] 19,905 19,553 -1.77%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 34,820 34,206 -1.76%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False] 64,094 62,978 -1.74%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] 7,157 7,032 -1.74%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False] 53.66 52.73 -1.73%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small] 3,502 3,563 +1.72%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] 8.5387 8.6829 +1.69%
benchmarks/test_envs_benchmark.py::test_parallel 0.5294 0.5383 +1.69%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 194.69 197.93 +1.66%
benchmarks/test_objectives_benchmarks.py::test_values[vec_td_lambda_return_estimate-True-False] 850.20 836.25 -1.64%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None] 640.69 630.33 -1.62%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-None] 52.73 53.58 +1.61%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 972.71 987.65 +1.54%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[100-img_shape0-atari] 16.31 16.56 +1.53%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-False] 0.2247 0.2214 -1.49%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] 63,990 63,041 -1.48%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb_cuda[200-img_shape1-large_batch] 8.5148 8.6391 +1.46%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] 305.81 310.24 +1.45%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] 1,866 1,892 +1.44%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy] 347,269 352,259 +1.44%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 164.73 167.08 +1.43%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False] 27,507 27,900 +1.43%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] 161.69 163.98 +1.42%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True] 18,493 18,232 -1.41%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None] 863.11 875.13 +1.39%
benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True] 297.21 301.22 +1.35%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True] 23,994 23,673 -1.34%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True] 30,487 30,081 -1.33%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward] 219.56 216.65 -1.33%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] 21,836 22,126 +1.33%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False] 32,034 31,614 -1.31%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] 401.31 406.45 +1.28%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small] 6,036 5,959 -1.27%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] 30,806 30,421 -1.25%
benchmarks/test_collectors_benchmark.py::test_single_with_rb 5.9299 5.8560 -1.25%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None] 712.42 721.20 +1.23%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] 235.16 238.05 +1.23%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] 355.98 351.72 -1.20%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False] 34,433 34,838 +1.18%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant] 4,709 4,764 +1.17%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] 344.43 348.30 +1.12%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None] 815.39 824.39 +1.10%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[50-img_shape0-small] 869.44 878.75 +1.07%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True] 23.08 23.32 +1.05%
benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True] 48.81 48.32 -1.01%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] 48.73 49.21 +0.99%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari] 654.39 660.85 +0.99%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] 167.99 169.61 +0.96%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True] 0.5210 0.5260 +0.96%
benchmarks/test_collectors_benchmark.py::test_async 11.02 10.92 -0.94%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None] 608.96 614.59 +0.92%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True] 19,793 19,611 -0.92%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] 266.22 268.60 +0.89%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward] 69.21 68.59 -0.89%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False] 34,565 34,262 -0.88%
... ... ... Showing 120 of 202 comparisons, sorted by absolute change.

vmoens added 2 commits June 17, 2026 18:55
@vmoens
Update
02352a1 
[ghstack-poisoned]
@vmoens
Update
a8215d4 
[ghstack-poisoned]
vmoens added a commit that referenced this pull request Jun 18, 2026
@vmoens
[Feature] Add GRPO candidate selection and token advantages
1a87e90 
ghstack-source-id: ad56c75
Pull-Request: #3875
@vmoens
vmoens merged commit a8215d4 into gh/vmoens/287/base Jun 18, 2026
109 of 110 checks passed
@vmoens
vmoens deleted the gh/vmoens/287/head branch June 18, 2026 13:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Documentation Improvements or additions to documentation Feature New feature llm/ LLM-related PR, triggers LLM CI tests Objectives Trainers

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vmoens