Policy gradients

A tennis coach does not hand a player a table of how good every possible shot would be. The player tries a shot, sees how the point went, and becomes a little more likely to repeat what worked. Over thousands of points the habits improve, without ever writing the values down.

The methods in the last two lessons, value iteration and Q-learning, learned values and then picked the action with the best value. This lesson’s methods skip that step and adjust the policy itself.

Learning the policy directly

Give the policy adjustable numbers, called parameters $\theta$: for example, the weights of a neural network that outputs a probability for each action. Write the expected return of the policy as $J(\theta)$. To improve the policy, climb $J$: change $\theta$ a little in the direction of its gradient. This is a policy gradientPolicy gradientA family of methods that learn a policy directly, adjusting its parameters in the direction that increases expected return, estimated from sampled episodes.Open in glossary method.

Learning a policy directly has advantages. The policy can be random on purpose, which sometimes is the best you can do. It handles actions that are continuous, like how hard to turn a steering wheel, where picking the best of all actions is not simple. And a small change to $\theta$ changes behavior a little, while a small change in values can flip which action is best.

A corridor where the best policy is random

Here is a problem from Sutton and Barto’s textbook (Example 13.1). A corridor has three cells and a goal, and every step costs 1. In the first and third cells, choosing right moves right. In the second cell the actions are switched: right moves left. The catch is that the agent cannot tell the cells apart. It sees the same thing everywhere, so its policy is a single number: $p$, the chance of choosing right.

Always choosing right bounces between the first two cells forever; always choosing left walks into the wall forever. The best policy is a mix. Solving the corridor’s Bellman equations gives the expected return as a function of $p$:

which is highest at $p = 2 - \sqrt{2} \approx 0.59$, worth about $-11.66$. An agent that learns values and acts epsilon-greedily (with epsilon 0.1) is stuck choosing between going right 95% of the time, worth about $-44$, and going left 95% of the time, worth about $-82$. A policy gradient method can learn the right mix.

REINFORCE

To climb $J$ we need its gradient, but $J$ is an average over random episodes. How do you take the gradient of an average over outcomes you can only sample? The log-derivative trickLog-derivative trickThe identity that the gradient of p(x) equals p(x) times the gradient of ln p(x). It turns the gradient of an expected reward into an expectation that can be estimated from samples.Open in glossary turns it into an average you can estimate:

Here $G_t$ is the return from step $t$ onward. Read it as a recipe: after an episode, for each action taken, push up its log-probability in proportion to the return that followed it. Actions followed by high returns become more likely; actions followed by low returns, less. Ronald Williams named this family of algorithms REINFORCEREINFORCEThe basic policy gradient algorithm (Williams, 1992): after an episode, make each action taken more likely in proportion to the return that followed it, minus a baseline.Open in glossary in 1992.

In the corridor the policy is $p = \sigma(\theta)$, so $\nabla_\theta \ln \pi(\text{right}) = 1 - p$ and $\nabla_\theta \ln \pi(\text{left}) = -p$. Every step costs 1, so every return is negative, and short episodes have less negative returns than long ones.

import numpy as np

# The short corridor: cells 0, 1, 2, then the goal. Every step costs 1.
# In cell 1 the actions are switched. The agent cannot tell the cells apart,
# so its whole policy is one number: p, the chance of choosing "right".
def run_episode(p, rng, max_steps=1000):
    cell, actions = 0, []
    while cell < 3 and len(actions) < max_steps:
        right = rng.random() < p
        actions.append(right)
        forward = (not right) if cell == 1 else right
        cell = cell + 1 if forward else max(0, cell - 1)
    return actions

def J(p):  # exact expected return from the start
    return 2 * (p - 2) / (p * (1 - p))

rng = np.random.default_rng(0)
theta, alpha, baseline = np.log(0.1 / 0.9), 0.005, 0.0  # start at p = 0.1
for episode in range(1, 301):
    p = 1 / (1 + np.exp(-theta))
    actions = run_episode(p, rng)
    T = len(actions)
    grad = 0.0
    for t, right in enumerate(actions):
        G = -(T - t)                       # return from step t on
        score = (1 - p) if right else -p   # d/dtheta of ln pi(action)
        grad += (G - baseline) * score
    for t in range(T):                     # running average of returns
        baseline += 0.05 * (-(T - t) - baseline)
    theta += alpha * grad
    if episode % 50 == 0:
        p = 1 / (1 + np.exp(-theta))
        print(f"episode {episode}: p(right) = {p:.2f}, expected return {J(p):.1f}")

print(f"best possible: p = {2 - np.sqrt(2):.3f}, expected return {J(2 - np.sqrt(2)):.2f}")

episode 50: p(right) = 0.56, expected return -11.7
episode 100: p(right) = 0.51, expected return -11.9
episode 150: p(right) = 0.58, expected return -11.7
episode 200: p(right) = 0.60, expected return -11.7
episode 250: p(right) = 0.62, expected return -11.7
episode 300: p(right) = 0.52, expected return -11.9
best possible: p = 0.586, expected return -11.66

What the training run shows

It finds the mix. From either side, the policy climbs to around $p = 0.59$ and stays near it, jittering. Nothing told it that randomness was the answer; following the gradient of expected return led there.

Each step follows a noisy estimate. One episode is one sample, and episodes vary wildly in length. With a large learning rate, one long, unlucky episode can fling the policy from one side of the range to the other. Smaller steps average out the noise over many episodes, at the cost of slower learning.

It keeps jittering at the top. Even at the best policy, individual estimates are large and point in both directions; they only average to zero. With a fixed step size the policy never settles exactly.

Baselines: less noise, same average

All returns in the corridor are negative, so plain REINFORCE makes every action it took less likely, and relies on the good episodes being pushed down less than the bad ones. That works on average, but each estimate is noisy.

A baselineBaseline (policy gradient)A number subtracted from the return in a policy gradient estimate, often an estimate of the average return. Any baseline that ignores the action taken leaves the estimate unbiased; a good one makes it much less noisy.Open in glossary fixes this by comparing each return to a typical return. Subtract a number $b$ from every return:

Now an action is pushed up when its episode went better than usual and down when it went worse. Remarkably, this does not change the average of the estimate at all, as long as $b$ does not depend on the action taken. The reason is one line: averaged over the policy’s own actions, $\sum_a \pi(a)\, \nabla \ln \pi(a) = \sum_a \nabla \pi(a) = \nabla 1 = 0$. So subtracting $b$ subtracts something whose average is zero. What it can change is the spread.

import numpy as np

# One choice between two actions: action 0 pays 1, action 1 pays 3.
# The policy picks action 1 with probability p = sigmoid(theta).
rewards = np.array([1.0, 3.0])
theta = 0.4
p = 1 / (1 + np.exp(-theta))
probs = np.array([1 - p, p])

# The expected reward and its true gradient: dJ/dtheta = (3 - 1) p (1 - p).
J = probs @ rewards
true_grad = (rewards[1] - rewards[0]) * p * (1 - p)

# REINFORCE estimates the gradient from one sampled action A as
# (reward - b) * d/dtheta ln pi(A). For this policy the score is
# 1 - p for action 1 and -p for action 0.
score = np.array([-p, 1 - p])
print(f"p = {p:.3f}, true gradient = {true_grad:.4f}")
print("average score over actions is zero:", bool(np.isclose(probs @ score, 0)))
for b in [0.0, 2.0, J, 6.0]:
    estimates = (rewards - b) * score           # one estimate per possible action
    mean = probs @ estimates                    # exact average over actions
    spread = np.sqrt(probs @ (estimates - mean) ** 2)
    print(f"baseline {b:4.2f}: average {mean:.4f}, spread {spread:.4f}")

p = 0.599, true gradient = 0.4805
average score over actions is zero: True
baseline 0.00: average 0.4805, spread 0.8836
baseline 2.00: average 0.4805, spread 0.0967
baseline 2.20: average 0.4805, spread 0.1935
baseline 6.00: average 0.4805, spread 2.0574

In the corridor, a baseline near the expected return cuts the spread of the estimates by about a third, and a baseline far from it can make the spread much wider, while the average stays on the true gradient every time. The code shows the same thing exactly, with no sampling: every baseline gives the same average, and the spread changes. The baseline with the least spread is close to the average return but not always equal to it. In the code, 2.0 spreads less than the average reward, 2.2, and with only two actions the spread reaches zero at $b = 3 - 2p \approx 1.80$, where both possible estimates equal the true gradient.

From REINFORCE to modern methods

In practice the baseline is usually learned: a second network estimates the value of each state, and each action is judged by how much better its return was than that estimate. A method with a learned value alongside the policy is called an actor-critic. Proximal policy optimizationProximal policy optimizationA policy gradient algorithm that limits how far each update can move the policy, by clipping the change in action probabilities. Used for RLHF in InstructGPT.Open in glossary (PPO; Schulman and colleagues, 2017) adds one more safeguard you saw the need for: it clips its objective so that an update gains nothing from moving an action’s probability more than a set fraction (20% in the paper’s example) away from where it started.

Policy gradients appear throughout modern AI. DeepMind’s AlphaGo improved its policy network with policy gradient reinforcement learning from games against itself. And PPO is the algorithm OpenAI used to fine-tune InstructGPT on human preferences, the subject of the next lesson.