[BugFix] Fix TanhNormal sample scoring at tanh saturation

[gtnv]

Description

Fixes incorrect TanhNormal scores and gradients when finite-precision tanh saturation prevents atanh(action) from recovering the sampled preimage.

This adds atomic rsample_and_log_prob() and sample_and_log_prob() operations that compute the action and its score from the same pre-tanh sample without persistent cache state. Ordinary log_prob(value) remains history-independent and continues to handle externally supplied actions.

Shared sampling utilities apply the atomic operation when a distribution supports it and retain the existing sample-then-score fallback for other distributions. Composite distributions are handled component by component while preserving the configured log-probability aggregation behavior.

The joint operation is used by SafeProbabilisticModule and by internally generated sample-and-score paths in SAC, CQL, CrossQ, REDQ, A2C, PPO, Decision Transformer, GRPO, and V-trace. The public operations and their gradient, shape, and fallback contracts are included in the distribution API documentation.

No clamps, scale caps, flags, or training-specific defaults are added.

Validation

• Atomic TanhNormal, composite-distribution, and history-independence coverage: 19 passed.
• SafeProbabilisticModule integration: 2 passed.
• SAC integration, both loss versions: 2 passed.
• CQL actor integration: 1 passed.
• Project-pinned ufmt check passes for all edited Python files.
• git diff --check passes.
• The expanded implementation was additionally exercised by the author on H100 hardware and across 5,835 CPU tests.

Closes #2199.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

As of commit 870911a with merge base d7659c7 ():

NEW FAILURE - The following job has failed:

• Unit-tests on Linux / tests-cpu (3.10, bulk) / linux-job (gh)
  test/rb/test_storages.py::TestStorages::test__rand_given_ndim_recompile

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

[theap06]

LGTM! Thanks

[vmoens]

Reviewed latest head be7a48dba7, including the CUDA-capture follow-up commit and the current approval. The single-sample cache is invalidated on update(), detects eager in-place mutation via tensor versioning, uses a snapshot for compile/inference/capture paths, and falls back to the ordinary inverse path for external or stale samples; custom affine bounds are included because the composed transform is cached as a whole. The focused eager/compiled consistency and invalidation matrix passes locally (18 passed), and git diff --check is clean. No blocking finding from this review; the remaining platform-specific coverage is the CUDA-capture path already described in the PR.

[vmoens-ii]

[duplicated review]

Submitted from the wrong GitHub account; superseded by the review from vmoens.

[vmoens]

The fix appears correct only for the narrow rsample() -> log_prob() path on the exact tensor returned by the most recent sampling-like call. I do not think we should make these stateful semantics the default for TanhNormal.

The cache is single-entry and identity-based. A second call to rsample(), sample(), or deterministic_sample replaces it, and update() clears it. Cloning, slicing, or otherwise replacing the tensor object also misses it. Consequently, benign control-flow changes such as

a1 = dist.rsample()
a2 = dist.rsample()
log_prob1 = dist.log_prob(a1)

silently restore the saturated inverse path and its incorrect score/gradients. More generally, log_prob(value) becomes dependent on call history and Python object identity.

The gradient behavior is also inconsistent. For the unchanged cached sample, gradients with respect to loc and scale follow the cached preimage and appear mathematically correct. However, the cached log_prob is constructed from cached_preimage, not through the supplied value, so autograd.grad(log_prob.sum(), sampled_action) can report the action as unused. A cloned or external action instead uses the inverse path and has the ordinary action gradient. I do not think the same public operation should have different differentiation semantics based on identity and cache state.

There is also unconditional overhead on every sample: a full detach().clone(), retained references to the sample and preimage (and potentially their autograd graph), plus the snapshot. Compile/inference/capture additionally execute the equality reduction and selection/inverse machinery. Showing no allocator growth over repeated iterations rules out an unbounded leak, but does not establish acceptable peak memory, graph lifetime, or throughput.

I would prefer an explicit atomic operation such as rsample_and_log_prob(), computed directly from the pre-tanh value without persistent mutable state. The ordinary log_prob(external_value) can then remain a history-independent operation. TorchRL's probabilistic-module path can use the joint operation where required.

[vmoens]

Thanks for replacing the cache-based approach with an atomic sample-and-score operation. This is much safer, and the gradient path now appears correct because the sample and log-probability are computed from the same reparameterized preimage.

I still have two concerns before approving:

1. Documentation

rsample_and_log_prob() is a new public API, but it currently has only a one-line docstring. Please document:

• The numerical and gradient contract.
• Why this should be used instead of rsample() followed by log_prob().
• The meaning and shape of both returned tensors.
• Interaction with event dimensions, bounds, and safe_tanh.
• The conditions under which SafeProbabilisticModule uses this operation, and when it silently falls back to the old behavior.

An example demonstrating the saturated-action case would also be useful.

Adding it to the repo's API doc would be nice as well!

2. Library-wide adoption

The PR integrates the operation into SafeProbabilisticModule and SAC, but several objectives obtain the distribution directly and still generate a sample followed immediately by log_prob(sample). These include paths in CQL, CrossQ, REDQ, and the Monte Carlo entropy implementations in A2C, PPO, and Decision Transformer.

Those paths bypass SafeProbabilisticModule, so they can still encounter the same numerical problem with a saturated TanhNormal. Composite distributions and multi-output probabilistic modules also currently fall back to the two-step behavior.

I'm not sure what the proper API could be since that method isn't present everywhere (maybe we could expose it in all distributions within torchrl somehow!?).
We should audit the internally generated sample-and-score paths and use a shared helper (either free function or method) where applicable. This should only affect actions generated from the same distribution; scoring externally supplied or replay-buffer actions should remain a normal log_prob() call. Representative integration tests outside SAC would help ensure the fix is consistently adopted.

[gtnv]

@vmoens alright: I've implemented the changes and expanded to the other use-cases.

I also tested on h100s and double-checked that the issue reproduced on main, and that the added code works like it should.

Ran tests for tanhnormal cases, all passed.
Actor/SAC integration tests: 6 passed
CPU tests: 5,835 passed