[Refactor] Use non_blocking transfers in distribution modules

[vmoens]

Stack from ghstack (oldest at bottom):

• [Feature] Add _Contiguous module and reshape improvements to encoders/decoders #3306
• [BugFix] Fix SliceSampler for torch.compile compatibility #3298
• [Feature] Refactor Dreamer training with async collectors, profiling, and improved config #3308
• [Feature] Add noise argument and scan mode to RSSMRollout #3307
• [Feature] Add explicit dimensions and device support to RSSM modules #3305
• [Refactor] Use compile-aware helpers in Dreamer objectives #3304
• [BugFix] Fix value functions for torch.compile compatibility #3303
• [BugFix] Fix TDLambdaEstimator for torch.compile compatibility #3302
• [BugFix] Optimize DreamerEnv to avoid CUDA sync in done checks #3301
• [Feature] Add allow_done_after_reset parameter to ModelBasedEnvBase #3300
• [BugFix] Fix ModelBasedEnvBase for torch.compile compatibility #3299
• [Refactor] Add compile-aware timing and profiling helpers #3297
• [Feature] Support callable scale in IndependentNormal and TanhNormal distributions #3296
• -> [Refactor] Use non_blocking transfers in distribution modules #3295
• [Refactor] Add WEIGHT_SYNC_TIMEOUT constant for collector weight synchronization #3294
• [Refactor] Remove MUJOCO_EGL_DEVICE_ID auto-setting from torchrl init #3293

Add non_blocking=True when transferring tensors to CUDA devices
for better performance by allowing async transfers.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 4 New Failures, 1 Cancelled Job, 11 Pending, 1 Unrelated Failure

As of commit 58a3d69 with merge base df00d61 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t d574446855d7eeae9d20bf531ba7f8feaa8f9d4cdd66a8796e8940f4049dfb67 /exec failed with exit code 1
• Libs Tests on Linux / unittests-gym (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t ffa061e19b786db58a6107ce03e565048ccebeb25a64b5b197340f97ee0c222e /exec failed with exit code 1
• Lint / python-source-and-configs / linux-job (gh)
  RuntimeError: Command docker exec -t 0c4ec0ac397386cb00774c8a1a8411844bc875cd7ba109a2509acfef94d697ba /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  ##[error]The operation was canceled.

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

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

• Unit-tests on Linux / tests-cpu (3.14) / linux-job (gh) (trunk failure)
  test/test_transforms.py::TestConditionalSkip::test_trans_parallel_env_check[False-False]

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