Policy gradients

The policy gradient theorem#

Parameterise $\pi_\theta(a \mid s)$ and maximise $J(\theta) = \mathbb{E}{\tau \sim \pi\theta}[G(\tau)]$. The difficulty is that the trajectory distribution depends on $\theta$, so the gradient cannot pass through the expectation directly.

The log-derivative trick, using $\nabla_\theta p_\theta(x) = p_\theta(x) \nabla_\theta \log p_\theta(x)$:

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

Because $\log p_\theta(\tau) = \log p(s_0) + \sum_t \left[\log \pi_\theta(a_t \mid s_t) + \log P(s_{t+1} \mid s_t, a_t)\right]$ and the environment terms carry no $\theta$, they vanish under the gradient. The transition dynamics disappear entirely. This is why policy gradients are model-free without needing any approximation.

Sutton et al. (2000) give the general form: $$ \nabla_\theta J(\theta) = \mathbb{E}{s \sim d^{\pi}, a \sim \pi\theta}\left[ \nabla_\theta \log \pi_\theta(a \mid s) \, Q^{\pi_\theta}(s,a) \right] $$ with $d^\pi$ the discounted state visitation distribution. Symbols: $G_t$ the return from $t$, $Q^{\pi}$ the action-value under the current policy, $\tau$ a trajectory.

Baselines are free#

For any function $b(s)$ that does not depend on $a$: $$ \mathbb{E}{a \sim \pi\theta}\left[ \nabla_\theta \log \pi_\theta(a \mid s) \, b(s) \right] = b(s) \nabla_\theta \sum_a \pi_\theta(a \mid s) = b(s) \nabla_\theta 1 = 0 $$

So subtracting $b(s)$ from the return leaves the gradient unbiased while changing its variance. The variance-minimising baseline is $b^*(s) = \frac{\mathbb{E}[(\nabla \log \pi)^2 G]}{\mathbb{E}[(\nabla \log \pi)^2]}$, a gradient-magnitude-weighted average of returns. In practice $b(s) = V^\pi(s)$ is used, which is close enough and gives $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$, the advantage. That substitution is the step into actor-critic.

Note that the per-episode standardisation used in the code is not a state-dependent baseline — it is a batch normalisation of returns. It reduces scale sensitivity and it does introduce a small bias. It is nonetheless the near-universal practical choice.

Variance, quantified#

The REINFORCE estimator's variance grows roughly as $O(T^2)$ in the horizon under mild assumptions, because $G_t$ sums $T-t$ random rewards and appears in $T$ terms. Three standard reductions:

1. Reward-to-go. Replace $G(\tau)$ with $G_t$, dropping rewards that precede the action. Unbiased by causality.
2. Baseline subtraction. Unbiased, as shown above.
3. Bootstrapping. Replace $G_t$ with $r_t + \gamma V(s_{t+1})$. This does introduce bias and cuts variance sharply. GAE (Schulman et al., 2016) parameterises the whole spectrum with $\lambda$.

The on-policy constraint#

The expectation is over $\tau \sim \pi_\theta$ for the current $\theta$. After one gradient step the data is off-policy and the estimator is biased. Correcting with importance sampling gives $$ \nabla_\theta J = \mathbb{E}{\tau \sim \pi{\theta_{\text{old}}}}\left[ \frac{\pi_\theta(a\mid s)}{\pi_{\theta_{\text{old}}}(a \mid s)} \nabla_\theta \log \pi_\theta(a\mid s) A(s,a) \right] $$ but the importance ratio has unbounded variance over long trajectories. Constraining or clipping that ratio is precisely what TRPO and PPO do, and it is what makes limited data reuse safe.

Natural policy gradients#

The plain gradient is steepest ascent in Euclidean parameter space, which is the wrong geometry — equal parameter changes produce wildly unequal changes in the policy distribution. The natural gradient (Amari, 1998; Kakade, 2002) preconditions by the Fisher information matrix: $$ \tilde{\nabla}\theta J = F(\theta)^{-1} \nabla\theta J, \qquad F(\theta) = \mathbb{E}\left[ \nabla_\theta \log \pi_\theta \, \nabla_\theta \log \pi_\theta^\top \right] $$ This is steepest ascent under the KL divergence between policies, which is invariant to reparameterisation. Computing $F^{-1}$ is infeasible for large networks; TRPO approximates it with conjugate gradients, and PPO abandons it for a much cheaper clipping heuristic that works nearly as well.

Continuous actions#

For $\mathcal{A} \subseteq \mathbb{R}^d$, output a distribution — typically a diagonal Gaussian $\pi_\theta(a\mid s) = \mathcal{N}(\mu_\theta(s), \operatorname{diag}(\sigma_\theta^2))$. The log-probability is differentiable and the same estimator applies unchanged. A state-independent learned $\log\sigma$ is standard and usually more stable than a state-dependent one. Bounded action spaces need care: naive clipping biases the gradient, and a tanh-squashed Gaussian with the corresponding log-det-Jacobian correction (as in SAC) is the cleaner treatment.

References#

• Williams (1992), Simple statistical gradient-following algorithms for connectionist reinforcement learning — REINFORCE.
• Sutton, McAllester, Singh and Mansour (2000), Policy Gradient Methods for Reinforcement Learning with Function Approximation.
• Kakade (2002), A Natural Policy Gradient.
• Schulman et al. (2016), High-Dimensional Continuous Control Using Generalized Advantage Estimation — arxiv.org/abs/1506.02438.
• Greensmith, Bartlett and Baxter (2004), Variance Reduction Techniques for Gradient Estimates in Reinforcement Learning.

What to learn next#

• Actor-critic methods — learning the baseline instead of guessing it.
• PPO — how to reuse data without the policy exploding.
• Gradient descent — the optimiser underneath all of this.