[BugFix] Improve runtime diagnostics and CUDA timing

[vmoens]

Summary

• warn only once per replay buffer when a sample-time batch size conflicts with the constructor value
• include the fixed shared-memory key schema when AsyncEnvPool rejects unsupported input
• add synchronized CUDA timing through the timeit sync option
• emit the nested TransformedEnv default-change warning at most once per process

Tests

• focused regression tests for all four behaviors
• test/envs/test_special.py: 146 passed, 129 skipped
• test/transforms/test_compose_and_env.py: 173 passed, 16 skipped
• legacy timeit test: 1 passed
• ufmt, flake8, pydocstyle, pyupgrade, codespell, autoflake, docstring-argument checks, and git diff --check

The full replay-buffer and utility modules have unrelated prioritized-buffer failures in this local checkout because the TorchRL C++ segment-tree extension is unavailable. The pinned pre-commit LibCST 1.0.1 source build also requires an unavailable Rust compiler, so ufmt was checked with the same pinned Black/usort versions and a compatible LibCST wheel.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 47 Pending

As of commit 269cf71 with merge base c5ca88e ():
💚 Looks good so far! There are no failures yet. 💚

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

[vmoens]

Reviewed all four fixes and the tests.

Three of the four are good as they are: the once-per-buffer batch-size warning (the getattr for buffers unpickled from older versions is appreciated), the schema-aware KeyError in the shared-slot exchange, and the process-wide flag for the TransformedEnv FutureWarning. One change is needed before merging.

Needs a fix: timeit(sync=True) crashes without CUDA. torch.cuda.synchronize() is called unconditionally in __enter__, start and elapsed (torchrl/_utils.py:290, 311, 330). On a CPU-only build it raises AssertionError: Torch not compiled with CUDA enabled (just checked on the macOS wheel), and on a CUDA build without a device it raises RuntimeError: Found no NVIDIA driver. The flag is meant to be driven from a config value in scripts that also run on CPU (the DreamerV3 tests do exactly that), so guard it with torch.cuda.is_available() or take a device and synchronise only when it is CUDA. The unit test monkeypatches torch.cuda.synchronize, so it cannot see this; please add a CPU case that passes sync=True and checks the timer still records.

Minor:

• The FutureWarning text still promises a change "in version 0.9" while we are on 0.14 and the default has not moved. Since the line is being touched, either name the release where auto_unwrap actually flips (and register it for that release) or downgrade it to a one-time informational warning; per CLAUDE.md section 12 the target version has to be real.
• On what the flag buys: Python's default filter already de-duplicates identical warnings per call site, so the process-wide flag only removes repeats coming from different call sites. N worker processes still print N warnings, which is what the 64-worker log showed. Fine as a papercut fix, just noting the papercut is not fully closed.

Cross-PR: #4268 adds unconditional torch.cuda.synchronize(device) calls in the DreamerV3 training loop for the same timing purpose, and #4265 adds its own synchronisation. Once this lands, those should use timeit(..., sync=True) behind a flag (or CUDA events) instead of hand-rolled syncs.