[Algorithm] Add RLHF reward-model training recipe#3923
Open
theap06 wants to merge 4 commits into
Open
Conversation
Add a model-agnostic Bradley-Terry pairwise reward-model loss as a first-class objective under torchrl/objectives/llm/, following the SFTLoss/GRPOLoss conventions (_AcceptedKeys/set_keys, TensorClass output, runnable docstring with paper references). This resolves the "# TODO: move to objectives" left on GPT2RewardModel.compute_reward_loss: the loss now lives in objectives and owns a swappable score_network (any TensorDictModule producing a per-sequence scalar, e.g. AutoModelForSequenceClassification(num_labels=1) or GPT2RewardModel). It supports mean/sum/none reduction, an optional score-centering regularizer, and a detached accuracy metric for logging. - torchrl/objectives/llm/reward.py: RewardModelLoss, RewardModelLossOutput and the reward_model_loss helper. - Export both from torchrl/objectives/llm/__init__.py. - Docs: Reward Model Training sections in llms_objectives.rst and llms.rst. - Tests: TestRewardModel in test/llm/test_llm_objectives.py (CPU-only, exercises a nested-key input per the contributor guide). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a dedicated sota-implementations recipe that trains a scalar reward model from pairwise human-preference data using RewardModelLoss. This modernizes the legacy examples/rlhf/train_reward.py into the current sota-implementations and LLM-API style, and is model-agnostic: any Hugging Face AutoModelForSequenceClassification backbone can be used. - sota-implementations/reward_model_training/: reward_model.py (Hydra main), utils.py, config.yaml, requirements.txt, README.md. - Hermetic CI path: an empty model.name builds a tiny from-config model and an empty data.dataset_name generates synthetic preference pairs, so the smoke test needs no download and no `datasets` dependency (which the linux_sota env does not provide). - CI: add a reward_model_training smoke entry to .github/unittest/linux_sota/scripts/test_sota.py. - sota-check/run_reward_model_training.sh for full release runs. Depends on the RewardModelLoss objective (separate PR); the recipe imports it and the smoke test exercises it end to end. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3923
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
github-actions
Bot
added
Documentation
meta-cla
Bot
added
the
CLA Signed
Contributor
Benchmark Results: PR
|
| Benchmark | main ops | PR ops | Change |
|---|---|---|---|
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] |
48.27 | 196.65 | +307.39% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] |
195.10 | 36.26 | -81.41% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] |
186.20 | 56.49 | -69.66% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64] |
10.93 | 8.5178 | -22.05% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16] |
37.21 | 29.14 | -21.70% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] |
3,435 | 2,713 | -21.03% |
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward] |
852.36 | 990.80 | +16.24% |
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same] |
29.18 | 24.68 | -15.41% |
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same] |
20.25 | 23.28 | +14.95% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64] |
12.58 | 10.96 | -12.88% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] |
3,299 | 3,664 | +11.06% |
benchmarks/test_collectors_benchmark.py::test_single_with_rb |
7.3840 | 6.6205 | -10.34% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] |
2,980 | 2,678 | -10.14% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] |
2,124 | 1,914 | -9.85% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] |
3,323 | 3,639 | +9.49% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] |
789.62 | 716.80 | -9.22% |
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] |
22,590 | 24,360 | +7.84% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64] |
6.5785 | 6.0691 | -7.74% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] |
2,984 | 2,757 | -7.60% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] |
2,008 | 2,153 | +7.23% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] |
3,002 | 2,813 | -6.28% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] |
527.70 | 560.66 | +6.24% |
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] |
1,710 | 1,813 | +6.00% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward] |
110.31 | 116.80 | +5.88% |
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None] |
323.45 | 341.74 | +5.65% |
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True] |
1.3294 | 1.3940 | +4.86% |
benchmarks/test_objectives_benchmarks.py::test_redq_speed[True-None] |
224.07 | 234.95 | +4.86% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True] |
30,909 | 29,434 | -4.77% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] |
2,081 | 1,982 | -4.75% |
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] |
8.5345 | 8.1331 | -4.70% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] |
2,138 | 2,239 | +4.70% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] |
1,063 | 1,014 | -4.66% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] |
176.30 | 168.16 | -4.61% |
benchmarks/test_objectives_benchmarks.py::test_redq_speed[reduce-overhead-None] |
227.68 | 238.17 | +4.61% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] |
1,071 | 1,025 | -4.25% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] |
837.51 | 803.84 | -4.02% |
benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False] |
7,726 | 7,426 | -3.88% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari] |
691.99 | 718.69 | +3.86% |
benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None] |
115.00 | 119.29 | +3.74% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] |
30,913 | 29,866 | -3.39% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False] |
51.74 | 53.48 | +3.35% |
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward] |
63.50 | 61.44 | -3.25% |
benchmarks/test_collectors_benchmark.py::test_sync |
16.86 | 16.31 | -3.21% |
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-memmap_cpu_storage_cpu... |
80.82 | 83.39 | +3.18% |
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cpu] |
96.11 | 99.06 | +3.07% |
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] |
140.07 | 135.94 | -2.94% |
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] |
59.80 | 61.53 | +2.90% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] |
3,633 | 3,738 | +2.90% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False] |
49,261 | 50,685 | +2.89% |
benchmarks/test_envs_benchmark.py::test_parallel |
0.9285 | 0.9551 | +2.87% |
benchmarks/test_envs_benchmark.py::test_simple |
1.7691 | 1.7189 | -2.84% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch] |
330.38 | 339.72 | +2.83% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None] |
259.79 | 267.09 | +2.81% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] |
878.18 | 902.64 | +2.79% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True] |
38,528 | 37,507 | -2.65% |
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] |
119.04 | 122.18 | +2.65% |
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None] |
273.27 | 280.46 | +2.63% |
benchmarks/test_envs_benchmark.py::test_transformed |
0.9218 | 0.8977 | -2.61% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True] |
23.53 | 24.14 | +2.59% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] |
481.66 | 469.17 | -2.59% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True] |
20,630 | 20,102 | -2.56% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False] |
38,182 | 39,158 | +2.56% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari] |
273.85 | 280.77 | +2.53% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] |
309.04 | 316.80 | +2.51% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000] |
163.05 | 166.99 | +2.41% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False] |
29,064 | 29,764 | +2.41% |
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None] |
288.97 | 295.88 | +2.39% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] |
36,768 | 37,632 | +2.35% |
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None] |
342.32 | 350.27 | +2.32% |
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512] |
638.60 | 653.13 | +2.28% |
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch] |
10.43 | 10.19 | -2.27% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] |
531.84 | 519.77 | -2.27% |
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant] |
4,346 | 4,443 | +2.23% |
benchmarks/test_envs_benchmark.py::test_serial |
0.5736 | 0.5864 | +2.22% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True] |
24.46 | 25.00 | +2.18% |
benchmarks/test_collectors_benchmark.py::test_sync_preempt |
16.76 | 16.40 | -2.17% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] |
53.32 | 54.44 | +2.09% |
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None] |
549.78 | 561.21 | +2.08% |
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] |
1,819 | 1,857 | +2.08% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] |
140.54 | 143.41 | +2.04% |
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[10000000-cpu] |
53.49 | 52.41 | -2.02% |
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None] |
288.12 | 293.90 | +2.01% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] |
48.44 | 49.40 | +1.98% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None] |
158.94 | 162.09 | +1.98% |
benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True] |
100.96 | 99.00 | -1.95% |
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru] |
3.0542 | 3.1133 | +1.94% |
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None] |
89.80 | 88.09 | -1.90% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False] |
63,387 | 64,586 | +1.89% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True] |
25.52 | 25.98 | +1.83% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] |
5,150 | 5,055 | -1.83% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1] |
284.56 | 279.41 | -1.81% |
benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None] |
473.03 | 481.51 | +1.79% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] |
165.90 | 168.85 | +1.78% |
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-gru] |
1.3509 | 1.3272 | -1.76% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari] |
638.58 | 649.54 | +1.72% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False] |
54,980 | 55,917 | +1.70% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False] |
41,943 | 42,654 | +1.70% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] |
266.20 | 270.56 | +1.64% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small] |
3,541 | 3,485 | -1.60% |
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru] |
4.2813 | 4.3488 | +1.58% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small] |
4,443 | 4,375 | -1.54% |
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward] |
281.55 | 285.87 | +1.53% |
benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-None] |
37.86 | 38.44 | +1.53% |
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm] |
2.0192 | 2.0499 | +1.52% |
benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None] |
566.66 | 574.67 | +1.41% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True] |
32,407 | 32,859 | +1.40% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16] |
43.60 | 44.20 | +1.38% |
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True] |
0.5194 | 0.5265 | +1.37% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False] |
34,495 | 34,961 | +1.35% |
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512] |
2,241 | 2,270 | +1.32% |
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy] |
372,800 | 377,459 | +1.25% |
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-False] |
0.2254 | 0.2226 | -1.23% |
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward] |
241.99 | 244.91 | +1.21% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True] |
28,862 | 28,522 | -1.18% |
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward] |
133.69 | 132.15 | -1.15% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4] |
72.05 | 71.22 | -1.15% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] |
65,265 | 64,522 | -1.14% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True] |
19,474 | 19,694 | +1.13% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] |
160.05 | 161.83 | +1.11% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] |
21,925 | 22,165 | +1.10% |
| ... | ... | ... | Showing 120 of 216 comparisons, sorted by absolute change. |
GPU
Compared 226 benchmarks. Regressions over 5%: 10. Improvements over 5%: 15.
| Benchmark | main ops | PR ops | Change |
|---|---|---|---|
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] |
44.28 | 188.22 | +325.10% |
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-None] |
53.12 | 99.37 | +87.06% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] |
195.36 | 48.67 | -75.09% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] |
1,009 | 726.93 | -27.95% |
benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None] |
105.89 | 77.95 | -26.38% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] |
1,006 | 757.70 | -24.69% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] |
3,421 | 2,784 | -18.61% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] |
3,647 | 3,090 | -15.28% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] |
3,022 | 2,584 | -14.49% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] |
3,365 | 2,898 | -13.89% |
benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True] |
325.13 | 286.13 | -11.99% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch] |
781.84 | 708.29 | -9.41% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] |
4,041 | 4,364 | +7.99% |
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward] |
268.99 | 289.75 | +7.72% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img] |
404.73 | 435.93 | +7.71% |
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] |
347.66 | 374.02 | +7.58% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward] |
335.51 | 360.39 | +7.42% |
benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True] |
46.60 | 49.83 | +6.92% |
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward] |
217.30 | 231.08 | +6.34% |
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None] |
748.80 | 792.14 | +5.79% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None] |
222.81 | 235.66 | +5.77% |
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru] |
46.92 | 49.52 | +5.56% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False] |
74,689 | 78,632 | +5.28% |
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True] |
1.3060 | 1.3730 | +5.13% |
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward] |
148.80 | 156.25 | +5.01% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] |
756.70 | 720.74 | -4.75% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] |
2,006 | 2,100 | +4.67% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] |
132.28 | 138.36 | +4.60% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True] |
22,483 | 23,514 | +4.58% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] |
2,876 | 2,747 | -4.51% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] |
167.59 | 174.97 | +4.40% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True] |
31,052 | 32,390 | +4.31% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] |
481.59 | 501.69 | +4.17% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] |
477.65 | 496.97 | +4.04% |
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] |
1,838 | 1,911 | +3.99% |
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-None] |
387.87 | 402.55 | +3.78% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img] |
565.49 | 544.16 | -3.77% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] |
2,066 | 1,988 | -3.76% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False] |
54,039 | 55,999 | +3.63% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-4] |
143.10 | 148.23 | +3.59% |
benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512] |
47.91 | 49.59 | +3.50% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] |
21,373 | 22,091 | +3.36% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16] |
17.61 | 18.20 | +3.34% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] |
47.75 | 49.33 | +3.32% |
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant] |
4,872 | 5,034 | +3.32% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] |
19,041 | 19,660 | +3.25% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] |
1,905 | 1,967 | +3.25% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] |
35,861 | 37,010 | +3.20% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64] |
12.38 | 12.78 | +3.18% |
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None] |
843.55 | 869.89 | +3.12% |
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None] |
337.00 | 347.46 | +3.11% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] |
433.51 | 446.88 | +3.08% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] |
52.70 | 54.28 | +3.01% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward] |
133.00 | 136.97 | +2.98% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64] |
10.69 | 11.00 | +2.96% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-1] |
185.29 | 190.64 | +2.89% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64] |
4.4342 | 4.5620 | +2.88% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True] |
36,128 | 37,132 | +2.78% |
benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False] |
12.31 | 12.65 | +2.76% |
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None] |
817.79 | 840.32 | +2.76% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True] |
27,976 | 28,738 | +2.73% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4] |
70.17 | 72.08 | +2.72% |
benchmarks/test_envs_benchmark.py::test_simple |
1.2432 | 1.2094 | -2.72% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] |
3,332 | 3,419 | +2.61% |
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] |
836.76 | 858.47 | +2.60% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True] |
23.83 | 23.22 | -2.58% |
benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False] |
11,842 | 12,131 | +2.44% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] |
33,455 | 34,270 | +2.44% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] |
3,585 | 3,671 | +2.42% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True] |
41,194 | 42,190 | +2.42% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1] |
277.71 | 284.34 | +2.39% |
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm] |
21.91 | 21.39 | -2.38% |
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] |
363.42 | 371.99 | +2.36% |
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward] |
240.56 | 246.02 | +2.27% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small] |
6,095 | 5,956 | -2.27% |
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy] |
361,830 | 353,804 | -2.22% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False] |
27,340 | 27,944 | +2.21% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False] |
35,709 | 34,927 | -2.19% |
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru] |
23.16 | 22.65 | -2.18% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] |
63,334 | 64,686 | +2.13% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False] |
32,638 | 31,955 | -2.09% |
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None] |
641.20 | 654.57 | +2.09% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64] |
7.1453 | 7.2938 | +2.08% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16] |
36.23 | 36.98 | +2.07% |
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-False] |
0.6007 | 0.5883 | -2.05% |
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward] |
329.13 | 335.85 | +2.04% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] |
29,561 | 30,159 | +2.02% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True] |
20,068 | 20,470 | +2.01% |
benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False] |
20.37 | 20.77 | +1.95% |
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None] |
510.28 | 520.22 | +1.95% |
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward] |
457.78 | 466.47 | +1.90% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] |
192.28 | 195.79 | +1.83% |
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward] |
69.27 | 70.53 | +1.82% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False] |
34,960 | 34,326 | -1.81% |
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] |
1,933 | 1,898 | -1.79% |
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward] |
480.14 | 488.41 | +1.72% |
benchmarks/test_envs_benchmark.py::test_serial |
0.4241 | 0.4169 | -1.69% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] |
2,878 | 2,830 | -1.69% |
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle] |
12,080 | 12,283 | +1.68% |
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None] |
419.69 | 426.73 | +1.68% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] |
165.88 | 168.67 | +1.68% |
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] |
7,292 | 7,173 | -1.63% |
benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-None] |
112.12 | 113.93 | +1.62% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False] |
39,071 | 38,445 | -1.60% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch] |
328.63 | 323.44 | -1.58% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1] |
509.13 | 517.01 | +1.55% |
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] |
274.10 | 278.34 | +1.54% |
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] |
24,016 | 24,367 | +1.46% |
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] |
712.07 | 722.45 | +1.46% |
benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward] |
80.69 | 81.87 | +1.45% |
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-4] |
47.44 | 48.12 | +1.44% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False] |
31,582 | 32,017 | +1.38% |
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None] |
736.77 | 726.67 | -1.37% |
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True] |
20,381 | 20,659 | +1.36% |
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None] |
741.62 | 751.62 | +1.35% |
benchmarks/test_collectors_benchmark.py::test_sync |
10.55 | 10.41 | -1.32% |
benchmarks/test_collectors_benchmark.py::test_single_pixels |
6.3131 | 6.3923 | +1.25% |
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-True] |
0.2114 | 0.2140 | +1.25% |
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small] |
3,479 | 3,435 | -1.25% |
benchmarks/test_collectors_benchmark.py::test_single |
6.7705 | 6.8549 | +1.25% |
| ... | ... | ... | Showing 120 of 226 comparisons, sorted by absolute change. |
vmoens
reviewed
Jun 30, 2026
vmoens
left a comment
vmoens
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for putting this together. I think the RewardModelLoss itself is a useful TorchRL primitive, good pieces to have in-tree
My main concern is the positioning of the recipe. Moving forward I'd like to be foreseeing on how we can bring torchrl to a state where we can do serious distributed training for LLM features. As written, this is a single-process training loop, with GPT-2 as the visible example. Fine for a tuto/smoke recipe, but it is not really a serious reward-model training stack for modern LLMs. I guess that, at scale, users will still want TRL/Accelerate/FSDP, TorchTitan/FSDP2, Megatron/NeMo, OpenRLHF, etc. for parallelism, sharded checkpointing, fault tolerance, packed/streamed datasets, mixed precision, activation checkpointing, and multi-node orchestration. So the end goal should be to interface with these.
Don't get me wrong, I do not think this PR should grow a Megatron/NeMo/TorchTitan integration. That would be too much scope. But I think we should be very explicit about what TorchRL is providing here:
- a canonical TensorDict preference-data format,
- a canonical reward-model loss,
- a small reference trainer,
- and eventually adapters/scorers that let externally trained reward models plug back into TorchRL GRPO/PPO/SFT pipelines.
Concretely, I would suggest adding a short README section along the lines of “Scaling and integration”. It could say that this recipe is a minimal single-node baseline, and that large-scale RM training should use an external backend while preserving the TorchRL data/scoring contract. Essentially we need to pave the way for the real stuff that is to come. If you're happy to work on this with me we can start a topic in discussion.
I would also avoid presenting GPT-2 as the “real” default. Keeping the tiny synthetic model for CI is great, but the docs/examples should probably use a small Qwen model.
So overall: I like the objective and the idea of a reference recipe. I would just make the scope honest: TorchRL should own the data contract, loss semantics, and scoring interface. Overall, we need to think about how to scale this (and the rest of the LLM stack) up.
… a small Qwen Addresses reviewer feedback on the RLHF reward-model recipe: - Add a 'Scaling and integration' README section making the scope explicit: TorchRL owns the TensorDict preference-data format, RewardModelLoss semantics, a small reference trainer, and (future) adapters/scorers into GRPO/PPO/SFT; large-scale RM training is delegated to external backends (TRL/Accelerate/FSDP, TorchTitan/FSDP2, Megatron/NeMo, OpenRLHF) while preserving the data/scoring contract. - Use Qwen/Qwen2.5-0.5B as the example default in the config and docs instead of gpt2. The tiny from-scratch synthetic model used by CI is unchanged.
Contributor
Author
|
@vmoens Thanks, this framing makes more sense! On the bigger picture, I fully agree the end goal is the adapters/scorers layer and interfacing with external backends. The way I'm picturing it: TensorDict is the shared format that holds the preference data, rollouts, scores, and rewards, and the backends (TRL/Accelerate/FSDP, etc.) bolt on at the edges. They consume the exported preference data, do the heavy distributed training, and hand back a checkpoint we re-import as a scorer. That way GRPO/SFT/reward all just read and write the same keys, and the parallelism can be done with the backends. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds a dedicated
sota-implementations/reward_model_training/recipe that trains a scalar reward model from pairwise human-preference data usingRewardModelLoss. This modernizes the legacyexamples/rlhf/train_reward.py(originally from #1597) into the currentsota-implementations+ LLM-API style, and is model-agnostic: any HFAutoModelForSequenceClassificationbackbone works.Stacked on #3922
This PR depends on #3922 (the
RewardModelLossobjective) and is stacked on top of it — the recipe importsRewardModelLossand the CI smoke test exercises it end to end. Please merge #3922 first; once it lands, the diff here reduces to just the recipe + CI files (the first commit shown belongs to #3922).Contents
sota-implementations/reward_model_training/:reward_model.py(Hydra main),utils.py,config.yaml,requirements.txt,README.md.model.namebuilds a tiny from-config model and an emptydata.dataset_namegenerates synthetic preference pairs, so the smoke test needs no download and no HFdatasetsdependency (which thelinux_sotaenv does not provide).reward_model_trainingsmoke entry in.github/unittest/linux_sota/scripts/test_sota.py.sota-check/run_reward_model_training.shfor full release runs.Tests
Ran the exact CI smoke command locally: builds a tiny model, trains on synthetic pairs, evaluates (val loss + accuracy), and checkpoints — exit 0.
ruff+blackclean.Real usage:
python reward_model.py model.name=gpt2 data.dataset_name=CarperAI/openai_summarize_comparisons.