Skip to content

[Feature] Introduce mujoco_playground environment wrapper#3751

Merged
vmoens merged 11 commits into
pytorch:mainfrom
itwasabhi:mujoco_playground0
Jun 17, 2026
Merged

[Feature] Introduce mujoco_playground environment wrapper#3751
vmoens merged 11 commits into
pytorch:mainfrom
itwasabhi:mujoco_playground0

Conversation

@itwasabhi

@itwasabhi itwasabhi commented May 14, 2026

Copy link
Copy Markdown
Contributor

Introduces a torch-rl wrapper of the mujoco playground environment. Wrapper also supports multi-agent decomposition.

Motivation and Context

#3733

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds core functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation (update in the documentation)
  • Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

  • I have read the CONTRIBUTION guide (required)
  • My change requires a change to the documentation.
  • I have updated the tests accordingly (required for a bug fix or a new feature).
  • I have updated the documentation accordingly.

@pytorch-bot

pytorch-bot Bot commented May 14, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 3 New Failures, 2 Unrelated Failures

As of commit 1d12044 with merge base 52cf596 (image):

NEW FAILURES - The following jobs have failed:

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

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

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

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 May 14, 2026
@github-actions

github-actions Bot commented May 14, 2026

Copy link
Copy Markdown
Contributor

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Introduce mujoco_playground environment wrapper.

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix Label Applied Example
[BugFix] BugFix [BugFix] Fix memory leak in collector
[Feature] Feature [Feature] Add new optimizer
[Doc] or [Docs] Documentation [Doc] Update installation guide
[Refactor] Refactoring [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Tests [Tests] Add unit tests for buffer
[Environment] or [Environments] Environments [Environments] Add Gymnasium support
[Data] Data [Data] Fix replay buffer sampling
[Performance] or [Perf] Performance [Performance] Optimize tensor ops
[BC-Breaking] bc breaking [BC-Breaking] Remove deprecated API
[Deprecation] Deprecation [Deprecation] Mark old function
[Quality] Quality [Quality] Fix typos and add codespell

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

@github-actions github-actions Bot added 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 Examples Trainers labels May 14, 2026
@itwasabhi itwasabhi changed the title Introduce mujoco_playground environment wrapper. [Feature] Introduce mujoco_playground environment wrapper. May 14, 2026
@github-actions github-actions Bot added the Feature New feature label May 14, 2026
@vmoens vmoens force-pushed the mujoco_playground0 branch from c7a92d5 to 15de4dd Compare May 19, 2026 10:21
@itwasabhi

itwasabhi commented May 19, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the fixes @vmoens

@itwasabhi itwasabhi force-pushed the mujoco_playground0 branch from 15de4dd to 916baf0 Compare May 21, 2026 21:26
@itwasabhi

itwasabhi commented May 21, 2026

Copy link
Copy Markdown
Contributor Author

failing mujoco_playground tests should now pass.

@itwasabhi

itwasabhi commented May 26, 2026

Copy link
Copy Markdown
Contributor Author

@vmoens any other requested changes? I believe the remaining test failures are unrelated.

itwasabhi and others added 4 commits June 4, 2026 10:07
@itwasabhi
Introduce mujoco_playground environment wrapper.
a4ccad2 
Introduces a torch-rl wrapper of the mujoco playground environment.
Wrapper also supports multi-agent decomposition.
@itwasabhi
Minor cleanup.
f445a5a
@vmoens @itwasabhi
[Feature] mujoco_playground wrapper review fixes
fb3a85b 
- CI: install `playground` (PyPI name) instead of import name
  `mujoco_playground`; the previous `pip install mujoco_playground` was the
  root cause of the current `unittests-mujoco-playground` CI failure.
- CI: drop the broken try/except JAX-init fallback in run_test.sh — it ran
  in a subprocess that exited immediately, so `JAX_PLATFORM_NAME` never
  reached pytest and any real "GPU not visible" error was hidden.
- Wrapper: docstrings no longer advertise a `state` field that the env
  never emits; the JAX state is intentionally kept on
  `self._current_state` rather than round-tripped through TensorDict (this
  is now documented as a `.. note::` block).
- Wrapper: freeze `MujocoPlaygroundAgentSpec` /
  `MujocoPlaygroundAgentMapping` dataclasses and deep-copy
  `KNOWN_MARL_MAPPINGS` entries on string lookup so users cannot mutate
  the module-level mapping by accident.
- Wrapper: emit a `UserWarning` when resolving a string against
  `KNOWN_MARL_MAPPINGS`, since those indices target Brax's observation
  layout and may not be semantically equivalent for mujoco_playground envs.
- Wrapper: document the policy contract for `homogenization_mode='max'`
  and `'concat'` (which action/obs entries are real vs padding/discarded).
- Wrapper: align `_MujocoPlaygroundMeta` num_workers handling with
  `_BraxMeta`; make `agent_mapping` and `config`/`config_overrides`
  keyword-only; accept `seed=None` in `_set_seed` (defaults to 0, matching
  the `_reset` fallback) instead of raising bare `Exception`; document
  `_listerize`'s inclusive-range semantics; drop unused
  `pixels_only`/`camera_id`/`render_kwargs` from `_build_env`.
- Example: rewrite `save_visualization` in
  `profile_mujoco_playground_collector.py` to snapshot
  `env._current_state` during a manual rollout, removing dependencies on
  non-existent `env._state_example` and `td["state"]`.
- Config: drop stale `categorical_action_encoding` field on
  `MujocoPlaygroundEnvConfig`; add `agent_mapping` and `num_workers` so
  the MARL and parallel-process knobs are reachable from the config
  system.
- Tests: replace three duplicated `_setup_jax` fixtures with a single
  module-level autouse fixture; introduce a session-scoped
  `marl_env_sizes` fixture to avoid constructing a throwaway env in every
  MARL test; expect the new MABrax warning on string lookups; add a
  new `TestMujocoPlaygroundDictObs` class covering reset /
  `check_env_specs` and a negative test that
  `homogenization_mode != 'none'` raises `NotImplementedError` for
  dict-obs envs; make `test_no_mapping_regression` actually exercise the
  MARL env it constructs.
- Docs: keep the `MOGymEnv` / `MOGymWrapper` pair adjacent in
  `envs_libraries.rst`; `__repr__` now uses `env_name=` to match the
  constructor kwarg.
@itwasabhi
fix failing tests
cc6c461
@itwasabhi itwasabhi force-pushed the mujoco_playground0 branch from 916baf0 to cc6c461 Compare June 4, 2026 09:10
@itwasabhi

itwasabhi commented Jun 4, 2026

Copy link
Copy Markdown
Contributor Author

@vmoens OK with these changes?

@vmoens vmoens changed the title [Feature] Introduce mujoco_playground environment wrapper. [Feature] Introduce mujoco_playground environment wrapper Jun 8, 2026
vmoens and others added 2 commits June 8, 2026 10:27
@vmoens @claude
[Feature] Auto-append FrameSkipTransform for envs without native fram…
23aa857 
…e_skip

Several env wrappers (Brax, Jumanji, MuJoCo Playground, ...) accept a
``frame_skip`` argument but silently ignore it because they don't loop
over it in their own ``_step``. Capture this once in the ``EnvBase``
metaclass: when ``frame_skip > 1`` and the env does not implement it
natively, wrap the env in a ``FrameSkipTransform``.

A new ``EnvBase._has_frame_skip`` class attribute (default ``False``)
gates the behaviour. Envs that implement frame-skipping in their own
``_step`` opt out by setting it to ``True``:

- ``GymLikeEnv`` (loops over ``wrapper_frame_skip``) -> True, which also
  covers Gym, DMControl, RoboHive, Habitat, IsaacGym, SafetyGymnasium.
- ``GenesisWrapper`` (loops over ``_frame_skip``) -> True.
- ``JumanjiWrapper`` overrides ``_step`` without the loop, so it resets
  the inherited flag back to False and now gets the auto-transform.

The default construction path is unchanged (the block only triggers for
``frame_skip > 1``), so there is no impact on existing envs. Building the
``TransformedEnv`` does not recurse: ``TransformedEnv`` is created via
``_TEnvPostInit`` which bypasses ``_EnvPostInit.__call__``.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@vmoens @claude
[Environment] mujoco_playground: review fixes (perf, docs, tests)
ccf6800 
- Cache the registry env listing (`available_envs`) once per process
  instead of re-enumerating the suites on every classproperty access.
- Precompute per-agent obs/action dims and index tensors in
  `_make_marl_specs` so the hot-path `_split_obs_for_agents` /
  `_reconstruct_global_action` no longer recompute `max_obs`/`raw_sizes`
  on every step. Behaviour is unchanged (identical shapes/values).
- Drop the redundant `wrapper_frame_skip = 1` line: `frame_skip > 1` is
  now handled by the auto-appended `FrameSkipTransform`.
- Document the two limitations surfaced in review: partial resets are not
  supported for vmapped batches (the JAX state lives on the instance), and
  `terminated` is aliased to `done` (no separate truncation signal).
- Document the public `MujocoPlaygroundAgentMapping` /
  `MujocoPlaygroundAgentSpec` dataclasses in the reference docs.
- Add a `frame_skip` test class and the missing
  `if __name__ == "__main__"` pytest entrypoint.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@github-actions github-actions Bot added Environments/jumanji Triggers jumanji environment tests only Transforms labels Jun 8, 2026
vmoens and others added 3 commits June 8, 2026 10:29
@vmoens
Merge remote-tracking branch 'origin/main' into pr3751-fixes
2ed4ed3
@vmoens @claude
[BugFix] frame_skip metaclass: read frame_skip from __dict__, not get…
79d6aec 
…attr

The auto-FrameSkipTransform block in `_EnvPostInit.__call__` read
`getattr(instance, "frame_skip", 1)` on every env construction. On a
batched env (`ParallelEnv`/`SerialEnv`) `frame_skip` is not a real
attribute, so this fell through to `BatchedEnv.__getattr__`, which for a
lazy (non-started) env calls `self.start()` and raises `RuntimeError`
("Trying to access attributes of closed/non started environments").
`getattr`'s default only swallows `AttributeError`, so the RuntimeError
propagated and broke construction of every lazy batched env (observed
across brax, jumanji, mujoco_playground num_workers tests and any plain
`ParallelEnv(...)` in the suite), in addition to spuriously starting the
workers.

Read `frame_skip` straight from `instance.__dict__` instead: it is set
there by `_EnvWrapper.__init__`, and batched/native envs that don't set
it simply default to 1 and are skipped without ever touching
`__getattr__`. Adds a regression test constructing lazy `SerialEnv` /
`ParallelEnv` of `CountingEnv`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@vmoens
linter
b3c41ee
@itwasabhi

itwasabhi commented Jun 17, 2026

Copy link
Copy Markdown
Contributor Author

Hi @vmoens would love to help push this one further. Can you help me understand
(a) merge/force push policy -- is the goal to keep merging in main for long-open PRs?
(b) flaky test policy -- I believe all new tests are passing, anything else stopping merging?

@vmoens

vmoens commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator

Hey @itwasabhi
Sorry for the delay on this. I made a few patches and relaunched the CI. Hopefully we're close to the finish line!

@github-actions

github-actions Bot commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Benchmark Results: PR 1d120447 vs main 52cf5961

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

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

CPU

Compared 192 benchmarks. Regressions over 5%: 7. Improvements over 5%: 11.

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,493 3,207 +28.63%
benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False] 6,675 7,920 +18.65%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,237 2,650 -18.13%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward] 239.57 279.10 +16.50%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 2,745 3,119 +13.63%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same] 23.84 21.00 -11.91%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,331 2,067 -11.33%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-sampler7-10000] 770.03 852.68 +10.73%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,942 3,245 +10.32%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 1,042 935.19 -10.28%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 724.30 797.66 +10.13%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,278 2,049 -10.02%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,234 2,015 -9.80%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward] 103.29 113.37 +9.75%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 480.31 515.78 +7.38%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 483.59 517.11 +6.93%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[pickle] 12,047 11,357 -5.72%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None] 260.71 275.27 +5.58%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] 23,233 22,072 -5.00%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None] 282.81 268.78 -4.96%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch] 772.56 734.35 -4.95%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3,653 3,482 -4.67%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None] 459.43 480.42 +4.57%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant] 4,523 4,319 -4.52%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] 119.11 124.31 +4.36%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 2,895 3,017 +4.21%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,127 2,996 -4.21%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] 399.64 416.21 +4.15%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] 7,061 6,774 -4.07%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True] 31,465 32,740 +4.05%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] 135.59 140.92 +3.93%
benchmarks/test_envs_benchmark.py::test_parallel 0.9795 0.9420 -3.83%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True] 29,900 31,022 +3.75%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False] 33,901 35,168 +3.74%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 2,927 2,818 -3.71%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 38.30 36.95 -3.53%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img] 383.10 396.41 +3.47%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True] 18,151 18,774 +3.43%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari] 263.73 272.67 +3.39%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True] 1.3544 1.3089 -3.36%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False] 63,222 65,174 +3.09%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None] 113.19 116.42 +2.85%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[100-img_shape0-atari] 29.82 30.63 +2.75%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True] 20,487 21,046 +2.73%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 492.06 505.16 +2.66%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,123 2,067 -2.64%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 33,456 34,339 +2.64%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] 136.26 132.72 -2.60%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True] 19,490 19,988 +2.55%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru] 4.1587 4.2632 +2.51%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True] 24.14 24.75 +2.50%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None] 333.64 325.30 -2.50%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-False] 0.5889 0.6036 +2.48%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small] 7,263 7,100 -2.24%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 659.28 644.84 -2.19%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] 36,362 37,145 +2.15%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari] 636.55 622.93 -2.14%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward] 58.67 57.42 -2.12%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward] 33.34 32.65 -2.04%
benchmarks/test_envs_benchmark.py::test_serial 0.5732 0.5849 +2.04%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 28.52 29.08 +1.94%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-None] 463.80 472.65 +1.91%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward] 63.47 62.30 -1.84%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False] 31,437 32,011 +1.83%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True] 40,161 40,855 +1.73%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] 308.78 314.11 +1.72%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-None] 162.71 159.98 -1.68%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward] 132.75 130.53 -1.67%
benchmarks/test_envs_benchmark.py::test_transformed 0.8873 0.9021 +1.66%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 184.67 187.72 +1.65%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-backward] 92.61 91.10 -1.63%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] 172.67 169.86 -1.63%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False] 49,876 50,682 +1.62%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False] 41,895 42,557 +1.58%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-None] 212.77 209.44 -1.56%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 165.51 168.08 +1.55%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False] 34,230 34,745 +1.50%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None] 82.90 84.10 +1.46%
benchmarks/test_envs_benchmark.py::test_simple 1.7920 1.7663 -1.43%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[10000000-cpu] 52.00 52.72 +1.39%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch] 15.27 15.48 +1.38%
benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False] 26.22 26.58 +1.38%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None] 90.29 89.05 -1.37%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-lstm] 0.9521 0.9392 -1.35%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] 29,918 30,309 +1.31%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-memmap_cpu_storage_cpu... 81.59 82.65 +1.30%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small] 4,468 4,411 -1.28%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 194.13 191.69 -1.26%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True] 23.44 23.73 +1.23%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img] 560.45 567.10 +1.19%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None] 84.71 83.70 -1.19%
benchmarks/test_objectives_benchmarks.py::test_redq_speed[False-None] 94.22 95.31 +1.16%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False] 49,833 50,398 +1.13%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 1,056 1,044 -1.13%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 166.60 168.48 +1.13%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-cudnn-True-0-gru] 1.4126 1.4286 +1.13%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward] 78.94 78.05 -1.12%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] 50.03 49.48 -1.10%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch] 13.42 13.57 +1.09%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None] 556.91 550.89 -1.08%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-backward] 524.50 518.91 -1.07%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-None] 124.24 122.94 -1.05%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 559.74 553.93 -1.04%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] 118.88 117.66 -1.03%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward] 253.12 250.52 -1.03%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False] 29,113 29,408 +1.01%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-no-buffers-True] 0.2155 0.2134 -0.98%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] 57.70 58.27 +0.97%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False] 76,480 77,223 +0.97%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant] 2,686 2,711 +0.96%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] 5,122 5,075 -0.91%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-no-buffers-False] 0.6788 0.6850 +0.91%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[reduce-overhead-None] 571.55 566.37 -0.91%
benchmarks/test_objectives_benchmarks.py::test_redq_speed[reduce-overhead-None] 222.69 224.71 +0.91%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward] 84.24 83.49 -0.89%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[numpy] 374,628 371,349 -0.88%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None] 711.18 705.05 -0.86%
benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False] 38.79 39.12 +0.85%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] 19,627 19,788 +0.82%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] 8.7494 8.6788 -0.81%
... ... ... Showing 120 of 192 comparisons, sorted by absolute change.

GPU

Compared 202 benchmarks. Regressions over 5%: 8. Improvements over 5%: 38.

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 36.50 47.26 +29.49%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None] 87.10 109.58 +25.81%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3,442 2,691 -21.81%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3,336 2,780 -16.66%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 2,901 3,296 +13.63%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 465.27 420.75 -9.57%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] 21,900 23,946 +9.34%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 532.16 485.28 -8.81%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False] 31,454 34,155 +8.59%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False] 31,600 34,295 +8.53%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-False] 38,199 41,449 +8.51%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False] 27,293 29,485 +8.03%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False] 49,223 53,069 +7.81%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True] 20,023 21,568 +7.71%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,170 2,009 -7.45%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False] 62,560 67,213 +7.44%
benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True] 295.16 316.43 +7.21%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True] 18,720 20,059 +7.16%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False] 34,539 36,969 +7.03%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] 418.13 389.02 -6.96%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] 63,524 67,923 +6.92%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,437 3,673 +6.85%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] 30,238 32,240 +6.62%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-True] 30,119 32,107 +6.60%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,141 2,282 +6.57%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True] 1.3346 1.4220 +6.55%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-True] 28,356 30,201 +6.51%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] 347.68 370.13 +6.46%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 499.71 531.68 +6.40%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] 22,899 24,317 +6.19%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] 237.88 252.35 +6.08%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img] 402.77 378.96 -5.91%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None] 400.96 424.50 +5.87%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-False] 44,929 47,554 +5.84%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False] 50,351 53,274 +5.81%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True] 22.20 23.47 +5.75%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True] 38,123 40,294 +5.69%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False] 54,749 57,850 +5.67%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-False] 42,686 45,102 +5.66%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] 142.84 134.94 -5.53%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] 47.13 49.72 +5.51%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 956.50 1,008 +5.34%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False] 76,580 80,542 +5.17%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True] 18,182 19,101 +5.05%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False] 57,224 60,104 +5.03%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small] 5,966 6,264 +5.00%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True] 23,509 24,664 +4.91%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False] 29,801 31,260 +4.90%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-True] 33,100 34,719 +4.89%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] 175.62 167.16 -4.81%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 2,103 2,202 +4.74%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] 19,948 20,877 +4.66%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True] 21,203 22,129 +4.37%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] 1,881 1,962 +4.28%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 185.17 193.00 +4.23%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-None] 798.68 832.46 +4.23%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] 8.7127 9.0696 +4.10%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True] 42,525 44,254 +4.07%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] 37,704 39,195 +3.95%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,602 2,703 +3.90%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 34,537 35,804 +3.67%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small] 4,304 4,458 +3.56%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-True] 20,988 21,726 +3.52%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[False-None] 110.17 113.93 +3.41%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward] 326.47 315.53 -3.35%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None] 771.37 796.82 +3.30%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari] 26.05 26.90 +3.27%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-True] 22.17 22.89 +3.26%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None] 96.80 99.86 +3.16%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None] 269.88 278.36 +3.14%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,033 2,096 +3.12%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img] 558.58 541.31 -3.09%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] 7,136 7,355 +3.07%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch] 728.04 705.83 -3.05%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 192.63 198.43 +3.01%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari] 276.18 268.02 -2.96%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-buffers-False] 0.6013 0.6189 +2.94%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False] 35,149 36,173 +2.91%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cpu_sampler] 87.29 89.72 +2.79%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 811.84 834.33 +2.77%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True] 21.30 21.88 +2.73%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,570 3,666 +2.69%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-backward] 68.65 70.49 +2.69%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] 811.37 832.54 +2.61%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[False-None] 644.26 661.02 +2.60%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512] 47.29 48.50 +2.57%
benchmarks/test_envs_benchmark.py::test_serial 0.4272 0.4381 +2.55%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-no-buffers-False] 0.6939 0.7111 +2.47%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] 347.14 355.67 +2.46%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward] 441.68 452.37 +2.42%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch] 15.16 15.52 +2.34%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb_cuda[100-img_shape0-atari] 17.27 17.66 +2.27%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cuda_samp... 1,442 1,474 +2.27%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward] 331.00 338.49 +2.26%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-buffers-True] 0.5228 0.5339 +2.12%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None] 345.04 352.33 +2.11%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cuda] 2,167 2,212 +2.05%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None] 676.95 690.73 +2.03%
benchmarks/test_collectors_benchmark.py::test_single_with_rb 5.9727 6.0919 +2.00%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] 1,932 1,893 -1.98%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-serial-no-buffers-True] 0.6057 0.6175 +1.96%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-memmap_cpu_storage_cud... 973.38 954.32 -1.96%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None] 729.88 743.99 +1.93%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] 52.44 53.44 +1.89%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[100-img_shape0-atari] 16.60 16.90 +1.82%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None] 828.40 843.46 +1.82%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-100000-10000-100-False] 51.06 51.98 +1.81%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari] 704.47 691.73 -1.81%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[100-img_shape0-atari] 30.02 30.56 +1.79%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True] 0.5441 0.5539 +1.79%
benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True] 46.98 47.81 +1.75%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant] 4,762 4,845 +1.74%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[200-img_shape1-large_batch] 13.45 13.67 +1.63%
benchmarks/test_collectors_benchmark.py::test_single_pixels 6.3022 6.4048 +1.63%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None] 513.67 521.94 +1.61%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape1-atari] 4,398 4,327 -1.61%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None] 864.60 878.47 +1.60%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch] 322.78 317.63 -1.59%
benchmarks/test_collectors_benchmark.py::test_single 6.7736 6.8794 +1.56%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari] 639.29 629.37 -1.55%
... ... ... Showing 120 of 202 comparisons, sorted by absolute change.

vmoens
vmoens approved these changes Jun 17, 2026
View reviewed changes

@vmoens vmoens left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM thanks!

@vmoens vmoens merged commit ae403e6 into pytorch:main Jun 17, 2026
105 of 110 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@vmoens vmoens vmoens approved these changes

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/jumanji Triggers jumanji environment tests only Environments Adds or modifies an environment wrapper Examples Feature New feature Trainers Transforms tutorials/

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@itwasabhi @vmoens