Introduction to Policy Gradient Methods

Previous sections reviewed value-based methods like Q-Learning, where the core idea is to learn the value of state-action pairs, $Q(s, a)$, and then derive a policy (often implicitly, like epsilon-greedy) based on these values. Policy Gradient methods offer a fundamentally different approach: they directly parameterize the policy and optimize its parameters to maximize expected return.

Let's denote the policy as $\pi_\theta(a|s)$, representing the probability of taking action $a$ in state $s$ under policy parameters $\theta$. The goal is to find the parameters $\theta$ that maximize a performance objective function, typically the expected total discounted reward starting from an initial state distribution. We can define this objective as:

$$J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T} \gamma^t r(s_t, a_t) \right]$$

Here, $\tau = (s_0, a_0, r_1, s_1, a_1, r_2, \dots)$ represents a trajectory generated by following policy $\pi_\theta$, $\gamma$ is the discount factor, and $r(s_t, a_t)$ is the reward received after taking action $a_t$ in state $s_t$.

The core challenge is computing the gradient of this objective function with respect to the policy parameters, $\nabla_\theta J(\theta)$. Directly differentiating the objective seems difficult because the expectation depends on the distribution of trajectories induced by $\pi_\theta$, which itself depends on $\theta$. The Policy Gradient Theorem provides a convenient way around this. It states that the gradient of the objective function can be expressed as:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t) Q^{\pi_\theta}(s_t, a_t) \right]$$

Alternatively, using the expectation notation over states and actions visited under $\pi_\theta$:

$$\nabla_\theta J(\theta) = \mathbb{E}_{s \sim d^{\pi_\theta}, a \sim \pi_\theta(a|s)} \left[ \nabla_\theta \log \pi_\theta(a|s) Q^{\pi_\theta}(s, a) \right]$$

where $d^{\pi_\theta}(s)$ is the stationary state distribution under policy $\pi_\theta$. The term $\nabla_\theta \log \pi_\theta(a|s)$ is sometimes called the "score function". It tells us how to adjust the parameters $\theta$ to increase the log-probability of taking action $a$ in state $s$. The theorem tells us to weight this direction by the action-value function $Q^{\pi_\theta}(s, a)$. Intuitively, this makes sense: we want to increase the probability of actions that lead to higher-than-average returns and decrease the probability of actions that lead to lower-than-average returns.

The REINFORCE Algorithm

The Policy Gradient Theorem provides the form of the gradient, but we still need $Q^{\pi_\theta}(s_t, a_t)$, which is generally unknown. The REINFORCE algorithm (also known as Monte Carlo Policy Gradient) offers a simple way to estimate this gradient using sampled trajectories.

REINFORCE uses the complete return from time step $t$ onwards, $G_t = \sum_{k=t}^{T} \gamma^{k-t} r_{k+1}$, as an unbiased sample estimate of $Q^{\pi_\theta}(s_t, a_t)$. Since we are sampling trajectories, we can approximate the expectation with an average over sampled trajectories. For a single trajectory, the update direction for each time step $t$ is $G_t \nabla_\theta \log \pi_\theta(a_t|s_t)$.

The basic REINFORCE algorithm proceeds as follows:

1. Initialize policy parameters $\theta$.
2. Repeat: a. Generate one full episode trajectory $\tau = (s_0, a_0, r_1, \dots, s_T, a_T, r_{T+1})$ by following the current policy $\pi_\theta$. b. For each time step $t = 0, 1, \dots, T$: i. Calculate the return from that time step onwards: $G_t = \sum_{k=t}^{T} \gamma^{k-t} r_{k+1}$. ii. Calculate the policy gradient estimate for this step: $g_t = G_t \nabla_\theta \log \pi_\theta(a_t|s_t)$. c. Update the policy parameters using the accumulated gradients from the episode (or often, by summing $g_t$ across the episode and taking a step): $\theta \leftarrow \theta + \alpha \sum_{t=0}^{T} g_t$, where $\alpha$ is the learning rate.

Here's a visual representation of the REINFORCE update logic:

A schematic overview of the REINFORCE algorithm: generating a trajectory using the current policy $\pi_\theta$, calculating returns $G_t$ and score functions $\nabla \log \pi_\theta(a_t|s_t)$ for each step, and then updating the policy parameters $\theta$.

Advantages and Disadvantages

Policy gradient methods like REINFORCE have several advantages:

1. Direct Optimization: They directly optimize the quantity we care about: the expected return.
2. Continuous Actions: They naturally handle continuous action spaces, where finding the maximum over actions (as required in Q-learning's update $\max_{a'} Q(s', a')$) is problematic. We can define $\pi_\theta(a|s)$ using distributions like Gaussians, whose parameters are outputs of a neural network.
3. Stochastic Policies: They can learn stochastic policies, which can be advantageous in partially observable environments or when facing adversaries. Value-based methods often learn deterministic policies (derived from the greedy choice over Q-values).

However, REINFORCE also suffers from a significant disadvantage:

1. High Variance: The Monte Carlo estimate $G_t$ of $Q^{\pi_\theta}(s_t, a_t)$ can have very high variance. A single trajectory might be unusually good or bad purely by chance, leading to noisy gradient estimates and slow, unstable learning.

This high variance is a primary motivation for developing more advanced policy gradient methods, such as Actor-Critic algorithms (covered in Chapter 3), which use a learned value function (the critic) to reduce the variance of the policy gradient estimate (from the actor).

Policy Representation with Function Approximation

Just as value functions can be approximated using linear functions or neural networks for large state spaces, the policy $\pi_\theta(a|s)$ can also be represented using function approximators. In deep reinforcement learning, $\pi_\theta$ is typically represented by a neural network that takes the state $s$ as input and outputs the parameters of the action distribution (e.g., probabilities for discrete actions, or mean and standard deviation for continuous actions). The parameters $\theta$ then correspond to the weights and biases of this neural network. This integration with function approximation allows policy gradient methods to scale to complex problems with high-dimensional state and action spaces, forming the bedrock of many advanced techniques explored later in this course.