skip to content

Why must a tanh-squashed continuous action correct its log-probability?

level: seniorimportance: nice to knowfreq 26%

answer

  1. a nonlinear map moves probability mass
  2. change of variables, Jacobian factor
  3. derivative of tanh is 1 - tanh squared
  4. subtract log of that per coordinate
  5. silently corrupts entropy, not the bounds

basics

~20 s

Squashing is a nonlinear change of variables, so it compresses probability mass near the bounds. The density of the squashed action equals the pre-squash density divided by the derivative 1 - tanh(u)^2, so the log-probability must have log(1 - tanh(u)^2) subtracted per action coordinate.

solid answer

~60 s

An actor for a bounded plant — say a 7-DoF arm commanded in joint torques with hard physical limits — samples an unbounded value `u` and emits `a = tanh(u)` scaled to the limits. `tanh` is monotone but not measure-preserving: near the bounds a wide band of `u` maps into a sliver of `a`, so density piles up there. The change-of-variables rule gives `p(a) = p(u) / |da/du|` with `da/du = 1 - tanh(u)^2`, hence log p(a) = log p(u) - sum over coordinates of log(1 - tanh(u)^2) Omitting that term does not break the actions — they are still bounded — it breaks every quantity computed *from* the density: entropy estimates, the temperature controller that tracks a target entropy, and the log-density term in an entropy-augmented critic target. The bug is silent, which is what makes it a good interview probe. Numerically, `1 - tanh(u)^2` underflows once `u` saturates, so the correction is computed with the stable identity `2 * (log 2 - u - softplus(-2u))` rather than by evaluating the square directly.

code

python · 14 lines
python
import math

mu, sigma = 0.6, 0.8  # pre-squash Gaussian for one action coordinate

def log_gauss(u):
    return -0.5 * ((u - mu) / sigma) ** 2 - math.log(sigma * math.sqrt(2 * math.pi))

du = 1e-5
print(' u   action   measured da/du   1-tanh^2   corrected log p(a)')
for u in (0.0, 1.0, 2.0, 3.0):
    stretch = (math.tanh(u + du / 2) - math.tanh(u - du / 2)) / du
    exact = 1 - math.tanh(u) ** 2
    corrected = log_gauss(u) - math.log(exact)
    print(f'{u:4.1f} {math.tanh(u):8.4f} {stretch:14.6f} {exact:11.6f} {corrected:16.4f}')

go deeper

for a junior

Recall that bounded actions are produced by passing an unbounded output through a saturating squash, and that doing so changes the probability of the action, so the recorded log-probability has to be adjusted.

for a middle

Be able to state the change-of-variables rule: the squashed density is the original divided by the derivative of the squash, so the log-probability loses the log of that derivative, summed over action coordinates.

for a senior

Show that you know what breaks and how silently: entropy estimates and any temperature controller reading them are corrupted worst exactly where the policy saturates. Mention the numerically stable form of the correction.

for a principal

Own the class of bug: a correctness defect that leaves losses finite and training apparently healthy while quietly moving the objective. Be ready to say what invariants or logged diagnostics you would require so a team catches this kind of density error before a deployment does.

## Why squash at all Real actuators have limits: a joint torque, a valve opening, a steering angle all live in a bounded interval. An unbounded policy output has to be mapped into that interval somehow. The two candidates are clipping and squashing, and they are not equivalent: - **Clipping** (`a = clip(u, -1, 1)`) has zero derivative outside the range, so any sample that lands outside contributes no gradient. Worse, it collapses a whole tail of the pre-squash distribution onto the single boundary value, producing a genuine point mass at the bound — the resulting distribution has no density there at all, and the log-density of a boundary action is undefined. - **Squashing** (`a = tanh(u)`, then rescaled to the physical range) is smooth, strictly monotone and invertible, has a nonzero derivative everywhere, and keeps a proper density on the open interval. The boundary is approached but never reached. Squashing is what bounded-action stochastic actors use, and it comes with a bill. ## The change-of-variables bill If `a = f(u)` with `f` smooth and strictly monotone, the density transforms as p_a(a) = p_u(u) / |df/du| evaluated at u = f^{-1}(a) and in log form log p_a(a) = log p_u(u) - log |df/du| For `f = tanh`, the derivative is `df/du = 1 - tanh(u)^2`, a number in (0, 1] that approaches 0 as `|u|` grows. Because that divisor is less than one, the squashed density is *larger* than the pre-squash density — mass gets compressed into a shorter interval — and the correction term is subtracted from the log-probability. For a multi-dimensional action with per-coordinate squashing the map is coordinate-wise, so the Jacobian is diagonal and its log-determinant is just the sum of the per-coordinate terms: log p(a) = log p(u) - sum_i log( 1 - tanh(u_i)^2 ) The geometric picture is worth carrying: take a small interval `du` around some `u`. Near `u = 0` it maps to an interval of nearly the same width. Near `u = 3` it maps to an interval roughly a hundred times narrower. Same probability, much less room, therefore much higher density. ## What actually breaks when you forget it Nothing visibly. The actions are still inside the bounds, training still runs, the losses are still finite. What is wrong is everything that reads the density: - **Entropy is mis-measured.** The uncorrected `-log p(u)` overstates the entropy of the emitted action, and it overstates it most exactly when the policy has saturated toward a bound. - **A temperature controller tracking a target entropy is driven by the wrong signal.** Since the error grows with saturation, the controller reads more randomness than the plant is receiving and drives the temperature the wrong way; the policy saturates further, the reported entropy stays plausible, and the agent quietly becomes a bang-bang controller. - **An entropy-augmented critic target is wrong** by the same term, so the soft values are biased in a state-dependent way. ## Numerical stability `1 - tanh(u)^2` underflows to exactly zero in floating point once `|u|` is moderately large, and `log 0` ruins the update. Two standard remedies: add a small epsilon inside the log, or use the exact identity log(1 - tanh(u)^2) = 2 * ( log 2 - u - softplus(-2u) ) which is stable for large positive and negative `u` alike. The identity is worth being able to reproduce; it follows from writing `1 - tanh(u)^2` as `4 e^{-2u} / (1 + e^{-2u})^2` and taking logs. A related practical point: because the correction blows up as `|u|` grows, a policy that saturates hard is being charged a large log-density, and it is common to lightly penalise the magnitude of the pre-squash output to keep the actor out of the saturation regime — where, as with any saturating nonlinearity, the gradient through it also vanishes. ## Evaluation-time action At deployment you usually want one action per state rather than a sample. Because `tanh` is strictly monotone it preserves quantiles, so `tanh(mu)` — the squashed mean of the pre-squash Gaussian — is exactly the **median** of the action distribution. It is not the mean of the squashed distribution, and it is not the mode either, precisely because of the Jacobian factor discussed above. The median is the right and available choice, and saying so is a good sign that you have understood the transform rather than memorised the formula.

  • Why not just clip the action into its range instead of squashing it?
    Clipping has zero gradient outside the range, so out-of-range samples teach the actor nothing, and it maps a whole tail onto the boundary value, creating a point mass. A distribution with a point mass has no density at the bound, so its log-probability is undefined there and every entropy or soft-value computation breaks. Squashing stays smooth, invertible and differentiable everywhere.
  • How do you compute the correction without it underflowing?
    1 - tanh(u)^2 becomes exactly zero in floating point once the input saturates, so log of it is not usable directly. Either add a small epsilon inside the log, or use the equivalent form 2*(log 2 - u - softplus(-2u)), which is numerically stable for large inputs of either sign and gives the same value.
  • At evaluation time, what single action should a squashed stochastic policy emit?
    The squashed mean of the pre-squash distribution. Because the squash is strictly monotone it preserves quantiles, so that value is exactly the median of the action distribution. It is not the mean of the squashed distribution and not its mode — both are shifted by the same Jacobian factor that forces the log-probability correction.

Squashing is like printing a map on a rubber sheet and then compressing the edges: the same countries are there, but the ones near the border are drawn far denser, and any measurement you take must divide by how much the sheet was stretched.

saying these in an interview costs you the question

  • Says the correction is unnecessary because the squash is monotone
  • Thinks omitting it lets actions escape their bounds
  • Adds the Jacobian log term instead of subtracting it
  • Computes log of one minus tanh squared with no stabilisation
  • Calls the squashed mean the mean of the action distribution

context

resolving links…