The Policy Gradient Theorem

Value-based methods like DQN have limitations, especially with continuous actions or when stochastic policies are needed. Policy gradient methods offer a different path: directly optimizing the parameters of the policy itself. But how do we actually do that? How do we know which way to adjust the policy parameters $\theta$ to improve performance?

Our goal is to find policy parameters $\theta$ that maximize the expected total return. We can define an objective function, often denoted as $J(\theta)$, which represents the expected return when following policy $\pi_\theta$. For instance, in an episodic task, this could be the expected cumulative discounted reward starting from an initial state distribution:

$$J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} [ \sum_{t=0}^{T-1} \gamma^t r_{t+1} ] = \mathbb{E}_{\tau \sim \pi_\theta} [ G_0 ]$$

Here, $\tau$ represents a full trajectory $(s_0, a_0, r_1, s_1, ...)$ generated by following the policy $\pi_\theta(a|s)$, $\gamma$ is the discount factor, $r_{t+1}$ is the reward received after taking action $a_t$ in state $s_t$, and $G_0$ is the total discounted return for the entire episode.

To maximize $J(\theta)$, we can use gradient ascent. The update rule for the policy parameters would look like this:

$$\theta \leftarrow \theta + \alpha \nabla_\theta J(\theta)$$

where $\alpha$ is the learning rate and $\nabla_\theta J(\theta)$ is the gradient of the objective function with respect to the policy parameters. The core challenge lies in calculating this gradient, $\nabla_\theta J(\theta)$. How does changing the policy parameters $\theta$ affect the expected total return? This involves not only how the policy chooses actions but also how the resulting trajectories and reward distributions change, which seems complicated to compute directly.

This is where the Policy Gradient Theorem comes in. It provides a way to compute this gradient without needing to know or differentiate the environment's dynamics (the transition probabilities $p(s'|s,a)$). This theorem is foundational for all policy gradient algorithms.

One common form of the Policy Gradient Theorem states that the gradient of the objective function is given by:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} [ \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t) G_t ]$$

Let's break down this important expression:

1. $\mathbb{E}_{\tau \sim \pi_\theta} [ ... ]$: The expectation is taken over trajectories $\tau$ generated by following the current policy $\pi_\theta$. This means we can estimate the gradient by running the policy in the environment and collecting sample trajectories.
2. $\sum_{t=0}^{T-1} ...$: We sum over all time steps within a trajectory.
3. $\nabla_\theta \log \pi_\theta(a_t|s_t)$: This is the gradient of the log-probability of taking the specific action $a_t$ in state $s_t$ under the current policy parameters $\theta$. This term tells us the direction in parameter space $\theta$ that most increases the probability of selecting action $a_t$ when in state $s_t$. It's calculable because we defined the policy $\pi_\theta$ (e.g., using a neural network whose output determines action probabilities) and can therefore compute its derivatives with respect to its parameters $\theta$.
4. $G_t = \sum_{k=t}^{T-1} \gamma^{k-t} r_{k+1}$: This is the total discounted return received from time step $t$ onwards in the trajectory $\tau$. It represents the long-term consequence of having taken action $a_t$ in state $s_t$.

The Log-Derivative Trick

You might wonder how we arrived at the $\nabla_\theta \log \pi_\theta(a_t|s_t)$ term. It comes from a useful identity sometimes called the "log-derivative trick" or "likelihood ratio trick". The gradient of the policy $\nabla_\theta \pi_\theta(a|s)$ can be rewritten as:

$$\nabla_\theta \pi_\theta(a|s) = \pi_\theta(a|s) \frac{\nabla_\theta \pi_\theta(a|s)}{\pi_\theta(a|s)} = \pi_\theta(a|s) \nabla_\theta \log \pi_\theta(a|s)$$

This reformulation is significant because it allows us to express the gradient $\nabla_\theta J(\theta)$ as an expectation involving $\pi_\theta$ itself. This means we can estimate the gradient using samples (actions $a_t$) drawn from the current policy $\pi_\theta$, without needing to explicitly calculate how changes in $\theta$ affect the state visitation distribution.

Intuition Behind the Theorem

The Policy Gradient Theorem provides an intuitive update mechanism. The term $\nabla_\theta \log \pi_\theta(a_t|s_t)$ indicates how to change $\theta$ to increase (or decrease) the probability of action $a_t$ in state $s_t$. The term $G_t$ acts as a weighting factor for this update direction.

• If the return $G_t$ following action $a_t$ was high (good outcome), the gradient update pushes $\theta$ in the direction indicated by $\nabla_\theta \log \pi_\theta(a_t|s_t)$, effectively increasing the probability of taking action $a_t$ again in state $s_t$ in the future.
• If the return $G_t$ was low (bad outcome), the update pushes $\theta$ in the opposite direction, decreasing the probability of taking $a_t$ in state $s_t$.

Essentially, the policy gradient method reinforces actions proportionally to the total return they ultimately lead to.

The Policy Gradient Theorem provides the theoretical justification for algorithms like REINFORCE, which we will explore next. REINFORCE directly uses Monte Carlo sampling (running full episodes) to estimate the expectation in the theorem and update the policy parameters. As we'll see, while straightforward, this approach can suffer from high variance in the gradient estimates, motivating further refinements like using baselines.