Policy Gradients in reinforcement learning

Module 16 · Advanced AI · 23 concepts

In depth / 15 lines

Policy Gradients, in more detail

Policy gradient methods learn the policy directly instead of learning values.

The policy is a network that outputs a probability for each action.

The agent samples from that distribution rather than always taking the best.

Training pushes up the probability of actions that led to high reward.

It pushes down the probability of actions that led to low reward.

The gradient is estimated from the sampled trajectories themselves.

That is the REINFORCE algorithm in its simplest form.

It handles continuous action spaces naturally, which value methods struggle with.

A robot arm angle has infinitely many values and cannot be enumerated in a table.

It also produces stochastic policies, which are the right answer in some games.

The weakness is variance, since a whole episode's reward is credited to every action in it.

The estimate is unbiased and extremely noisy.

A baseline subtracts the average expected reward, which cuts variance without adding bias.

The advantage measures how much better an action was than average, and works better still.

Actor-critic methods learn a value function alongside the policy to compute that advantage.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Policy Gradients

Optimising the policy directly

J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}}\big[ G(\tau) \big]

$\tau$ : a trajectory, one whole run through the environment.
$G(\tau)$ : its total return.

The objective is the expected return, and the parameters being optimised are the policy's own rather than a value function's.

\nabla_{\theta} J = \mathbb{E}\Big[ \nabla_{\theta} \log \pi_{\theta}(a \mid s) \; G \Big]

That is the policy gradient theorem, and it is remarkable because the environment's dynamics do not appear in it.

The log-derivative trick

\nabla_{\theta} p_{\theta}(\tau) = p_{\theta}(\tau) \, \nabla_{\theta} \log p_{\theta}(\tau)

The gradient of an expectation over a distribution that itself depends on $\theta$ is not an expectation of a gradient, which is the obstacle.

\log p_{\theta}(\tau) = \log p(s_0) + \sum_{t} \Big[ \log \pi_{\theta}(a_t \mid s_t) + \log P(s_{t+1} \mid s_t, a_t) \Big]

Taking the logarithm turns the product over timesteps into a sum, and the transition terms do not depend on $\theta$ , so they differentiate to zero.

Only the policy's own log probabilities survive, which is why the method works without a model of the environment.

What the update actually does

\theta \leftarrow \theta + \alpha \, \nabla_{\theta} \log \pi_{\theta}(a \mid s) \, G

Increase the log probability of actions that were followed by a high return, decrease it for those that were not, in proportion to the return.

If all returns are positive, every action's probability is pushed up, and learning happens only through the differences in magnitude, which is slow and noisy.

The baseline, which is not optional

\nabla_{\theta} J = \mathbb{E}\Big[ \nabla_{\theta} \log \pi_{\theta}(a \mid s) \, \big( G - b(s) \big) \Big]

Subtracting any function of the state leaves the gradient unbiased, because the expected score function is zero.

b(s) = V(s) \implies A(s,a) = Q(s,a) - V(s)

Using the state value as the baseline turns the weight into the advantage: how much better this action was than the average from that state.

That single change cuts the variance by orders of magnitude, and a learned $V$ is the critic in an actor-critic method.

One update, in order

Run the current policy to collect a batch of trajectories.
Compute the return from each timestep onward, not the whole episode's return.
Subtract the baseline to get advantages, and normalise them across the batch.
Take the gradient of the log probabilities weighted by those advantages.
Step, and discard the data, because it was collected under the old policy.

Step five is the expensive part. Policy gradients are on-policy, so every batch is used once, which is why they need far more environment interaction than Q-learning.

\mathcal{L} \mathrel{-}= \beta \, H\big(\pi_{\theta}(\cdot \mid s)\big)

An entropy bonus is usually added, because a policy that becomes deterministic early stops exploring and cannot recover.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

What do policy gradient methods learn directly?

Advanced AI

23 concepts in this module, then a graded test with every answer explained.

Study this module →