Log-Probability of a Softmax Policy

1. Background Knowledge

In reinforcement learning, a policy $\pi(a|s)$ maps states to a probability distribution over actions. When actions are discrete, this is typically parameterized by a softmax over a vector of unnormalized scores called logits. If $h_b$ is the logit for action $b$, the probability of action $a$ is:

$$\pi(a) = \frac{e^{h_a}}{\sum_b e^{h_b}}$$

Policy-gradient algorithms (e.g., REINFORCE, A2C) require the log-probability $\log \pi(a)$ rather than the probability itself. This is because gradients of the log-likelihood are numerically more stable and algebraically simpler. Taking the log of the softmax gives:

$$\log \pi(a) = h_a - \log \sum_b e^{h_b}$$

The term $\log \sum_b e^{h_b}$ is the log-sum-exp (LSE) of the logits. Directly computing $\sum_b e^{h_b}$ can overflow for large logits or underflow for very negative ones. The log-sum-exp trick stabilizes this by shifting all logits by a constant $m$ (typically $m = \max_b h_b$):

$$\text{LSE}(h) = m + \log \sum_b e^{h_b - m}$$

Since $m$ is the maximum, every $h_b - m \le 0$, so each exponential is in $(0, 1]$ and the sum cannot overflow. The shift $m$ is added back to preserve the correct value.

2. Algorithm Approach

The core pattern is the log-sum-exp trick:

1. Find the maximum logit $m = \max_b h_b$.
2. Subtract $m$ from every logit to get shifted values $h_b - m$.
3. Exponentiate the shifted values and sum them.
4. Take the natural log of that sum.
5. Add $m$ back to get the stable LSE.
6. Subtract the LSE from the chosen action's logit $h_a$ to get $\log \pi(a)$.

This is a standard numerical-stability technique used throughout machine learning (e.g., in cross-entropy loss, attention mechanisms, and softmax layers).

3. Step-by-Step Strategy

• Extract the max: Compute $m = \max(\texttt{logits})$. This is a single pass over the array.
• Shift and exponentiate: For each logit $h_b$, compute $e^{h_b - m}$. Because $h_b - m \le 0$, these values are bounded above by 1.
• Sum the exponentials: Accumulate $\sum_b e^{h_b - m}$.
• Log and unshift: Compute $m + \log(\text{sum})$ to obtain the stable LSE.
• Subtract for log-probability: Return $\texttt{logits}[a] - \text{LSE}$.

In code, this looks roughly like:

import numpy as np

def log_prob(logits, a):
 m = np.max(logits)
 shifted = logits - m
 lse = m + np.log(np.sum(np.exp(shifted)))
 return logits[a] - lse

The key insight is that you never compute $\sum e^{h_b}$ directly; you always work with the shifted version.

4. Common Pitfalls

• Skipping the shift: Computing $\log(\sum e^{h_b})$ directly will overflow if any $h_b > \sim 709$ (the log of the max float64) or underflow if all $h_b$ are very negative.
• Using the wrong max: The shift must be the maximum of the logits, not the mean or zero. Using a smaller shift can still cause overflow.
• Index confusion: Make sure you subtract the LSE from logits[a] (the chosen action's logit), not from the max or some other element.
• Integer logits: If logits is an integer array, logits - m stays integer and np.exp may behave unexpectedly. Cast to float first.
• Empty or single-element edge cases: With one action, the LSE equals that logit and $\log \pi(a) = 0$, which is correct ($\pi = 1$). Verify your implementation handles this gracefully.

5. Time & Space Complexity

• Time: $O(n)$ where $n$ is the number of actions (length of logits). You scan once for the max, once for the shifted exponentials, and once for the sum. These are all linear passes.
• Space: $O(n)$ if you materialize the shifted array, or $O(1)$ extra space if you compute the sum in a single loop without storing intermediate values. In practice, vectorized NumPy implementations use $O(n)$ temporary storage.