Policy Gradient Theorem

• Related: Reinforcement Learning, RL, PPO

The Core Problem

In Reinforcement Learning, we want to maximize the expected total reward, denoted as $J(\theta )$. The objective function is defined as: $J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]$ where $\tau$ is a trajectory (a sequence of states and actions), ${\pi }_{\theta }$ is our policy parameterized by $\theta$, and $R(\tau )$ is the total reward of that trajectory.

To maximize this, we need the gradient of the objective with respect to our parameters: ${\nabla }_{\theta }J(\theta )$.

The Dilemma

The expected return depends on two things:

1. Our policy ${\pi }_{\theta }(a|s)$ (which we control).
2. The environment’s transition dynamics $P(s^{\prime }|s,a)$ (which we do not know and cannot differentiate).

If we naively try to take the derivative of the expectation, we run into the problem that the distribution of trajectories itself depends on $\theta$.

The Theorem

The Policy Gradient Theorem proves that we can compute the gradient of the expected return without knowing the derivative of the state distribution (the environment dynamics).

The gradient elegantly simplifies to: ${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[{\sum }_{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )\right]$

The Derivation (Log-Derivative Trick)

The proof relies on a simple calculus trick.

1. Expand the expectation into a sum over all trajectories: ${\nabla }_{\theta }J(\theta )={\nabla }_{\theta }{\sum }_{\tau }P(\tau |\theta )R(\tau )={\sum }_{\tau }{\nabla }_{\theta }P(\tau |\theta )R(\tau )$
2. Apply the log-derivative trick: Recall that ${\nabla }_x\log f(x)=\frac{{\nabla }_{x}f(x)}{f(x)}$, which means ${\nabla }_{x}f(x)=f(x){\nabla }_x\log f(x)$. Substitute this into our equation: ${\sum }_{\tau }P(\tau |\theta ){\nabla }_{\theta }\log P(\tau |\theta )R(\tau )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\nabla }_{\theta }\log P(\tau |\theta )R(\tau )]$
3. Decompose the trajectory probability: The probability of a trajectory is the product of initial state probability, policy probabilities, and environment transition probabilities: $P(\tau |\theta )=P(s_0){\prod }_{t=0}^T{\pi }_{\theta }(a_t|s_t)P(s_{t+1}|s_t,a_t)$
4. Take the log: The product becomes a sum: $\log P(\tau |\theta )=\log P(s_0)+{\sum }_{t=0}^T\log {\pi }_{\theta }(a_t|s_t)+{\sum }_{t=0}^T\log P(s_{t+1}|s_t,a_t)$
5. Take the gradient: Because the initial state and environment dynamics do not depend on $\theta$, their derivatives are zero. This is the crucial step. They vanish entirely: ${\nabla }_{\theta }\log P(\tau |\theta )={\sum }_{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)$

Why this is brilliant

• We completely bypassed the environment dynamics.
• We only need to differentiate our own policy network.
• This formula mathematically justifies trial-and-error: we run trajectories in the environment, collect rewards $R(\tau )$, and push the log-probabilities of the actions we took up or down proportional to the reward received.