Skip to content

Policy Gradient Fundamentals

The foundation for TRPO, PPO, and everything that follows

Prerequisites: Basic RL concepts (states, actions, rewards, policies, value functions).

1. The Policy Gradient Idea

In value-based RL (Q-learning, SARSA), we learn a value function and derive a policy from it. Policy gradient methods take a different approach: parameterize the policy directly and optimize it by gradient ascent on expected return.

We have a policy \(\pi_\theta(a \mid s)\) — a probability distribution over actions given a state, parameterized by \(\theta\) (typically the weights of a neural network). The objective is to maximize expected cumulative reward:

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

where \(\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \ldots)\) is a trajectory sampled by running policy \(\pi_\theta\) in the environment, and \(\gamma \in [0, 1]\) is the discount factor.

The idea is simple: compute \(\nabla_\theta J(\theta)\) and take gradient ascent steps.

2. The Policy Gradient Theorem

The key theoretical result (Sutton et al., 2000) gives us a tractable expression for the gradient:

\[\nabla_\theta J(\theta) = \mathbb{E}_{s_t, a_t \sim \pi_\theta} \left[ \nabla_\theta \log \pi_\theta(a_t \mid s_t) \; \Psi_t \right]\]

where \(\Psi_t\) is some measure of how good the action was. Different choices of \(\Psi_t\) give different algorithms:

Choice of \(\Psi_t\) Name Properties
\(R(\tau) = \sum_{t'=0}^{T} \gamma^{t'} r_{t'}\) REINFORCE High variance, no bias
\(\hat{R}_t = \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'}\) Reward-to-go Lower variance (causality)
\(\hat{A}_t = Q(s_t, a_t) - V(s_t)\) Advantage Lowest variance (baselined)

2.1 The score function \(\nabla_\theta \log \pi_\theta(a_t \mid s_t)\)

This term is called the score function (or log-likelihood gradient). It points in the direction in parameter space that increases the probability of action \(a_t\) in state \(s_t\).

2.2 The advantage function \(A(s_t, a_t)\)

The advantage function is the most common choice for \(\Psi_t\):

\[A(s_t, a_t) = Q(s_t, a_t) - V(s_t)\]
  • \(Q(s_t, a_t)\): expected return from taking action \(a_t\) in state \(s_t\), then following \(\pi_\theta\)
  • \(V(s_t)\): expected return from state \(s_t\) under \(\pi_\theta\) (averaged over all actions)

Intuition: The advantage measures how much better this action was compared to the average action at this state.

  • \(A > 0\): the action was better than average → increase its probability
  • \(A < 0\): the action was worse than average → decrease its probability

The gradient thus reinforces good actions and suppresses bad ones, weighted by how much better or worse they are.

3. The REINFORCE Algorithm

The simplest policy gradient algorithm (Williams, 1992):

  1. Collect a trajectory \(\tau\) by running \(\pi_\theta\)
  2. For each time step \(t\), compute the reward-to-go \(\hat{R}_t\)
  3. Update: \(\theta \leftarrow \theta + \alpha \sum_{t} \nabla_\theta \log \pi_\theta(a_t \mid s_t) \, \hat{R}_t\)

This works, but has high variance — the gradient estimates are noisy because \(\hat{R}_t\) depends on all future randomness (stochastic transitions, stochastic actions).

3.1 Variance reduction via baselines

Subtracting a baseline \(b(s_t)\) from the return does not change the expected gradient (it's still unbiased), but can dramatically reduce variance:

\[\nabla_\theta J(\theta) = \mathbb{E}\left[ \nabla_\theta \log \pi_\theta(a_t \mid s_t) \left( \hat{R}_t - b(s_t) \right) \right]\]

The optimal baseline is approximately \(V(s_t)\) — the expected return from state \(s_t\). This gives us the advantage:

\[\hat{A}_t = \hat{R}_t - V(s_t)\]

In practice, we learn \(V_\phi(s)\) (a neural network with parameters \(\phi\)) alongside the policy. This is the actor-critic architecture:

  • Actor: the policy \(\pi_\theta\) (chooses actions)
  • Critic: the value function \(V_\phi\) (evaluates states, provides the baseline)

4. The Vanilla Policy Gradient Algorithm

Putting it together:

For \(k = 0, 1, 2, \ldots\) do:

  1. Collect trajectories \(\mathcal{D}_k = \{\tau_i\}\) by running \(\pi_{\theta_k}\)
  2. Compute rewards-to-go \(\hat{R}_t\) for each time step
  3. Compute advantage estimates \(\hat{A}_t = \hat{R}_t - V_{\phi_k}(s_t)\)
  4. Update policy via gradient ascent:
\[\theta_{k+1} = \theta_k + \alpha \, \hat{g}_k\]

where:

\[\hat{g}_k = \frac{1}{|\mathcal{D}_k|} \sum_{\tau \in \mathcal{D}_k} \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) \Big|_{\theta_k} \hat{A}_t\]
  1. Update value function by regression:
\[\phi_{k+1} = \arg\min_\phi \frac{1}{|\mathcal{D}_k| T} \sum_{\tau \in \mathcal{D}_k} \sum_{t=0}^{T} \left( V_\phi(s_t) - \hat{R}_t \right)^2\]

5. The Hidden Danger: Destructive Updates

The algorithm above looks clean. In practice, it is extremely fragile.

Nothing in vanilla policy gradient prevents catastrophic updates. If:

  • Advantage estimates are noisy (they always are)
  • The learning rate \(\alpha\) is a bit too large (it will be)
  • The batch is slightly unrepresentative (it always is)

Then \(\pi_{\theta_{k+1}}\) can be radically different from \(\pi_{\theta_k}\).

This breaks two things simultaneously:

  1. Data mismatch: The data in \(\mathcal{D}_k\) was sampled from \(\pi_{\theta_k}\), not \(\pi_{\theta_{k+1}}\). If the new policy is very different, the gradient estimate computed from old data is meaningless.

  2. Value function invalidation: The critic \(V_{\phi_k}\) was trained to predict returns under \(\pi_{\theta_k}\). Under a very different policy, these predictions are wrong, so the advantage estimates become garbage.

The result: training instability, performance collapse, oscillation.

Unlike supervised learning — where a bad gradient step just temporarily increases the loss — a bad policy gradient step changes the data distribution itself (because the policy determines which states the agent visits). This creates a vicious cycle: bad update → bad data → worse update.

5.1 Why step size in parameter space is the wrong metric

The fundamental issue is that the learning rate \(\alpha\) controls the step size in parameter space (\(\|\Delta\theta\|\)), but what matters is the step size in policy space — how much the distribution \(\pi_\theta(\cdot \mid s)\) changes.

Two parameter vectors close in Euclidean distance can produce wildly different policies, and vice versa. The geometry of parameter space doesn't match the geometry of policy space.

This is the core problem that TRPO and PPO solve — by constraining updates in policy space rather than parameter space.

6. Generalized Advantage Estimation (GAE)

Before moving to TRPO and PPO, there's one more important ingredient: how to compute good advantage estimates.

The simple advantage \(\hat{A}_t = \hat{R}_t - V(s_t)\) uses the full Monte Carlo return, which is unbiased but high-variance. The 1-step TD advantage \(\hat{A}_t = r_t + \gamma V(s_{t+1}) - V(s_t)\) is low-variance but biased (because \(V\) is approximate).

GAE (Schulman et al., 2016) interpolates between these extremes using a parameter \(\lambda \in [0, 1]\):

\[\hat{A}_t^{\text{GAE}(\gamma, \lambda)} = \sum_{l=0}^{T-t} (\gamma \lambda)^l \, \delta_{t+l}\]

where \(\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)\) is the 1-step TD error.

\(\lambda\) Behavior Bias Variance
\(\lambda = 0\) 1-step TD advantage High bias Low variance
\(\lambda = 1\) Monte Carlo advantage No bias High variance
\(\lambda \approx 0.95\) Typical sweet spot Low bias Moderate variance

GAE is used in virtually all modern policy gradient implementations (TRPO, PPO, GRPO).

7. The Importance Sampling Connection

One more concept that bridges vanilla PG to TRPO/PPO: importance sampling.

After collecting data from \(\pi_{\theta_k}\), we want to evaluate the gradient for a different policy \(\pi_\theta\). The importance sampling ratio:

\[r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_k}(a_t \mid s_t)}\]

reweights the old data to account for the policy change. The surrogate objective becomes:

\[L(\theta) = \mathbb{E}_{s_t, a_t \sim \pi_{\theta_k}} \left[ r_t(\theta) \, \hat{A}_t \right]\]

This is exact when \(\theta = \theta_k\) (where \(r_t = 1\)), and a good approximation when \(\theta\) is close to \(\theta_k\).

The problem: nothing prevents \(r_t(\theta)\) from becoming very large (or very small), which makes the surrogate objective a poor approximation and leads to destructive updates. This is precisely what TRPO constrains (via KL divergence) and PPO clips (via the \([1-\epsilon, 1+\epsilon]\) bound).

8. Notation Reference

Symbol Meaning
\(\pi_\theta(a \mid s)\) Policy: probability of action \(a\) in state \(s\), parameterized by \(\theta\)
\(J(\theta)\) Expected cumulative reward under policy \(\pi_\theta\)
\(V(s)\) State value function: expected return from state \(s\)
\(Q(s, a)\) Action value function: expected return from taking \(a\) in \(s\)
\(A(s, a) = Q(s,a) - V(s)\) Advantage function
\(\hat{R}_t\) Reward-to-go: \(\sum_{t'=t}^{T} \gamma^{t'-t} r_{t'}\)
\(\hat{A}_t\) Advantage estimate (e.g., via GAE)
\(\gamma\) Discount factor
\(\alpha\) Learning rate
\(r_t(\theta)\) Importance sampling ratio: \(\pi_\theta(a_t \mid s_t) / \pi_{\theta_k}(a_t \mid s_t)\)
\(\tau\) Trajectory: \((s_0, a_0, r_0, s_1, \ldots)\)
\(\mathcal{D}_k\) Batch of trajectories collected at iteration \(k\)

9. Summary

What policy gradient methods get right

  1. Direct optimization of the quantity we care about (expected return)
  2. Natural handling of continuous actions (no need for \(\max_a Q(s,a)\))
  3. Stochastic policies — exploration is built into the policy itself
  4. Theoretical foundation — the policy gradient theorem provides an unbiased gradient estimator

What goes wrong

  1. High variance — gradient estimates are noisy, requiring large batches or variance reduction (baselines, GAE)
  2. Destructive updates — no mechanism to prevent catastrophically large policy changes
  3. Sample inefficiency — on-policy methods discard data after each update
  4. Step size sensitivity — parameter-space learning rates don't correspond to policy-space changes

The path forward

The destructive update problem is the most critical. Two solutions emerged:

  • TRPO: Constrain updates via KL divergence (trust region). Theoretically principled but computationally expensive (conjugate gradient + line search).
  • PPO: Clip the importance sampling ratio to softly enforce a trust region. Simple, first-order, and nearly as stable as TRPO.

References

  • [Williams, 1992] — Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. Machine Learning, 8(3-4), 229-256. (REINFORCE)
  • [Sutton et al., 2000] — Policy Gradient Methods for Reinforcement Learning with Function Approximation. NeurIPS 2000. (Policy gradient theorem)
  • [Schulman et al., 2016] — High-Dimensional Continuous Control Using Generalized Advantage Estimation. ICLR 2016. arXiv:1506.02438 (GAE)
  • [OpenAI Spinning Up] — spinningup.openai.com/en/latest/spinningup/rl_intro3.html (Policy optimization introduction)

Next: Trust Region Policy Optimization (TRPO) — how to constrain policy updates for guaranteed improvement.

Supplement: Parameterizing the Policy — how \(\pi_\theta(a \mid s)\) is actually implemented, from Gaussians to normalizing flows