[Feature] Add joint sample-and-score distribution adapter

[vmoens]

Summary

• add ensure_rsample_and_log_prob() as an identity-preserving capability adapter
• expose sample_and_log_prob() and rsample_and_log_prob() on arbitrary PyTorch and third-party distributions through a lightweight proxy
• adapt distributions at TorchRL probabilistic-module and objective boundaries, allowing consumers to use the method contract without concrete-distribution branches
• add unwrap_distribution() for unavoidable exact-type dispatch such as torch.distributions.kl_divergence
• document the method-oriented API and preserve the existing free-function fallback
• cover native identity, wrapped gradients, torch.compile, serialization, composites, non-reparameterizable distributions, actor/objective integration, and sampling overhead

Motivation

#4080 introduces atomic sample-and-score operations for TanhNormal and generic free-function fallbacks. Distribution-consuming modules are modular, however, and cannot assume which concrete distribution they will receive. This follow-up establishes a uniform internal method contract:

action, log_prob = dist.rsample_and_log_prob()

Distributions that already provide the method retain object identity. Other distributions are adapted without monkey-patching PyTorch classes or maintaining TorchRL copies of the torch.distributions namespace.

Compatibility

The proxy delegates ordinary distribution behavior, including sampling, scoring, entropy, support, statistics, expansion, enumeration, mapping behavior for composite distributions, and distribution-specific attributes. isinstance checks remain transparent. APIs that dispatch on exact concrete types must use unwrap_distribution(); the PPO KL path does this explicitly.

Structured CompositeDistribution samples continue to use the centralized component aggregation in the existing free functions. Distributions with has_rsample=False expose sample_and_log_prob() and continue to reject reparameterized sampling.

Stacking

This draft is stacked on #4080 and currently includes its commits. Once #4080 lands, this branch should be rebased onto main, leaving only the adapter follow-up in the diff.

Validation

• pytest -q test/test_distributions.py test/modules/test_actor.py: 1087 passed, 1 skipped
• pytest -q test/objectives/test_cql.py: 408 passed, 3 skipped
• representative SAC, CrossQ, REDQ, A2C, OnlineDT, and V-trace regressions passed
• a direct non-skipped KLPENPPOLoss execution passed, covering explicit unwrapping for KL dispatch
• sampling microbenchmark: adapted method 19.5 us mean; free function 20.0 us mean in the local run
• pinned ufmt, flake8, pydocstyle, and git diff --check passed

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit 2373262 with merge base c5ab3c9 ():

NEW FAILURE - The following job has failed:

• Unit-tests on Linux / tests-cpu (3.10, bulk) / linux-job (gh)
  API rate limit exceeded

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.13, bulk) / linux-job (gh) (trunk failure)
  test/rb/test_storages.py::TestStorages::test__rand_given_ndim_recompile

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