Skip to content

[Feature] Add mjlab environment wrapper#3910

Merged
vmoens merged 1 commit into
pytorch:mainfrom
vmoens:codex/mjlab-wrapper
Jun 27, 2026
Merged

[Feature] Add mjlab environment wrapper#3910
vmoens merged 1 commit into
pytorch:mainfrom
vmoens:codex/mjlab-wrapper

Conversation

@vmoens

@vmoens vmoens commented Jun 26, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add MJLabWrapper and MJLabEnv for mjlab manager-based RL environments.
  • Handle device and flat vectorized batch-size validation explicitly.
  • Default to TorchRL-driven partial resets by disabling mjlab native auto-reset, with opt-in native auto-reset support for invalid terminal observations.
  • Add render() and single-env from_pixels support via mjlab render_mode="rgb_array".
  • Add mjlab docs/config exports, optional dependency extra, tests, and a path-filtered Linux mjlab CI workflow.

Testing

  • python3 -m py_compile torchrl/envs/libs/mjlab.py test/libs/test_mjlab.py torchrl/trainers/algorithms/configs/envs_libs.py torchrl/trainers/algorithms/configs/__init__.py torchrl/envs/libs/__init__.py torchrl/envs/__init__.py
  • PYTHONPATH=$PWD /opt/homebrew/bin/python3.13 -m pytest test/libs/test_mjlab.py -q

Note: the real mjlab CUDA Cartpole smoke test is skipped locally without mjlab/CUDA and runs in the dedicated mjlab Linux workflow.

@pytorch-bot

pytorch-bot Bot commented Jun 26, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⏳ No Failures, 45 Pending

As of commit 539389c with merge base 98cfefd (image):
💚 Looks good so far! There are no failures yet. 💚

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

@meta-cla 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 26, 2026
@github-actions github-actions Bot added Feature New feature Documentation Improvements or additions to documentation CI Has to do with CI setup (e.g. wheels & builds, tests...) Environments Adds or modifies an environment wrapper Trainers and removed Feature New feature labels Jun 26, 2026
@vmoens vmoens force-pushed the codex/mjlab-wrapper branch from 59d561c to bcd3b4c Compare June 26, 2026 23:16
@github-actions github-actions Bot added the Feature New feature label Jun 26, 2026
@vmoens

vmoens commented Jun 26, 2026

Copy link
Copy Markdown
Collaborator Author

cc @lin-erica

@vmoens vmoens force-pushed the codex/mjlab-wrapper branch from bcd3b4c to bf703ef Compare June 27, 2026 00:10
@github-actions

github-actions Bot commented Jun 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Benchmark Results: PR 539389cb vs main 5b025138

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

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

CPU

Compared 216 benchmarks. Regressions over 5%: 14. Improvements over 5%: 13.

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 452.53 3,476 +668.07%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 52.05 191.91 +268.70%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 200.32 39.71 -80.18%
benchmarks/test_objectives_benchmarks.py::test_values[vec_td1_return_estimate-False-False] 87.02 55.42 -36.31%
benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True] 87.19 55.86 -35.93%
benchmarks/test_objectives_benchmarks.py::test_values[vec_td_lambda_return_estimate-True-False] 86.77 55.68 -35.83%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,822 3,459 +22.57%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same] 24.86 30.27 +21.74%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch] 11.53 13.84 +19.99%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari] 22.67 27.13 +19.65%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 3,034 3,511 +15.71%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3,241 3,738 +15.33%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 33.86 29.44 -13.04%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,395 2,107 -12.05%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] 130.40 145.46 +11.55%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 569.43 504.39 -11.42%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same] 23.96 21.26 -11.26%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,091 2,320 +10.96%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,141 2,816 -10.34%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,260 2,037 -9.85%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,156 2,338 +8.45%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3,372 3,109 -7.78%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64] 7.1328 6.6752 -6.42%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,478 3,259 -6.29%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-1] 184.28 193.93 +5.23%
benchmarks/test_objectives_benchmarks.py::test_redq_speed[False-backward] 59.95 56.82 -5.22%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward] 969.77 1,020 +5.14%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward] 58.00 60.89 +4.98%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 35,780 34,115 -4.65%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True] 1.4014 1.3369 -4.61%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] 7,022 7,326 +4.33%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward] 62.97 65.64 +4.25%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None] 83.56 87.04 +4.16%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward] 246.92 257.14 +4.14%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64] 10.39 10.82 +4.12%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None] 284.02 295.73 +4.12%
benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False] 7,595 7,888 +3.85%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None] 89.87 93.32 +3.83%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari] 702.88 729.55 +3.79%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False] 49,457 51,310 +3.75%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] 5,363 5,562 +3.72%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None] 369.17 355.68 -3.65%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None] 336.09 348.34 +3.64%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 1,056 1,093 +3.59%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None] 116.13 120.24 +3.54%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img] 568.32 588.41 +3.54%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] 48.88 50.59 +3.49%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None] 491.57 474.96 -3.38%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False] 52.86 54.64 +3.37%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 1,078 1,042 -3.35%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 478.81 462.95 -3.31%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16] 17.69 18.27 +3.31%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True] 18,677 18,061 -3.30%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] 8.6039 8.8864 +3.28%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 519.06 502.16 -3.26%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari] 647.39 668.39 +3.24%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle] 12,212 12,605 +3.23%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 675.60 654.03 -3.19%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] 64,705 66,710 +3.10%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant] 4,273 4,403 +3.04%
benchmarks/test_envs_benchmark.py::test_transformed 0.9081 0.9345 +2.90%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4] 70.72 72.68 +2.78%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None] 284.60 292.50 +2.77%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] 22,519 21,899 -2.76%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None] 712.02 731.48 +2.73%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru] 3.0513 2.9682 -2.73%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None] 163.68 168.03 +2.66%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False] 34,357 35,235 +2.55%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False] 42,676 43,745 +2.50%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward] 442.06 431.03 -2.49%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru] 4.3324 4.2248 -2.48%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] 1,784 1,827 +2.43%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None] 116.71 119.55 +2.43%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-gru] 1.4275 1.3939 -2.36%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1] 280.54 274.00 -2.33%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] 127.18 130.11 +2.31%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None] 756.84 739.44 -2.30%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None] 279.06 285.31 +2.24%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True] 28,804 29,440 +2.21%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-4] 47.69 48.72 +2.15%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False] 57,285 58,516 +2.15%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64] 4.4383 4.5308 +2.08%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1] 643.00 656.22 +2.05%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward] 79.29 80.91 +2.05%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small] 4,425 4,514 +2.03%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-4] 164.30 167.49 +1.94%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-16] 11.98 12.22 +1.93%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 2,984 3,041 +1.92%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] 60.55 61.70 +1.91%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-None] 50.62 51.57 +1.87%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True] 19,311 18,951 -1.86%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] 54.48 55.46 +1.79%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1] 476.82 485.33 +1.78%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] 143.07 145.61 +1.77%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] 177.72 174.63 -1.74%
benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True] 100.09 101.83 +1.74%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] 265.59 269.98 +1.65%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] 38,163 38,774 +1.60%
benchmarks/test_envs_benchmark.py::test_serial 0.5891 0.5980 +1.51%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 814.88 827.04 +1.49%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] 425.11 431.40 +1.48%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1] 516.38 523.96 +1.47%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward] 93.69 92.33 -1.46%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True] 42,915 42,291 -1.45%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True] 31,022 30,585 -1.41%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True] 24,681 25,028 +1.41%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512] 114.05 115.63 +1.38%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-64] 2.9993 3.0402 +1.36%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] 1,940 1,914 -1.36%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] 311.48 315.68 +1.35%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward] 84.48 85.61 +1.33%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True] 20,716 20,991 +1.33%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None] 129.75 128.04 -1.32%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy] 360,582 365,309 +1.31%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 172.70 170.46 -1.30%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True] 21,003 21,273 +1.29%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 171.00 173.14 +1.25%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] 122.40 123.91 +1.24%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None] 289.65 293.20 +1.22%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None] 337.26 341.35 +1.21%
... ... ... Showing 120 of 216 comparisons, sorted by absolute change.

GPU

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

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 463.09 2,016 +335.34%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 189.29 47.74 -74.78%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 2,831 3,697 +30.61%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 2,943 3,749 +27.39%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None] 105.73 85.14 -19.47%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 2,827 3,364 +19.00%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch] 11.23 13.34 +18.77%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari] 22.10 26.23 +18.71%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,652 3,022 +13.94%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3,017 3,341 +10.76%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,074 2,256 +8.80%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,101 3,320 +7.06%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 843.18 783.82 -7.04%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] 341.74 365.79 +7.04%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 461.69 493.24 +6.83%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward] 448.00 478.14 +6.73%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward] 373.11 398.14 +6.71%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,174 2,037 -6.31%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512] 47.50 44.51 -6.30%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 787.68 739.99 -6.05%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img] 573.98 539.40 -6.02%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] 256.92 272.36 +6.01%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 2,934 3,104 +5.80%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True] 19,603 20,722 +5.71%
benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False] 19.91 18.82 -5.44%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,891 3,038 +5.10%
benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True] 289.11 303.84 +5.10%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1] 271.06 284.05 +4.79%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle] 12,415 11,822 -4.77%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] 24,139 22,995 -4.74%
benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True] 48.16 46.03 -4.42%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward] 144.96 151.07 +4.21%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] 1,863 1,941 +4.20%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 1,053 1,010 -4.03%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward] 215.31 223.81 +3.95%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] 4,052 4,211 +3.95%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] 36,399 37,828 +3.92%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False] 55,728 57,909 +3.92%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant] 4,588 4,766 +3.87%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] 242.36 251.64 +3.83%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 521.56 501.75 -3.80%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward] 232.68 241.20 +3.66%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True] 22.02 21.26 -3.44%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True] 19,231 19,884 +3.40%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] 140.29 135.71 -3.27%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None] 756.44 780.83 +3.22%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small] 6,132 5,934 -3.22%
benchmarks/test_envs_benchmark.py::test_simple 1.2391 1.1991 -3.22%
benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False] 11.82 11.44 -3.17%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,072 2,137 +3.16%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cuda] 2,171 2,236 +3.00%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True] 22,662 23,338 +2.98%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True] 36,742 37,782 +2.83%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] 1,877 1,929 +2.78%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64] 4.4143 4.5335 +2.70%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1] 645.84 629.02 -2.60%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None] 107.82 105.06 -2.56%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari] 278.52 271.47 -2.53%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False] 49,219 50,419 +2.44%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True] 32,420 33,201 +2.41%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] 830.47 850.11 +2.36%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward] 127.76 130.74 +2.33%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] 29,687 30,373 +2.31%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64] 12.30 12.59 +2.29%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-4] 158.25 161.86 +2.28%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True] 40,691 41,614 +2.27%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch] 331.77 324.50 -2.19%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small] 4,286 4,380 +2.18%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64] 10.60 10.83 +2.17%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-backward] 77.62 79.26 +2.11%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward] 887.98 906.57 +2.09%
benchmarks/test_collectors_benchmark.py::test_single_with_rb_pixels 5.3047 5.1942 -2.08%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-4] 46.98 47.96 +2.08%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-False] 0.5969 0.6092 +2.06%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-False] 44,106 45,006 +2.04%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-backward] 82.70 84.38 +2.04%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cuda_samp... 1,496 1,465 -2.03%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-1] 188.60 192.34 +1.98%
benchmarks/test_envs_benchmark.py::test_serial 0.4184 0.4265 +1.94%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] 53.47 54.50 +1.93%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 33,186 33,824 +1.92%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 191.53 195.15 +1.89%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None] 104.58 102.62 -1.87%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1] 479.93 470.99 -1.86%
benchmarks/test_envs_benchmark.py::test_parallel 0.5423 0.5520 +1.79%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True] 20,231 20,588 +1.76%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16] 17.72 18.04 +1.76%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True] 18,364 18,684 +1.74%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True] 27,889 28,373 +1.73%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] 8.0071 8.1437 +1.71%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] 409.44 402.69 -1.65%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] 19,356 19,658 +1.56%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False] 75,427 76,587 +1.54%
benchmarks/test_envs_benchmark.py::test_transformed 0.7085 0.6979 -1.50%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-16] 49.31 50.04 +1.49%
benchmarks/test_collectors_benchmark.py::test_single 6.6742 6.5748 -1.49%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 510.58 502.99 -1.49%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1] 521.06 513.47 -1.46%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True] 17,772 18,018 +1.39%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None] 43.35 43.93 +1.34%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] 21,531 21,819 +1.34%
benchmarks/test_objectives_benchmarks.py::test_values[vec_td_lambda_return_estimate-True-False] 852.20 840.84 -1.33%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None] 614.11 621.87 +1.26%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None] 605.35 612.90 +1.25%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None] 87.99 89.01 +1.15%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm] 21.42 21.18 -1.13%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small] 3,536 3,496 -1.12%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] 7,194 7,114 -1.12%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False] 62,261 62,956 +1.12%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari] 692.93 700.20 +1.05%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 1,276 1,263 -1.03%
benchmarks/test_collectors_benchmark.py::test_async 10.92 10.81 -1.02%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False] 41,680 42,106 +1.02%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True] 29,948 30,252 +1.01%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] 63,739 63,106 -0.99%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 991.99 1,002 +0.98%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-4] 19.44 19.25 -0.96%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-16] 11.86 11.97 +0.96%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] 309.95 306.97 -0.96%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False] 31,847 31,542 -0.96%
... ... ... Showing 120 of 226 comparisons, sorted by absolute change.

@vmoens vmoens force-pushed the codex/mjlab-wrapper branch 2 times, most recently from 65e6b65 to 1605ace Compare June 27, 2026 22:34
@vmoens vmoens force-pushed the codex/mjlab-wrapper branch from 1605ace to 539389c Compare June 27, 2026 22:53
@vmoens vmoens merged commit 09482a9 into pytorch:main Jun 27, 2026
110 checks passed
@vmoens vmoens deleted the codex/mjlab-wrapper branch June 27, 2026 23:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

CI Has to do with CI setup (e.g. wheels & builds, tests...) CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Documentation Improvements or additions to documentation Environments Adds or modifies an environment wrapper Feature New feature Trainers

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vmoens