[Feature] Unify TorchRL checkpointing by vmoens · Pull Request #3967 · pytorch/rl

vmoens · 2026-07-09T13:21:56Z

Summary

add the manifest-driven torchrl.checkpoint API with directory and ZIP archive containers, atomic replacement, partial loading, adapter/version checks, scoped adapters, operation options, migrations, and process-global RNG state
bind unified checkpoints across the base Trainer and PPO, A2C, REINFORCE, SAC, offline-to-online SAC, DQN, DDPG, IQL, CQL, and TD3 while retaining versioned legacy backend compatibility
migrate DQN CartPole and rlrender checkpoint helpers, including full-state resume and policy-only rendering that avoids replay-buffer payloads
persist target-update, evaluator, logger, hook, and trainer progress and add API documentation plus a Sphinx tutorial

Testing

uv run pytest -q test/test_checkpoint.py test/test_render.py
uv run pytest -q test/test_trainer.py -k 'unified_trainer_checkpoint or subclass_exposes_checkpoint or stepper_checkpoint_roundtrip'
uv run pytest -q test/test_configs.py -k 'checkpoint_config_parity or (TrainerConfigs and trainer_config)'
uv run pytest -q test/collectors/test_evaluator.py -k checkpoint_state
uv run pytest -q test/test_loggers.py -k checkpoint_state
DQN CartPole archive save and resume smoke tests
uv run python tutorials/sphinx-tutorials/checkpointing.py
pre-commit hooks on all changed files

pytorch-bot · 2026-07-09T13:22:00Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 6 New Failures, 2 Unrelated Failures

As of commit 868b54a with merge base a74b478 ():

NEW FAILURES - The following jobs have failed:

Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
RuntimeError: Command docker exec -t 551ea20050e8c163560fc7c45a97150175cbdc7860f3db2ab88882d787cc00a3 /exec failed with exit code 1
Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
RuntimeError: Command docker exec -t 8e65290da311f49258937ea9b88faa4029c35a846b1a3b3bd31a8b8caf52e28e /exec failed with exit code 1
Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
test/test_configs.py::TestDataConfigs::test_writer_ensemble_config
Unit-tests on Linux / tests-cpu (3.13) / linux-job (gh)
RuntimeError: Command docker exec -t 8c149ad4e0d176443030ce1033a77b434ff413c333458ff24fc18ed361734f91 /exec failed with exit code 1
Unit-tests on Linux / tests-cpu (3.14) / linux-job (gh)
RuntimeError: Command docker exec -t 2fa48c4d12675bdac04cc3f75f51353361a2a6cc2d40af3c6ebe9e0573552f36 /exec failed with exit code 1
Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh)
TypeError: TensorDictBase.save() got an unexpected keyword argument 'archive'

FLAKY - The following job failed but was likely due to flakiness present on trunk:

Unit-tests on Linux / tests-olddeps (3.10, 11.8) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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

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

Unit-tests on Linux / tests-optdeps (3.12, 13.0) / linux-job (gh) (trunk failure)
##[error]The operation was canceled.

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

github-actions · 2026-07-09T14:23:38Z

Benchmark Results: PR `868b54aa` vs main `a74b478f`

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

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

CPU

Compared 220 benchmarks. Regressions over 5%: 11. Improvements over 5%: 22.

Benchmark	main ops	PR ops	Change
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]`	187.12	56.61	-69.75%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None]`	62.42	92.97	+48.93%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	1,810	2,367	+30.77%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	2,863	3,726	+30.12%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16]`	28.26	36.25	+28.30%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	1,782	2,265	+27.10%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,761	3,499	+26.71%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]`	47.75	37.58	-21.30%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	1,982	2,305	+16.29%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None]`	407.65	468.08	+14.83%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	2,911	3,329	+14.38%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-32-512]`	34.11	29.33	-14.01%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same]`	1,689	1,921	+13.70%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same]`	2,768	2,396	-13.46%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari]`	5,059	4,403	-12.98%
`benchmarks/test_collectors_benchmark.py::test_single_with_rb`	7.3974	6.5280	-11.75%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,718	2,961	+8.94%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	132.70	144.37	+8.79%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru]`	2.9174	3.1719	+8.72%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	392.77	423.74	+7.89%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	940.58	1,012	+7.59%
`benchmarks/test_envs_benchmark.py::test_parallel`	0.8930	0.9606	+7.57%
`benchmarks/test_envs_benchmark.py::test_simple`	1.8154	1.6832	-7.28%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]`	1,068	991.93	-7.08%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	749.34	801.96	+7.02%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]`	560.18	523.20	-6.60%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	5.9934	5.6498	-5.73%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm]`	1.9357	2.0454	+5.67%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-memmap_cpu_storage_cpu...`	79.16	83.64	+5.65%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	800.95	843.54	+5.32%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors]`	24,091	22,822	-5.27%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False]`	47.05	49.49	+5.19%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	497.10	522.61	+5.13%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]`	535.46	509.20	-4.90%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None]`	271.21	284.51	+4.90%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save]`	7,233	6,886	-4.79%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward]`	111.82	116.81	+4.46%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True]`	0.5302	0.5069	-4.41%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True]`	23.10	24.08	+4.23%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True]`	31,157	29,842	-4.22%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru]`	4.1331	4.3057	+4.18%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]`	1,053	1,097	+4.17%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None]`	276.54	287.98	+4.14%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True]`	24.78	25.79	+4.04%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward]`	124.01	128.83	+3.89%
`benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False]`	7,972	7,662	-3.88%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False]`	51,130	49,260	-3.66%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant]`	4,481	4,318	-3.66%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward]`	525.65	506.98	-3.55%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch]`	10.05	10.40	+3.54%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None]`	280.12	289.72	+3.42%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-lstm]`	3.1237	3.0196	-3.33%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	885.56	914.77	+3.30%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None]`	82.68	85.32	+3.19%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-True]`	0.2138	0.2071	-3.16%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False]`	39,742	38,494	-3.14%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True]`	20,377	21,015	+3.13%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch]`	774.39	750.38	-3.10%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None]`	1,805	1,751	-3.03%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None]`	718.93	697.66	-2.96%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[reduce-overhead-None]`	219.68	226.16	+2.95%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None]`	115.22	118.59	+2.93%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False]`	30,452	29,571	-2.89%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,775	3,667	-2.84%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch]`	310.66	319.42	+2.82%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,318	3,225	-2.81%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512]`	2,323	2,258	-2.80%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[True-None]`	220.12	226.17	+2.75%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True]`	1.3222	1.2861	-2.73%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	118.27	121.40	+2.65%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]`	168.13	163.87	-2.53%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False]`	31,796	32,598	+2.52%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch]`	332.25	340.41	+2.46%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]`	163.21	159.23	-2.44%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None]`	82.63	84.63	+2.41%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-1]`	78.75	80.63	+2.39%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None]`	282.66	289.35	+2.37%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None]`	116.42	119.01	+2.22%
`benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False]`	38.18	37.34	-2.21%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward]`	60.60	59.28	-2.17%
`benchmarks/test_collectors_benchmark.py::test_sync_preempt`	16.24	16.59	+2.17%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None]`	160.52	163.95	+2.13%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	164.26	160.79	-2.11%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None]`	567.18	578.95	+2.07%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	3,152	3,087	-2.07%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	165.16	161.75	-2.06%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-4]`	19.76	20.16	+2.03%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512]`	111.18	113.39	+1.99%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False]`	52.30	53.34	+1.98%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None]`	701.59	715.31	+1.96%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari]`	692.94	706.40	+1.94%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1]`	496.63	506.25	+1.94%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[False-None]`	98.15	96.25	-1.94%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False]`	65,190	63,931	-1.93%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1]`	462.19	471.10	+1.93%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle]`	12,341	12,103	-1.93%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512]`	669.56	656.75	-1.91%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None]`	580.59	569.48	-1.91%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True]`	23,591	24,030	+1.86%
`benchmarks/test_envs_benchmark.py::test_transformed`	0.9181	0.9010	-1.86%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward]`	88.79	90.44	+1.85%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True]`	33,159	32,553	-1.83%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True]`	30,824	30,266	-1.81%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1]`	264.85	269.65	+1.81%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False]`	58,532	57,478	-1.80%
`benchmarks/test_replaybuffer_benchmark.py::test_replay_buffer_direct_client_identity`	7,840,394	7,981,503	+1.80%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True]`	22,226	21,832	-1.77%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-4]`	143.48	145.98	+1.75%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-backward]`	94.15	92.54	-1.72%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False]`	35,613	35,004	-1.71%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None]`	257.74	262.11	+1.69%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant]`	2,648	2,603	-1.69%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]`	167.51	164.72	-1.67%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True]`	29,093	28,613	-1.65%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small]`	7,151	7,037	-1.59%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-gru]`	1.4203	1.4423	+1.55%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward]`	78.94	80.16	+1.55%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-False]`	1.6053	1.5805	-1.54%
`benchmarks/test_collectors_benchmark.py::test_async`	17.63	17.36	-1.52%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64]`	4.4661	4.5338	+1.52%
...	...	...	Showing 120 of 220 comparisons, sorted by absolute change.

GPU

Compared 230 benchmarks. Regressions over 5%: 12. Improvements over 5%: 12.

Benchmark	main ops	PR ops	Change
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]`	44.12	198.02	+348.85%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]`	514.99	58.95	-88.55%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None]`	59.88	88.17	+47.24%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None]`	76.65	102.27	+33.42%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None]`	84.34	99.67	+18.17%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,315	1,975	-14.69%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,141	2,708	-13.77%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,064	2,328	+12.83%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,201	2,797	-12.62%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	447.73	494.43	+10.43%
`benchmarks/test_collectors_benchmark.py::test_sync`	9.6111	10.51	+9.32%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	6.9096	6.2736	-9.20%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]`	473.63	517.23	+9.20%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari]`	4,108	4,481	+9.07%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,318	3,073	-7.37%
`benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True]`	415.50	445.72	+7.27%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,895	2,687	-7.19%
`benchmarks/test_collectors_benchmark.py::test_sync_preempt`	10.44	9.7098	-7.01%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	3,186	2,987	-6.23%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	882.61	828.83	-6.09%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False]`	32,570	30,613	-6.01%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	3,117	2,933	-5.90%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward]`	377.18	398.19	+5.57%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]`	985.23	1,039	+5.45%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True]`	37,041	38,882	+4.97%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False]`	64,125	67,297	+4.95%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward]`	223.69	234.61	+4.88%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True]`	31,407	29,896	-4.81%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False]`	39,666	37,778	-4.76%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512]`	49.01	46.70	-4.72%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-False]`	0.6052	0.5767	-4.72%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	843.38	806.08	-4.42%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img]`	425.84	407.01	-4.42%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch]`	764.22	730.96	-4.35%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True]`	21,566	20,631	-4.34%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	263.54	274.66	+4.22%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img]`	175.74	169.30	-3.66%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True]`	0.5257	0.5067	-3.62%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True]`	20,344	19,613	-3.59%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None]`	1,907	1,975	+3.54%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img]`	590.65	569.76	-3.54%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img]`	407.94	393.65	-3.50%
`benchmarks/test_envs_benchmark.py::test_simple`	1.2446	1.2014	-3.47%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64]`	10.36	10.71	+3.36%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True]`	41,555	42,934	+3.32%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-lstm]`	81.28	78.67	-3.22%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same]`	3,444	3,554	+3.21%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512]`	1,310	1,270	-3.00%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16]`	41.01	42.21	+2.91%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None]`	1,855	1,909	+2.90%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[50-img_shape0-small]`	858.86	883.38	+2.85%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage]`	8.7534	9.0002	+2.82%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari]`	662.26	644.51	-2.68%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	788.80	767.80	-2.66%
`benchmarks/test_envs_benchmark.py::test_serial`	0.4193	0.4302	+2.59%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari]`	702.63	684.68	-2.56%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	362.60	371.78	+2.53%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None]`	856.85	878.53	+2.53%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch]`	14.79	15.16	+2.51%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-16]`	47.59	48.77	+2.47%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None]`	110.94	108.28	-2.40%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None]`	833.60	853.31	+2.36%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1]`	458.81	448.11	-2.33%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small]`	4,470	4,370	-2.24%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward]`	145.75	149.02	+2.24%
`benchmarks/test_envs_benchmark.py::test_transformed`	0.7134	0.6974	-2.24%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-None]`	399.44	390.81	-2.16%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16]`	35.50	36.27	+2.16%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward]`	451.73	441.99	-2.16%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]`	167.24	170.80	+2.13%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64]`	12.12	12.38	+2.13%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None]`	624.42	637.35	+2.07%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]`	165.69	169.11	+2.06%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False]`	29,839	29,236	-2.02%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True]`	22,285	21,843	-1.98%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False]`	56,153	55,050	-1.96%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True]`	19,545	19,168	-1.93%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1]`	608.92	620.63	+1.92%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	2,195	2,237	+1.89%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward]`	247.56	242.91	-1.88%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[100-img_shape0-atari]`	29.41	29.95	+1.86%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True]`	34,170	33,538	-1.85%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward]`	350.69	357.16	+1.84%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None]`	696.56	708.73	+1.75%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small]`	3,522	3,461	-1.74%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False]`	34,402	33,806	-1.73%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64]`	4.4994	4.5760	+1.70%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512]`	1,263	1,242	-1.64%
`benchmarks/test_collectors_benchmark.py::test_single_pixels`	6.1525	6.2527	+1.63%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False]`	27,480	27,045	-1.58%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	162.15	164.64	+1.54%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch]`	134.92	136.96	+1.52%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	162.50	164.92	+1.49%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True]`	28,822	28,395	-1.48%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None]`	637.81	628.73	-1.42%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None]`	433.07	439.23	+1.42%
`benchmarks/test_replaybuffer_benchmark.py::TestWindowingTransformsBenchmark::test_catframes_offline`	637.01	646.07	+1.42%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None]`	121.22	119.50	-1.42%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-4]`	183.69	181.11	-1.41%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru]`	22.63	22.32	-1.39%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	895.57	883.20	-1.38%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm]`	21.41	21.11	-1.38%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512]`	1,283	1,300	+1.37%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1]`	507.08	500.32	-1.33%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-False]`	1.5976	1.6187	+1.32%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]`	167.42	169.53	+1.26%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,810	2,774	-1.26%
`benchmarks/test_collectors_benchmark.py::test_single`	6.7584	6.6736	-1.25%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-triton-False-0-gru]`	20.31	20.05	-1.25%
`benchmarks/test_envs_benchmark.py::test_parallel`	0.5430	0.5497	+1.24%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari]`	20.16	20.40	+1.21%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-False]`	0.2250	0.2277	+1.20%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,103	3,141	+1.20%
`benchmarks/test_collectors_benchmark.py::test_sync_pixels`	10.28	10.16	-1.17%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[200-img_shape1-large_batch]`	6.6851	6.6077	-1.16%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	450.95	456.18	+1.16%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]`	184.46	186.59	+1.16%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16]`	17.92	18.12	+1.15%
`benchmarks/test_replaybuffer_benchmark.py::test_replay_buffer_direct_client_identity`	7,673,908	7,586,730	-1.14%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None]`	343.40	339.50	-1.13%
...	...	...	Showing 120 of 230 comparisons, sorted by absolute change.

…ication # Conflicts: # torchrl/trainers/trainers.py

…ication

[Feature] Unify TorchRL checkpointing

6d7ab16

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 Jul 9, 2026

github-actions Bot added Feature New feature Documentation Improvements or additions to documentation Objectives Collectors sota-implementations/ tutorials/ Record Trainers Integrations/torch_geometric Integrations labels Jul 9, 2026

vmoens added 5 commits July 9, 2026 16:39

[BugFix] Harden unified checkpoint integrations

17f587b

Merge remote-tracking branch 'origin/main' into codex/checkpoint-unif…

f962bed

…ication # Conflicts: # torchrl/trainers/trainers.py

[Quality] Remove checkpoint version branding

0fda1aa

[Feature] Default checkpoint state to TensorDict

78f02a9

Merge remote-tracking branch 'origin/main' into codex/checkpoint-unif…

868b54a

…ication

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature] Unify TorchRL checkpointing#3967

[Feature] Unify TorchRL checkpointing#3967
vmoens wants to merge 6 commits into
mainfrom
codex/checkpoint-unification

vmoens commented Jul 9, 2026

Uh oh!

pytorch-bot Bot commented Jul 9, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 9, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

vmoens commented Jul 9, 2026

Summary

Testing

Uh oh!

pytorch-bot Bot commented Jul 9, 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/3967

❌ 6 New Failures, 2 Unrelated Failures

Uh oh!

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

Uh oh!

Benchmark Results: PR 868b54aa vs main a74b478f

CPU

GPU

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

pytorch-bot Bot commented Jul 9, 2026 •

edited

Loading

github-actions Bot commented Jul 9, 2026 •

edited

Loading

Benchmark Results: PR `868b54aa` vs main `a74b478f`