[Feature] Add RewardModelLoss objective for RLHF reward-model training by theap06 · Pull Request #3922 · pytorch/rl

theap06 · 2026-06-30T07:19:04Z

What

Adds 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, a TensorClass output, runnable docstring with paper references).

Why

The Bradley-Terry reward loss already existed as GPT2RewardModel.compute_reward_loss (torchrl/modules/models/llm.py), explicitly tagged # TODO: move to objectives and tied to GPT2. This promotes it to a proper objective that:

lives in objectives/ like every other loss;
is model-agnostic — it owns a swappable score_network (any TensorDictModule producing a per-sequence scalar, e.g. AutoModelForSequenceClassification(num_labels=1) or GPT2RewardModel);
is TensorDict-first with configurable keys.

torchrl/objectives/llm/reward.py: RewardModelLoss, RewardModelLossOutput, and the reward_model_loss helper. Supports mean/sum/none reduction, an optional score-centering regularizer, and a detached accuracy metric for logging.
Exports 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, includes a nested-key input case per the contributor guide).

Tests

pytest test/llm/test_llm_objectives.py -k RewardModel → 9 passed. ruff + black clean; the public-API __all__ guard passes.

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>

pytorch-bot · 2026-06-30T07:19:07Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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 · 2026-06-30T10:40:59Z

Benchmark Results: PR `9c2cfb2a` vs main `74abd6c3`

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

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

CPU

Compared 216 benchmarks. Regressions over 5%: 11. Improvements over 5%: 18.

Benchmark	main ops	PR ops	Change
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]`	48.44	194.75	+302.00%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]`	190.95	38.01	-80.10%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]`	187.42	55.64	-70.32%
`benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False]`	4,833	7,723	+59.80%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	829.46	992.81	+19.69%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	2,955	3,495	+18.28%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same]`	28.89	23.93	-17.16%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,749	3,196	+16.24%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	359.16	414.90	+15.52%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	2,655	3,047	+14.76%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	459.22	523.59	+14.02%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	731.30	831.15	+13.65%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same]`	20.17	22.65	+12.30%
`benchmarks/test_collectors_benchmark.py::test_single_with_rb`	7.3806	6.5132	-11.75%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64]`	10.58	9.3800	-11.33%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,611	2,906	+11.29%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	697.18	773.10	+10.89%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,623	3,255	-10.16%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,504	3,187	-9.06%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]`	466.86	507.58	+8.72%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	3,368	3,093	-8.16%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	844.47	910.55	+7.82%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors]`	22,000	23,596	+7.26%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle]`	12,366	11,498	-7.02%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	6.6916	6.2538	-6.54%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None]`	657.17	697.01	+6.06%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None]`	1,728	1,827	+5.74%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save]`	7,187	6,792	-5.50%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None]`	256.10	270.04	+5.44%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-gru]`	1.3930	1.3278	-4.68%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari]`	5,163	4,931	-4.50%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-lstm]`	0.9713	0.9277	-4.49%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]`	1,038	1,083	+4.31%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]`	484.90	505.53	+4.25%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]`	1,044	1,087	+4.14%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True]`	33,036	34,354	+3.99%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm]`	2.0980	2.0145	-3.98%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward]`	57.75	60.00	+3.90%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru]`	4.3106	4.1527	-3.66%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	119.18	123.52	+3.64%
`benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True]`	97.65	94.22	-3.51%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None]`	177.45	183.62	+3.48%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-True]`	0.2116	0.2189	+3.48%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-gru]`	1.4831	1.4323	-3.43%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[reduce-overhead-None]`	218.36	225.58	+3.31%
`benchmarks/test_envs_benchmark.py::test_simple`	1.7554	1.6983	-3.25%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage]`	8.2027	7.9455	-3.14%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True]`	22,855	23,569	+3.13%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru]`	3.0141	3.1080	+3.12%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small]`	7,242	7,021	-3.05%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward]`	60.64	62.40	+2.91%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-lstm]`	3.1850	3.0924	-2.91%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	138.44	142.44	+2.89%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	2,066	2,007	-2.85%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[50-img_shape0-small]`	846.77	870.56	+2.81%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None]`	84.62	86.98	+2.78%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None]`	463.79	476.35	+2.71%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[10000000-cpu]`	54.09	52.68	-2.62%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False]`	65,254	63,557	-2.60%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None]`	567.61	582.34	+2.60%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	1,869	1,821	-2.56%
`benchmarks/test_collectors_benchmark.py::test_sync_preempt`	16.23	16.65	+2.55%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None]`	266.47	273.26	+2.55%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward]`	83.36	85.46	+2.52%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True]`	21,435	21,973	+2.51%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch]`	136.63	139.94	+2.43%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1]`	506.91	519.15	+2.41%
`benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False]`	24.89	24.30	-2.33%
`benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False]`	37.01	36.15	-2.32%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False]`	41,725	42,691	+2.32%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512]`	568.99	555.83	-2.31%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None]`	338.69	346.43	+2.29%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1]`	468.87	479.51	+2.27%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True]`	36,771	37,590	+2.23%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False]`	62,890	64,288	+2.22%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False]`	75,898	77,572	+2.21%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cpu]`	96.94	99.07	+2.19%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-lstm]`	0.8790	0.8597	-2.19%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512]`	109.45	107.08	-2.17%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari]`	268.58	274.34	+2.14%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True]`	0.5147	0.5257	+2.14%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None]`	289.31	295.38	+2.10%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True]`	24.39	24.90	+2.08%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward]`	61.86	63.13	+2.05%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64]`	4.4417	4.5299	+1.98%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False]`	31,111	31,715	+1.94%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch]`	715.86	729.64	+1.93%
`benchmarks/test_envs_benchmark.py::test_transformed`	0.8998	0.8826	-1.91%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None]`	330.04	336.22	+1.87%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None]`	550.72	560.99	+1.87%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None]`	161.09	164.03	+1.82%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward]`	87.58	89.15	+1.79%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward]`	510.40	519.40	+1.76%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	1,930	1,896	-1.75%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False]`	48.29	49.11	+1.70%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[False-None]`	95.49	93.87	-1.69%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None]`	1,720	1,749	+1.69%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward]`	133.51	135.75	+1.68%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None]`	85.27	86.67	+1.64%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward]`	242.31	246.17	+1.59%
`benchmarks/test_objectives_benchmarks.py::test_values[vec_td1_return_estimate-False-False]`	54.28	55.13	+1.57%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-64]`	1.2192	1.2383	+1.57%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-4]`	19.43	19.73	+1.56%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch]`	303.68	308.37	+1.54%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True]`	20,235	20,537	+1.49%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None]`	88.59	89.89	+1.47%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]`	157.69	159.97	+1.45%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-16]`	4.8641	4.9338	+1.43%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True]`	31,825	32,274	+1.41%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img]`	391.01	396.32	+1.36%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]`	164.06	166.26	+1.34%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-False]`	43,826	44,406	+1.32%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	161.33	163.42	+1.30%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16]`	17.85	18.08	+1.28%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari]`	684.43	693.18	+1.28%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None]`	279.40	282.91	+1.26%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False]`	33,764	34,180	+1.23%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True]`	29,880	29,527	-1.18%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False]`	34,131	34,529	+1.17%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None]`	294.00	297.40	+1.16%
...	...	...	Showing 120 of 216 comparisons, sorted by absolute change.

GPU

Compared 226 benchmarks. Regressions over 5%: 6. Improvements over 5%: 16.

Benchmark	main ops	PR ops	Change
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	1,781	503.44	-71.73%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None]`	97.97	128.36	+31.01%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]`	45.23	51.68	+14.26%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	3,053	2,686	-12.01%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,143	2,794	-11.11%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	794.10	708.37	-10.80%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari]`	3,914	4,325	+10.50%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,817	3,109	+10.36%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,623	2,879	+9.78%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	2,775	3,018	+8.75%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,595	2,813	+8.38%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	915.01	990.81	+8.28%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cpu_sampler]`	84.47	90.88	+7.58%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward]`	372.06	396.83	+6.66%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	259.53	276.13	+6.40%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	454.81	480.97	+5.75%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant]`	5,003	4,727	-5.51%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]`	483.32	508.73	+5.26%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True]`	35,756	37,634	+5.25%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16]`	42.28	44.48	+5.20%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None]`	41.18	43.26	+5.04%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same]`	5.6293	5.3469	-5.02%
`benchmarks/test_collectors_benchmark.py::test_sync_pixels`	10.35	9.8386	-4.94%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	1,925	2,018	+4.86%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	2,708	2,839	+4.83%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None]`	426.43	407.20	-4.51%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same]`	6.7637	6.4608	-4.48%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	363.99	347.85	-4.43%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None]`	616.77	592.51	-3.93%
`benchmarks/test_envs_benchmark.py::test_simple`	1.2548	1.2058	-3.91%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512]`	1,337	1,389	+3.89%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512]`	1,302	1,352	+3.84%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	847.67	815.52	-3.79%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]`	985.94	1,022	+3.70%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None]`	1,911	1,841	-3.64%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cuda]`	2,320	2,239	-3.50%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-4]`	164.81	159.08	-3.48%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False]`	60,475	62,569	+3.46%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari]`	655.15	632.78	-3.42%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward]`	155.27	150.05	-3.36%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None]`	731.68	707.17	-3.35%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,266	3,161	-3.21%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False]`	30,762	31,696	+3.04%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512]`	1,257	1,295	+3.02%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None]`	834.62	809.59	-3.00%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small]`	4,394	4,265	-2.94%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True]`	19,090	19,639	+2.87%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	474.92	488.38	+2.83%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16]`	36.23	37.26	+2.82%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True]`	41,345	40,217	-2.73%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward]`	135.39	131.73	-2.70%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]`	194.28	189.11	-2.66%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	807.56	828.88	+2.64%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small]`	6,159	5,999	-2.60%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False]`	26,588	27,267	+2.55%
`benchmarks/test_collectors_benchmark.py::test_sync_preempt`	10.44	10.18	-2.49%
`benchmarks/test_envs_benchmark.py::test_transformed`	0.7205	0.7033	-2.39%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False]`	49.09	47.92	-2.38%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	1,973	2,020	+2.36%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True]`	21,689	21,185	-2.32%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]`	167.76	171.61	+2.30%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb_cuda[200-img_shape1-large_batch]`	8.9306	8.7267	-2.28%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None]`	1,935	1,891	-2.25%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False]`	74,572	76,222	+2.21%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small]`	3,473	3,549	+2.20%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False]`	30,631	31,297	+2.17%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None]`	347.41	339.99	-2.14%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False]`	33,606	34,314	+2.11%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy]`	353,998	346,614	-2.09%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True]`	23.30	22.81	-2.08%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[100-img_shape0-atari]`	13.62	13.34	-2.06%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	7.1539	7.2989	+2.03%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch]`	741.89	726.87	-2.03%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant]`	4,761	4,665	-2.01%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None]`	103.90	105.98	+2.00%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False]`	36,763	37,494	+1.99%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64]`	12.51	12.76	+1.98%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True]`	33,622	34,285	+1.97%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True]`	17,572	17,917	+1.96%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True]`	22.36	21.92	-1.95%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True]`	23.93	23.47	-1.94%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	162.01	165.15	+1.94%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward]`	79.54	81.06	+1.90%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[200-img_shape1-large_batch]`	6.8017	6.6738	-1.88%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True]`	19,065	19,421	+1.87%
`benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False]`	12,135	11,912	-1.83%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-None]`	113.27	111.21	-1.82%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru]`	50.28	49.37	-1.81%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True]`	31,684	32,235	+1.74%
`benchmarks/test_collectors_benchmark.py::test_async`	10.85	11.04	+1.74%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True]`	27,805	28,273	+1.68%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage]`	8.8933	8.7510	-1.60%
`benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True]`	318.49	313.41	-1.60%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False]`	48,137	48,895	+1.57%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-16]`	48.86	49.63	+1.57%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False]`	53,828	52,994	-1.55%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64]`	10.82	10.99	+1.50%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None]`	835.43	823.05	-1.48%
`benchmarks/test_collectors_benchmark.py::test_async_pixels`	10.80	10.64	-1.45%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-False]`	0.6063	0.5975	-1.44%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle]`	12,248	12,422	+1.42%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True]`	20,371	20,087	-1.40%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None]`	839.29	850.92	+1.39%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb_cuda[100-img_shape0-atari]`	17.87	17.63	-1.35%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,350	2,319	-1.33%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward]`	224.43	227.41	+1.33%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False]`	48,123	48,755	+1.31%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch]`	333.73	329.40	-1.30%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1]`	282.90	286.56	+1.30%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False]`	56,260	55,540	-1.28%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward]`	352.92	357.41	+1.27%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None]`	87.76	88.88	+1.27%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save]`	7,240	7,148	-1.27%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None]`	374.92	370.26	-1.24%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None]`	102.22	103.39	+1.15%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm]`	22.18	21.93	-1.14%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img]`	426.78	431.57	+1.12%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-None]`	54.70	54.10	-1.10%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img]`	174.19	172.28	-1.09%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward]`	276.93	279.89	+1.07%
...	...	...	Showing 120 of 226 comparisons, sorted by absolute change.

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 30, 2026

github-actions Bot added Documentation Improvements or additions to documentation Objectives llm/ LLM-related PR, triggers LLM CI tests Feature New feature labels Jun 30, 2026

theap06 mentioned this pull request Jun 30, 2026

[Algorithm] Add RLHF reward-model training recipe #3923

Open

theap06 marked this pull request as draft June 30, 2026 07:22

theap06 marked this pull request as ready for review June 30, 2026 07:30

vmoens added 2 commits June 30, 2026 12:59

[Feature] Address RewardModelLoss review feedback

9c2cfb2

[Quality] Apply RewardModelLoss lint formatting

449174f

theap06 mentioned this pull request Jul 2, 2026

[Feature] Add Learner primitive (LocalLearner, FSDP2Learner) #3926

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature] Add RewardModelLoss objective for RLHF reward-model training#3922

[Feature] Add RewardModelLoss objective for RLHF reward-model training#3922
theap06 wants to merge 3 commits into
pytorch:mainfrom
theap06:reward-model-loss

theap06 commented Jun 30, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Jun 30, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 30, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

theap06 commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Why

Contents

Tests

Uh oh!

pytorch-bot Bot commented Jun 30, 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/3922

Uh oh!

github-actions Bot commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmark Results: PR 9c2cfb2a vs main 74abd6c3

CPU

GPU

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

theap06 commented Jun 30, 2026 •

edited

Loading

pytorch-bot Bot commented Jun 30, 2026 •

edited

Loading

github-actions Bot commented Jun 30, 2026 •

edited

Loading

Benchmark Results: PR `9c2cfb2a` vs main `74abd6c3`