[Feature] On-policy KL distillation loss for LLMs by theap06 · Pull Request #3921 · pytorch/rl

theap06 · 2026-06-30T00:47:53Z

##Summary

Implements DistillationLoss, an on-policy KL knowledge-distillation objective
for LLM policies, under torchrl/objectives/llm/.
A student policy is distilled toward a frozen, typically stronger teacher by
minimizing reverse KL(student ‖ teacher) on the student's own tokens, restricted
to assistant (response) tokens. It mirrors SFTLoss and reuses the same teacher
log-prob plumbing as GRPOLoss/SFTLoss.

Add DistillationLoss, an on-policy reverse-KL knowledge-distillation objective for LLM policies, mirroring SFTLoss. The student is distilled toward a frozen (typically stronger) teacher by minimizing KL(student || teacher) on the student's own tokens, restricted to assistant tokens. Teacher per-token log-probs are read from a tensordict key produced upstream by RetrieveLogProb, the same pattern GRPO/SFT use for reference log-probs. - distillation.py: DistillationLoss, DistillationLossOutput, and the reverse_kl_token_estimate / distillation_loss helpers (Schulman k3 estimator). - Register in objectives/llm/__init__.py; document in docs/source/reference/llms_objectives.rst. - Tests: model-free unit tests for the estimator and reduction, plus transformers-gated integration tests mirroring TestSFT. Reverse-KL, token-log-probability form. Forward-KL / JSD variants need the teacher's full vocabulary distribution and are left as follow-ups. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

pytorch-bot · 2026-06-30T00:47:57Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 3 New Failures

As of commit b89c142 with merge base c8a3e37 ():

NEW FAILURES - The following jobs have failed:

Build M1 Wheels / pytorch/rl / build-wheel-py3_10-cpu (gh)
Process completed with exit code 1.
Build M1 Wheels / pytorch/rl / upload / upload-wheel-py3_10-cpu (gh)
Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.10_cpu_
Unit-tests on Linux / tests-optdeps (3.12, 13.0) / linux-job (gh)
RuntimeError: Command docker exec -t 925773ab10cdeaeac4a1426ab125b57eabedc9cc64b6cc6427697afe7b522fac /exec failed with exit code 2

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

github-actions · 2026-06-30T10:46:44Z

Benchmark Results: PR `b89c1420` vs main `74abd6c3`

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

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

CPU

Compared 216 benchmarks. Regressions over 5%: 22. Improvements over 5%: 27.

Benchmark	main ops	PR ops	Change
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]`	36.08	199.27	+452.34%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,228	425.47	-80.90%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]`	187.94	53.10	-71.75%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]`	919.37	1,117	+21.45%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16]`	35.76	28.32	-20.80%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	827.69	987.39	+19.29%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,055	2,477	-18.94%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-None]`	182.89	216.66	+18.46%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,129	3,705	+18.42%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	2,208	1,845	-16.46%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-32-512]`	39.12	33.16	-15.24%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None]`	157.53	181.48	+15.20%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same]`	29.29	25.02	-14.57%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same]`	20.21	23.10	+14.28%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,839	2,440	-14.06%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	805.33	692.49	-14.01%
`benchmarks/test_collectors_benchmark.py::test_single`	7.7939	8.8587	+13.66%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,091	1,808	-13.52%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-None]`	45.16	51.22	+13.42%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	755.80	662.16	-12.39%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari]`	5,217	4,603	-11.76%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	7.0067	6.1900	-11.66%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	2,035	1,801	-11.50%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64]`	10.54	9.3748	-11.07%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	3,147	3,495	+11.06%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None]`	147.00	163.10	+10.95%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]`	528.63	471.47	-10.81%
`benchmarks/test_collectors_benchmark.py::test_single_with_rb`	7.2264	6.4905	-10.18%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	109.69	120.49	+9.84%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,033	2,735	-9.82%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	515.51	467.05	-9.40%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward]`	30.81	33.58	+8.99%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward]`	123.01	134.02	+8.95%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward]`	78.05	84.97	+8.87%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward]`	73.05	79.52	+8.85%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]`	507.71	547.08	+7.75%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	834.74	771.81	-7.54%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	390.54	419.84	+7.50%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,999	2,777	-7.43%
`benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True]`	89.37	95.95	+7.36%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,856	3,051	+6.84%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-backward]`	26.55	28.33	+6.67%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru]`	2.8441	3.0246	+6.35%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-None]`	36.05	38.26	+6.15%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	132.09	139.87	+5.89%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64]`	12.10	11.41	-5.71%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None]`	323.00	339.82	+5.21%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None]`	326.40	343.17	+5.14%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward]`	232.43	244.12	+5.03%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True]`	1.3098	1.3734	+4.86%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-gru]`	1.3400	1.2776	-4.66%
`benchmarks/test_envs_benchmark.py::test_simple`	1.7912	1.7096	-4.56%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1]`	258.67	270.38	+4.53%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward]`	276.06	288.20	+4.40%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None]`	85.29	89.03	+4.38%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small]`	3,397	3,545	+4.37%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch]`	771.67	738.16	-4.34%
`benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False]`	7,936	7,592	-4.33%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[reduce-overhead-None]`	216.25	225.59	+4.32%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small]`	4,275	4,458	+4.29%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False]`	79,884	76,739	-3.94%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari]`	628.58	653.14	+3.91%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True]`	39,496	37,979	-3.84%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm]`	1.9085	1.9790	+3.70%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img]`	175.55	169.19	-3.62%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-4]`	19.67	20.35	+3.50%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True]`	29,129	28,127	-3.44%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None]`	117.25	121.22	+3.38%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None]`	667.41	688.97	+3.23%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward]`	60.74	62.68	+3.18%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch]`	14.57	15.03	+3.15%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False]`	46.26	47.72	+3.15%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward]`	124.48	128.34	+3.10%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img]`	403.71	391.21	-3.10%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None]`	1,766	1,821	+3.09%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors]`	22,831	23,528	+3.05%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save]`	7,129	6,913	-3.02%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-16]`	4.9456	5.0935	+2.99%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False]`	51.87	53.42	+2.98%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4]`	69.63	71.71	+2.98%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward]`	112.58	115.85	+2.90%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cpu]`	97.69	100.51	+2.89%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False]`	27,479	28,260	+2.84%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None]`	334.94	344.42	+2.83%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch]`	307.33	316.00	+2.82%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64]`	4.4008	4.5225	+2.76%
`benchmarks/test_envs_benchmark.py::test_transformed`	0.8798	0.9040	+2.75%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True]`	19,561	19,030	-2.71%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1]`	464.21	476.58	+2.66%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True]`	32,179	33,032	+2.65%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True]`	24.23	24.87	+2.62%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,071	2,991	-2.61%
`benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False]`	35.94	35.01	-2.59%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False]`	51.70	53.03	+2.57%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch]`	329.20	337.62	+2.56%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward]`	245.58	251.76	+2.52%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True]`	25.04	25.66	+2.49%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False]`	50,827	49,572	-2.47%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward]`	60.39	61.88	+2.47%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16]`	17.55	17.98	+2.45%
`benchmarks/test_envs_benchmark.py::test_serial`	0.5722	0.5862	+2.43%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[10000000-cpu]`	52.96	51.68	-2.43%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True]`	37,981	38,903	+2.43%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1]`	624.66	639.66	+2.40%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari]`	19.81	20.29	+2.40%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16]`	42.50	43.51	+2.38%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1]`	501.24	512.91	+2.33%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None]`	456.52	467.12	+2.32%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-1]`	78.54	80.30	+2.25%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-False]`	0.6046	0.5910	-2.24%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-64]`	1.2471	1.2748	+2.22%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False]`	32,099	31,401	-2.18%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None]`	116.76	119.19	+2.08%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari]`	687.23	701.00	+2.00%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None]`	293.59	287.75	-1.99%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]`	193.01	196.83	+1.98%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-None]`	119.35	121.71	+1.97%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward]`	85.49	87.17	+1.97%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward]`	499.21	508.95	+1.95%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img]`	421.53	413.48	-1.91%
...	...	...	Showing 120 of 216 comparisons, sorted by absolute change.

GPU

Compared 226 benchmarks. Regressions over 5%: 20. Improvements over 5%: 17.

Benchmark	main ops	PR ops	Change
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari]`	2,596	4,012	+54.55%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None]`	75.14	103.01	+37.10%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,636	3,317	+25.85%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	3,272	2,520	-22.99%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,282	2,652	-19.18%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	2,991	3,519	+17.64%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None]`	120.88	100.67	-16.72%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch]`	281.00	327.29	+16.47%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch]`	664.83	752.38	+13.17%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	3,183	2,817	-11.51%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,207	2,840	-11.46%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]`	993.84	889.50	-10.50%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,755	3,036	+10.19%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	1,012	916.61	-9.41%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,612	3,278	-9.22%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,301	2,111	-8.26%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1]`	462.51	424.40	-8.24%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,138	1,963	-8.21%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]`	523.69	480.73	-8.20%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	423.04	457.18	+8.07%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	2,044	2,207	+8.00%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	744.17	686.60	-7.74%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img]`	387.60	417.12	+7.61%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	829.32	770.89	-7.05%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]`	42.19	39.28	-6.91%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward]`	287.63	267.82	-6.89%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True]`	0.4946	0.5274	+6.64%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]`	997.01	932.64	-6.46%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cpu_sampler]`	84.22	89.34	+6.08%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari]`	650.69	689.94	+6.03%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True]`	1.2965	1.3698	+5.65%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False]`	31,593	33,319	+5.46%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small]`	6,185	5,854	-5.36%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small]`	4,122	4,343	+5.35%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	782.15	823.31	+5.26%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True]`	22,845	21,683	-5.09%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	374.18	355.33	-5.04%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch]`	130.25	136.13	+4.51%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle]`	11,909	12,431	+4.38%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward]`	361.08	345.70	-4.26%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False]`	56,383	58,686	+4.09%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16]`	41.08	42.72	+3.99%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru]`	54.44	52.28	-3.96%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch]`	14.55	15.11	+3.84%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	487.61	469.01	-3.82%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False]`	39,078	37,589	-3.81%
`benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True]`	501.03	482.85	-3.63%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True]`	22.28	21.49	-3.53%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant]`	5,018	4,845	-3.45%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img]`	568.43	548.91	-3.43%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward]`	152.37	147.18	-3.41%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None]`	1,948	2,014	+3.38%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True]`	32,390	33,446	+3.26%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-4]`	158.55	153.41	-3.24%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cuda]`	2,310	2,237	-3.17%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari]`	268.88	277.13	+3.07%
`benchmarks/test_envs_benchmark.py::test_simple`	1.2548	1.2169	-3.02%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-None]`	401.01	388.96	-3.01%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img]`	175.78	170.58	-2.96%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None]`	838.36	862.70	+2.90%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]`	164.64	169.25	+2.80%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]`	164.33	168.74	+2.68%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru]`	23.73	23.11	-2.61%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward]`	244.52	250.77	+2.56%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True]`	18,319	18,774	+2.48%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward]`	371.67	363.02	-2.33%
`benchmarks/test_envs_benchmark.py::test_serial`	0.4252	0.4343	+2.15%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward]`	69.09	67.61	-2.14%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None]`	279.98	274.09	-2.10%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False]`	64,267	62,932	-2.08%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False]`	28,933	29,528	+2.06%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None]`	523.39	534.09	+2.04%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True]`	19,436	19,833	+2.04%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True]`	23.61	23.13	-2.04%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[200-img_shape1-large_batch]`	6.7539	6.6177	-2.02%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-4]`	183.08	179.42	-2.00%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None]`	1,915	1,953	+1.99%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None]`	833.73	850.24	+1.98%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[100-img_shape0-atari]`	29.60	30.18	+1.95%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward]`	133.56	130.98	-1.93%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None]`	109.73	111.83	+1.91%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False]`	48,799	49,698	+1.84%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16]`	35.10	35.74	+1.81%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-True]`	0.2126	0.2164	+1.81%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False]`	76,993	78,356	+1.77%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	6.9776	7.0992	+1.74%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4]`	70.91	69.68	-1.74%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same]`	5.7322	5.6327	-1.73%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward]`	68.11	69.29	+1.73%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1]`	508.18	499.58	-1.69%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-1]`	193.92	190.79	-1.61%
`benchmarks/test_collectors_benchmark.py::test_single_with_rb_pixels`	5.2758	5.3587	+1.57%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False]`	49.07	48.30	-1.57%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	272.94	277.13	+1.54%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True]`	17,681	17,952	+1.54%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None]`	712.86	723.09	+1.44%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True]`	33,261	33,739	+1.43%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	161.44	163.71	+1.41%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage]`	8.4102	8.5272	+1.39%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-no-buffers-False]`	0.6782	0.6876	+1.38%
`benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True]`	49.91	50.59	+1.37%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None]`	443.28	449.32	+1.36%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-buffers-False]`	0.5873	0.5952	+1.34%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small]`	3,481	3,527	+1.34%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1]`	618.79	610.50	-1.34%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch]`	10.45	10.31	-1.33%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None]`	105.61	104.21	-1.33%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False]`	54,443	55,147	+1.29%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None]`	834.49	845.20	+1.28%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy]`	380,408	375,552	-1.28%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-4]`	19.76	20.01	+1.25%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward]`	239.48	236.51	-1.24%
`benchmarks/test_collectors_benchmark.py::test_async`	10.83	10.96	+1.23%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False]`	34,533	34,950	+1.21%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari]`	20.07	20.31	+1.20%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True]`	20,916	20,666	-1.19%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False]`	31,734	31,370	-1.15%
`benchmarks/test_collectors_benchmark.py::test_sync`	10.40	10.52	+1.14%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True]`	30,395	30,049	-1.14%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None]`	97.78	98.89	+1.13%
...	...	...	Showing 120 of 226 comparisons, sorted by absolute change.

…tion Adds DistillationLoss, DistillationLossOutput and the per-token reverse_kl_token_estimate k3 estimator to torchrl.objectives.llm. The loss distills a student policy toward a frozen teacher by minimizing KL(student || teacher) (on-policy, default) or KL(teacher || student) (off-policy corpora), estimated from per-token log-probs on shared tokens with the same k3 approximation used by the GRPO and SFT KL regularizers. Teacher log-probs are consumed from the tensordict as written by RetrieveLogProb (or an offline scoring pass) and detached; student log-probs are computed by the wrapped model. Supports assistant-only masking with tokenizer fallback, sequence length normalization and mean/sum/none reductions. Closes pytorch#3921

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 Feature New feature Documentation Improvements or additions to documentation Objectives llm/ LLM-related PR, triggers LLM CI tests and removed Feature New feature labels Jun 30, 2026

[BugFix] Fix distillation masks and teacher detaching

8a743a8

github-actions Bot added the Feature New feature label Jun 30, 2026

[Quality] Apply distillation test formatting

b89c142

theap06 mentioned this pull request Jul 2, 2026

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

Open

theap06 mentioned this pull request Jul 21, 2026

[Feature] Add DistillationLoss for token-level LLM knowledge distillation #4038

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature] On-policy KL distillation loss for LLMs - #3921

[Feature] On-policy KL distillation loss for LLMs#3921
theap06 wants to merge 3 commits into
pytorch:mainfrom
theap06:llm-distillation-loss

theap06 commented Jun 30, 2026

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

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/3921

❌ 3 New Failures

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 b89c1420 vs main 74abd6c3

CPU

GPU

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Jun 30, 2026 •

edited

Loading

github-actions Bot commented Jun 30, 2026 •

edited

Loading

Benchmark Results: PR `b89c1420` vs main `74abd6c3`