[Feature] Add _Contiguous module and reshape improvements to encoders/decoders by vmoens · Pull Request #3306 · pytorch/rl

vmoens · 2026-01-08T08:27:25Z

Stack from ghstack (oldest at bottom):

Add _Contiguous helper module for torch.compile inductor compatibility
Refactor ObsEncoder.forward and ObsDecoder.forward to use flatten/unflatten with contiguous()
Add _maybe_record_function_decorator for profiling

[ghstack-poisoned]

pytorch-bot · 2026-01-08T08:27:27Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 6e4366a with merge base 6680c17 ():
💚 Looks good so far! There are no failures yet. 💚

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

[ghstack-poisoned]

…/decoders - Add _Contiguous helper module for torch.compile inductor compatibility - Refactor ObsEncoder.forward and ObsDecoder.forward to use flatten/unflatten with contiguous() - Add _maybe_record_function_decorator for profiling ghstack-source-id: 86ef6a9 Pull-Request: #3306

[ghstack-poisoned]

…/decoders - Add _Contiguous helper module for torch.compile inductor compatibility - Refactor ObsEncoder.forward and ObsDecoder.forward to use flatten/unflatten with contiguous() - Add _maybe_record_function_decorator for profiling ghstack-source-id: 485ff5d Pull-Request: #3306

[ghstack-poisoned]

…/decoders - Add _Contiguous helper module for torch.compile inductor compatibility - Refactor ObsEncoder.forward and ObsDecoder.forward to use flatten/unflatten with contiguous() - Add _maybe_record_function_decorator for profiling ghstack-source-id: d21f2a3 Pull-Request: #3306

[ghstack-poisoned]

…/decoders - Add _Contiguous helper module for torch.compile inductor compatibility - Refactor ObsEncoder.forward and ObsDecoder.forward to use flatten/unflatten with contiguous() - Add _maybe_record_function_decorator for profiling ghstack-source-id: a92513f Pull-Request: #3306

[ghstack-poisoned]

github-actions · 2026-07-06T10:04:48Z

Benchmark Results: PR `6e4366ac` vs main `e1bc0ef8`

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

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

CPU

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

Benchmark	main ops	PR ops	Change
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]`	36.10	196.11	+443.26%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,386	423.85	-82.23%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]`	190.35	56.22	-70.46%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	2,888	3,604	+24.78%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same]`	20.60	23.76	+15.33%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same]`	29.60	25.22	-14.81%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-32-512]`	33.89	28.87	-14.80%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	369.92	424.25	+14.69%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,334	2,001	-14.25%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,195	2,765	-13.46%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	3,405	2,963	-12.97%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,729	3,260	-12.57%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	429.65	380.28	-11.49%
`benchmarks/test_collectors_benchmark.py::test_single_with_rb`	7.3128	6.5352	-10.63%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	2,265	2,038	-10.02%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	854.78	775.41	-9.29%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	802.35	729.64	-9.06%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	3,177	2,894	-8.91%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,417	3,121	-8.69%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	7.0333	6.4754	-7.93%
`benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True]`	95.66	103.02	+7.70%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]`	536.12	499.98	-6.74%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None]`	80.26	85.20	+6.16%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	3,342	3,139	-6.07%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage]`	8.4093	8.8925	+5.75%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-gru]`	1.4192	1.3403	-5.56%
`benchmarks/test_envs_benchmark.py::test_simple`	1.8016	1.7032	-5.46%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True]`	38,584	36,515	-5.36%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False]`	52,634	49,999	-5.01%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None]`	1,735	1,821	+5.00%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512]`	109.68	115.13	+4.97%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True]`	38,637	36,771	-4.83%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None]`	686.99	719.00	+4.66%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None]`	1,801	1,885	+4.66%
`benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False]`	7,878	8,245	+4.66%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None]`	555.58	581.39	+4.65%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cpu]`	99.77	95.40	-4.38%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	138.66	144.55	+4.25%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-16]`	47.73	49.66	+4.04%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1]`	499.70	519.47	+3.96%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-False-0-lstm]`	0.8969	0.8620	-3.89%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle]`	12,313	11,839	-3.85%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400]`	1,052	1,092	+3.79%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None]`	274.07	283.59	+3.47%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1]`	279.07	269.56	-3.41%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64]`	12.11	12.51	+3.31%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]`	1,062	1,097	+3.28%
`benchmarks/test_objectives_benchmarks.py::test_values[vec_td_lambda_return_estimate-True-False]`	57.66	55.78	-3.25%
`benchmarks/test_envs_benchmark.py::test_transformed`	0.9045	0.8751	-3.25%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch]`	140.29	135.74	-3.24%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1]`	467.11	482.06	+3.20%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None]`	115.75	119.41	+3.17%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors]`	23,907	23,154	-3.15%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward]`	59.34	61.17	+3.08%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False]`	27,729	26,900	-2.99%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy]`	372,455	361,351	-2.98%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True]`	1.3079	1.3468	+2.98%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True]`	22,218	21,559	-2.96%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True]`	28,770	27,935	-2.90%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None]`	707.13	727.64	+2.90%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16]`	17.71	18.22	+2.84%
`benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False]`	37.99	39.06	+2.82%
`benchmarks/test_collectors_benchmark.py::test_sync_preempt`	16.27	16.72	+2.79%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64]`	4.4238	4.5462	+2.77%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True]`	24,074	23,412	-2.75%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-1]`	192.16	187.09	-2.64%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None]`	122.83	125.97	+2.56%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch]`	749.01	767.21	+2.43%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm]`	1.9806	2.0286	+2.42%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None]`	708.81	725.50	+2.36%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None]`	89.85	91.93	+2.32%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None]`	159.68	163.36	+2.31%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant]`	4,315	4,217	-2.29%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward]`	89.05	91.08	+2.28%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[reduce-overhead-None]`	219.32	224.29	+2.27%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None]`	179.31	183.36	+2.26%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True]`	33,124	32,379	-2.25%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True]`	19,951	19,506	-2.23%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward]`	63.16	64.56	+2.22%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	917.57	897.23	-2.22%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True]`	30,255	29,588	-2.21%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	985.95	1,008	+2.20%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small]`	3,413	3,340	-2.16%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True]`	18,611	18,215	-2.13%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False]`	38,346	37,537	-2.11%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img]`	403.02	411.42	+2.08%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True]`	20,219	20,637	+2.07%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None]`	280.39	286.17	+2.06%
`benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False]`	25.67	26.20	+2.04%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward]`	58.73	59.91	+2.01%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16]`	35.88	36.60	+2.01%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64]`	10.46	10.67	+2.00%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru]`	2.9348	2.9932	+1.99%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False]`	57,363	58,489	+1.96%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400]`	193.72	197.53	+1.96%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512]`	656.94	669.68	+1.94%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-gru]`	1.4798	1.4518	-1.89%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[False-None]`	95.41	97.21	+1.89%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True]`	35,081	34,428	-1.86%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img]`	563.54	573.95	+1.85%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False]`	43,037	42,267	-1.79%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None]`	283.84	288.91	+1.79%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False]`	52.01	52.94	+1.78%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16]`	42.11	42.86	+1.77%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4]`	70.87	72.12	+1.76%
`benchmarks/test_collectors_benchmark.py::test_async`	17.22	17.52	+1.75%
`benchmarks/test_objectives_benchmarks.py::test_redq_speed[True-None]`	219.77	223.54	+1.72%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512]`	2,354	2,314	-1.69%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img]`	423.04	430.10	+1.67%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-buffers-False]`	0.5928	0.5832	-1.62%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True]`	24.98	25.38	+1.61%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-64]`	1.2424	1.2621	+1.58%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small]`	4,343	4,411	+1.56%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]`	544.77	553.10	+1.53%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True]`	42,280	41,635	-1.53%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari]`	632.94	642.54	+1.52%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-64]`	3.0234	3.0678	+1.47%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant]`	2,572	2,534	-1.47%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	119.59	121.33	+1.46%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None]`	463.13	456.37	-1.46%
...	...	...	Showing 120 of 216 comparisons, sorted by absolute change.

GPU

Compared 226 benchmarks. Regressions over 5%: 14. Improvements over 5%: 19.

Benchmark	main ops	PR ops	Change
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400]`	44.64	194.08	+334.75%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	1,966	505.61	-74.28%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400]`	186.44	51.06	-72.61%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None]`	76.81	102.75	+33.76%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	3,656	2,677	-26.79%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None]`	85.22	105.51	+23.81%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	3,304	2,548	-22.89%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000]`	2,887	3,538	+22.53%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,672	3,050	-16.93%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400]`	861.74	997.61	+15.77%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,023	2,310	+14.20%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000]`	821.84	718.08	-12.63%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000]`	3,167	2,808	-11.31%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000]`	2,145	1,933	-9.88%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	2,935	2,662	-9.31%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari]`	3,978	4,338	+9.06%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400]`	440.54	479.13	+8.76%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000]`	758.33	693.64	-8.53%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward]`	232.05	248.61	+7.14%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward]`	344.94	369.33	+7.07%
`benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True]`	529.69	493.66	-6.80%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward]`	341.48	362.98	+6.30%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False]`	48,119	51,142	+6.28%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward]`	123.45	131.09	+6.19%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors]`	23,824	22,388	-6.03%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward]`	335.77	355.99	+6.02%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400]`	536.32	506.32	-5.59%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512]`	1,279	1,348	+5.40%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save]`	7,207	6,821	-5.36%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512]`	1,310	1,378	+5.22%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None]`	819.24	861.30	+5.13%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400]`	822.14	863.97	+5.09%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward]`	143.56	150.75	+5.01%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img]`	564.56	536.57	-4.96%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[100-img_shape0-atari]`	30.08	28.61	-4.88%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None]`	216.51	226.91	+4.81%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000]`	2,825	2,692	-4.72%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None]`	603.16	630.95	+4.61%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None]`	716.88	749.86	+4.60%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small]`	6,065	5,788	-4.58%
`benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant]`	4,554	4,743	+4.14%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None]`	724.09	753.95	+4.12%
`benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cpu_sampler]`	85.56	88.94	+3.95%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle]`	12,289	11,811	-3.89%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000]`	1,953	1,878	-3.84%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward]`	270.03	279.99	+3.69%
`benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None]`	822.36	851.50	+3.54%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000]`	3,070	2,962	-3.51%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False]`	31,127	32,215	+3.50%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True]`	23.28	22.48	-3.41%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward]`	897.46	867.08	-3.38%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True]`	0.4952	0.4791	-3.25%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True]`	21.54	20.84	-3.25%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small]`	3,473	3,362	-3.18%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16]`	35.33	36.45	+3.15%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None]`	801.83	826.78	+3.11%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64]`	11.94	12.31	+3.09%
`benchmarks/test_collectors_benchmark.py::test_single_with_rb_pixels`	5.2789	5.1165	-3.08%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward]`	226.81	233.76	+3.07%
`benchmarks/test_collectors_benchmark.py::test_single_with_rb`	5.9026	5.7251	-3.01%
`benchmarks/test_envs_benchmark.py::test_transformed`	0.7099	0.6887	-2.98%
`benchmarks/test_envs_benchmark.py::test_simple`	1.2297	1.1932	-2.97%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch]`	15.21	14.76	-2.91%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True]`	36,405	37,451	+2.87%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512]`	47.25	48.60	+2.85%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000]`	164.45	168.82	+2.66%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True]`	22.87	22.27	-2.63%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False]`	55,191	56,624	+2.60%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True]`	33,490	34,349	+2.57%
`benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru]`	49.70	48.54	-2.33%
`benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward]`	256.63	262.57	+2.31%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000]`	166.09	169.87	+2.28%
`benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None]`	122.63	119.89	-2.23%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy]`	338,521	331,064	-2.20%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img]`	407.98	416.55	+2.10%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch]`	136.70	139.50	+2.05%
`benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward]`	445.72	454.85	+2.05%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400]`	533.44	522.71	-2.01%
`benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False]`	11.98	12.22	+2.01%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	162.90	166.15	+1.99%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True]`	17,998	17,642	-1.98%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1]`	628.14	615.87	-1.95%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward]`	66.43	67.69	+1.89%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None]`	389.52	396.71	+1.85%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward]`	453.38	445.09	-1.83%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False]`	27,335	26,837	-1.82%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False]`	41,293	42,026	+1.78%
`benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512]`	1,234	1,256	+1.77%
`benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-ListStorage-None-4000]`	161.01	163.82	+1.75%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None]`	108.53	110.41	+1.74%
`benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False]`	19.21	19.54	+1.72%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False]`	39,364	38,687	-1.72%
`benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True]`	46.87	47.67	+1.72%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-4]`	156.91	154.37	-1.62%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1]`	513.10	504.80	-1.62%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[100-img_shape0-atari]`	13.36	13.14	-1.62%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari]`	649.44	639.33	-1.56%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward]`	368.85	374.46	+1.52%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False]`	75,725	76,863	+1.50%
`benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None]`	619.45	628.68	+1.49%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64]`	7.0207	7.1232	+1.46%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True]`	41,068	41,667	+1.46%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True]`	22,141	21,819	-1.45%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari]`	20.44	20.14	-1.45%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[200-img_shape1-large_batch]`	6.6728	6.5762	-1.45%
`benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage]`	8.4392	8.3184	-1.43%
`benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-True]`	0.2119	0.2089	-1.43%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True]`	20,733	20,438	-1.42%
`benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None]`	504.73	511.89	+1.42%
`benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch]`	308.92	304.55	-1.42%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True]`	23,093	22,780	-1.36%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-16]`	47.64	48.28	+1.35%
`benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True]`	29,398	29,795	+1.35%
`benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None]`	265.71	269.27	+1.34%
`benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward]`	317.59	321.76	+1.31%
`benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1]`	468.23	462.15	-1.30%
`benchmarks/test_envs_benchmark.py::test_serial`	0.4225	0.4280	+1.29%
`benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None]`	45.18	44.60	-1.28%
`benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None]`	87.34	88.45	+1.27%
`benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb_cuda[200-img_shape1-large_batch]`	8.7330	8.6252	-1.23%
...	...	...	Showing 120 of 226 comparisons, sorted by absolute change.

Update

8e6d742

[ghstack-poisoned]

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 Jan 8, 2026

Update

2761e9f

[ghstack-poisoned]

Update

9826818

[ghstack-poisoned]

Update

348c94d

[ghstack-poisoned]

Update

c2236ba

[ghstack-poisoned]

github-actions Bot added Feature New feature Modules labels Jun 8, 2026

github-actions Bot added Integrations/torch_geometric Integrations labels Jun 8, 2026

Update

9e50fbd

[ghstack-poisoned]

Update

6e4366a

[ghstack-poisoned]

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature] Add _Contiguous module and reshape improvements to encoders/decoders#3306

[Feature] Add _Contiguous module and reshape improvements to encoders/decoders#3306
vmoens wants to merge 7 commits into
gh/vmoens/187/basefrom
gh/vmoens/187/head

vmoens commented Jan 8, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Jan 8, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

vmoens commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Jan 8, 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/3306

✅ No Failures

Uh oh!

github-actions Bot commented Jul 6, 2026

Benchmark Results: PR 6e4366ac vs main e1bc0ef8

CPU

GPU

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

vmoens commented Jan 8, 2026 •

edited

Loading

pytorch-bot Bot commented Jan 8, 2026 •

edited

Loading

Benchmark Results: PR `6e4366ac` vs main `e1bc0ef8`