Policy Gradient Methods › The Policy Gradient Theorem

Here is the core challenge stated precisely: we need to compute $\nabla_\theta J(\theta)$ so we can update our parameters via gradient ascent. But the expectation involves sampling from the policy itself, which depends on $\theta$. This creates a circular dependency that seems difficult to resolve. The gradient depends on how changing $\theta$ affects which sequences get sampled, but sampling is a discrete, non-differentiable operation. We cannot simply backpropagate through the sampling process the way we would through a continuous neural network layer.

The policy gradient theorem provides an elegant solution to this problem. Rather than trying to differentiate through sampling directly, it rewrites the gradient in a form that we can estimate using samples from the current policy. The key insight is that we can express the gradient of an expectation in terms of an expectation of gradients that we can compute.

Let's derive the policy gradient step by step. Starting with our objective written as an explicit sum:

$$J(\theta) = \sum_{y} \pi_\theta(y|x) R(x, y)$$

We take the gradient with respect to $\theta$. Since $R(x,y)$ is produced by a frozen reward model and does not depend on $\theta$, only the policy probability is differentiated:

$$\nabla_\theta J(\theta) = \sum_{y} \nabla_\theta \pi_\theta(y|x) \cdot R(x, y)$$

This is still a sum over all sequences, which we cannot compute directly. We need to transform it into an expectation. The trick is to multiply and divide by $\pi_\theta(y|x)$, introducing a factor of 1 without changing the value:

$$\nabla_\theta J(\theta) = \sum_{y} \pi_\theta(y|x) \cdot \frac{\nabla_\theta \pi_\theta(y|x)}{\pi_\theta(y|x)} \cdot R(x, y)$$

Now we apply the log-derivative trick (also called the likelihood ratio trick or REINFORCE trick). This technique appears throughout machine learning and statistics whenever we need to compute gradients of expectations. The core identity is:

$$\frac{\nabla_\theta \pi_\theta(y|x)}{\pi_\theta(y|x)} = \nabla_\theta \log \pi_\theta(y|x)$$

This follows directly from the chain rule of differentiation. Recall that $\frac{d}{d\theta} \log f(\theta) = \frac{1}{f(\theta)} \frac{df}{d\theta}$, which rearranges to $\frac{df/d\theta}{f(\theta)} = \frac{d}{d\theta} \log f(\theta)$. The logarithm converts a multiplicative relationship into an additive one, which proves extremely useful when we deal with sequence probabilities that are products of many terms.

Substituting back:

$$\nabla_\theta J(\theta) = \sum_{y} \pi_\theta(y|x) \cdot \nabla_\theta \log \pi_\theta(y|x) \cdot R(x, y)$$

A sum over all possible sequences weighted by their probabilities is exactly an expectation under the policy:

$$\nabla_\theta J(\theta) = \mathbb{E}_{y \sim \pi_\theta}[\nabla_\theta \log \pi_\theta(y|x) \cdot R(x, y)]$$

where:

• $\nabla_\theta J(\theta)$: gradient of the objective function with respect to parameters
• $\nabla_\theta \log \pi_\theta(y|x)$: gradient of the log-probability of the sequence, called the "score function"
• $R(x, y)$: the reward, acting as a scalar weight on the gradient

This is the policy gradient theorem for our setting. It tells us something remarkable: to compute the gradient of expected reward, we can sample sequences from our policy, compute the gradient of their log-probabilities, and multiply by their rewards. The score function $\nabla_\theta \log \pi_\theta(y|x)$ points in the direction in parameter space that would increase the probability of sequence $y$. Multiplying by the reward scales this direction: high rewards give large positive contributions, low rewards give small or negative contributions.

The transformation from an intractable sum to a Monte Carlo estimate is the key practical insight. We cannot enumerate all possible sequences; for even modest vocabulary sizes and sequence lengths, this would involve more sequences than atoms in the universe. But we can sample sequences from our policy, and the policy gradient theorem tells us that averaging over these samples gives us an unbiased estimate of the true gradient.

The unbiasedness guarantee is important. It means that if we sample enough sequences, our gradient estimate will converge to the true gradient. We might not get the exact gradient from a single sample, but in expectation across many samples, we get the right answer. This is what allows reinforcement learning to work despite the discrete, non-differentiable nature of sampling.