Skip to content

fix(tanh_normal): keep TanhNormal finite on saturated samples - #31

Open
mctigger wants to merge 1 commit into
mainfrom
fix/tanh-normal-saturation
Open

fix(tanh_normal): keep TanhNormal finite on saturated samples#31
mctigger wants to merge 1 commit into
mainfrom
fix/tanh-normal-saturation

Conversation

@mctigger

Copy link
Copy Markdown
Owner

float32 tanh rounds to exactly ±1 once |x| exceeds about 9, which a Normal with a scale of a few units samples routinely. ClampedTanhTransform fed such values through a bare atanh (±inf) and a log(1 - y² + 1e-6) log-det (floored at log 1e-6 with a 1/eps derivative), so log_prob became -inf and the Monte-Carlo entropy/mode of the SamplingDistribution wrapper went non-finite with NaN gradients. A TanhNormal actor head initialised at std ≈ 5 hit this on every state.

Follow DreamerV2's TanhBijector:

  • _inverse clamps y to the largest magnitude strictly below 1 that its dtype represents before atanh, so x stays finite and the gradient through a saturated sample is zero instead of infinite;
  • log_abs_det_jacobian uses 2 (log 2 - x - softplus(-2x)) evaluated on x (the same identity torch's TanhTransform uses), exact for every finite x.

Off saturation the density is unchanged (matches the analytic form to ~1e-6). Adds transform-level and distribution-level regression tests for the saturated regime, including half precision.

float32 tanh rounds to exactly ±1 once |x| exceeds about 9, which a Normal
with a scale of a few units samples routinely. ClampedTanhTransform fed such
values through a bare atanh (±inf) and a log(1 - y² + 1e-6) log-det (floored
at log 1e-6 with a 1/eps derivative), so log_prob became -inf and the
Monte-Carlo entropy/mode of the SamplingDistribution wrapper went non-finite
with NaN gradients. A TanhNormal actor head initialised at std ≈ 5 hit this on
every state.

Follow DreamerV2's TanhBijector:
- _inverse clamps y to the largest magnitude strictly below 1 that its dtype
  represents before atanh, so x stays finite and the gradient through a
  saturated sample is zero instead of infinite;
- log_abs_det_jacobian uses 2 (log 2 - x - softplus(-2x)) evaluated on x
  (the same identity torch's TanhTransform uses), exact for every finite x.

Off saturation the density is unchanged (matches the analytic form to ~1e-6).
Adds transform-level and distribution-level regression tests for the
saturated regime, including half precision.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MUeQAbKsxvBA7JFLTU5vDv
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant