[BugFix] Fix TanhNormal sample scoring at tanh saturation

[gtnv]

Description

Fix TanhNormal.rsample() -> log_prob() when finite-precision tanh saturation prevents atanh(action) from recovering the generating preimage, corrupting SAC scores and gradients.

TanhNormal now retains one forward-produced preimage and reuses it only for its unchanged sample. External or mutated values use the ordinary inverse path, and update() clears the cache. Compilation, inference, and CUDA capture validate the sample through a tensor snapshot. There are no clamps, flags, or training-specific defaults.

Testing

• Direct final-head checks: main misses the score by 16,312 and the scale gradient by 165,992; this PR is within 7.63e-6 and 9.54e-7 in eager, Dynamo eager, and Inductor.

• Matched MultiWalker SAC, 3 seeds × 2M frames: every main run develops score corruption and exploding losses; corrected runs keep score error at 0.0 and raw critic loss at or below 23.6. E2E ran at efb3a592b; final head only adds the CUDA-capture dispatch validated below.

• Final head be7a48dba: focused H100 18 passed; distributions 974 passed; pre-commit passes. Objectives at the behavior-identical E2E head: 7528 passed, 2324 skipped, 6 deselected.
• Safe/unsafe tanh with ordinary/custom bounds passes inference and CUDA graph within 7.63e-6; scale-gradient error is at most 4.55e-7, with zero allocated-memory growth across 20,000 iterations.

The cached sample keeps its generating loc/scale gradient path. Cloned and external actions retain ordinary action gradients.

Closes #2199

cc @vmoens @theap06

[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.

✅ No Failures

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

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.