ML10.4 · LESSON 4 OF 6

Policy gradients

Distinguish sequential decision-making, temporal prediction, recommendation and causal estimation.

THE IDEA

Learn the policy directly

Q-learning learns values and derives behaviour from them. A policy-gradient method learns the behaviour itself. The policy is a probability distribution over actions with adjustable numbers, here one preference per action turned into probabilities with softmax. After each outcome you nudge the preferences so that actions that led to good results become more likely.

REINFORCE on a bandit

pi = softmax(theta)                     # action probabilities
action ~ pi,  reward R observed
theta = theta + lr * (R - baseline) * (onehot(action) - pi)

The term (onehot(action) − pi) is the gradient of the log-probability of the chosen action. Multiplying by R − baseline pushes the chosen action up when it did better than expected and down when it did worse.

Two actions pay 1 with probability 0.7 (action 0) and 0.3 (action 1). The policy is a softmax over two numbers theta.

action 0

theta = 0

pi = 0.5

action 1

theta = 0

pi = 0.5

learning rate0.5

Press Sample one step to see the update arithmetic.

For the current policy (pi_0 = 0.5), the single-sample update of theta_0 is random:

no baseline: mean 0.1, variance 0.115

baseline = expected reward 0.5: mean 0.1, variance 0.0525

Same average push, much less noise per sample when the baseline is subtracted.

Press Sample one step to draw an action and a reward, see the arithmetic, and move the policy. Single updates are noisy, and sometimes they push the wrong way.

Why a baseline helps

Without a baseline, every reward of 1 pushes the chosen action up, even if every action usually pays. Subtracting a baseline close to the average reward turns the signal into "better or worse than usual". The expected direction of the update stays the same, but the random swings shrink, so learning is steadier.

Predict first. The baseline is 0.6. The chosen action gave reward 0. What does the update do to that action's probability?

seed 1seed 2seed 3seed 4seed 5seed 6seed 7seed 8mean of 8

learning rate0.2

20 seeded runs of 300 steps

mean final pi(best): 0.979

runs ending on the worse action: 0 of 20

spread (std of final pi): 0.006

Many runs with and without a baseline at the same learning rate. Look at both the average and the spread between runs.

Chapter context

These topics change the structure of the learning problem. Their first lesson is to identify what a prediction or action is allowed to depend on and how feedback is produced. Each section is a bridge to a larger specialist subject.

TRY IT YOURSELF

Part 1. Do it hands-on

Each setting runs many seeded training runs of the same two-action problem. Tune the learning rate and the baseline to satisfy the three goals.

seed 1seed 2seed 3seed 4seed 5seed 6seed 7seed 8mean of 8

learning rate0

20 seeded runs of 300 steps

mean final pi(best): 0.5

runs ending on the worse action: 0 of 20

spread (std of final pi): 0

To do: Find a learning rate where, after 300 steps, the policy picks the better action with average probability above 0.9
To do: Show the danger of a big step: with no baseline, find a learning rate where at least 1 in 20 runs ends up preferring the worse action
To do: Turn the baseline on at a learning rate of 2.0 or more, where it removes the wrong-action runs and gives a smaller seed-to-seed spread than the same rate without it

Part 2. Do it in code

Write softmax(theta) that is stable for huge values (subtract the maximum first) and returns a numpy array that sums to 1.
Write reinforce_step(theta, action, reward, baseline, lr): return a new theta array with theta_j + lr * (reward - baseline) * (onehot(action)_j - pi_j), where pi = softmax(theta). Do not modify the input.
Write train_bandit(probs, steps, lr, seed, use_baseline): theta starts at zeros, rng = random.Random(seed), b = 0.0. Each step t = 1..steps: pi = softmax(theta); u = rng.random(); the action is the first index whose cumulative probability exceeds u (the last index if rounding leaves none); reward = 1 if rng.random() < probs[action] else 0; update with baseline b (or 0 when use_baseline is False); then b += (reward - b) / t. Return the final softmax(theta).

import random
import numpy as np

def softmax(theta):
    pass

def reinforce_step(theta, action, reward, baseline, lr):
    pass

def train_bandit(probs, steps, lr, seed, use_baseline):
    pass

print(softmax(np.array([0.0, 0.0])))

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Must a stochastic policy assign probability one to its current favorite action?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

G_t=\sum_{k\ge0}\gamma^k r_{t+k+1},\quad Q(s,a)\leftarrow Q(s,a)+\alpha[r+\gamma\max_{a'}Q(s',a')-Q(s,a)].

For a terminal transition the continuation term is zero. A policy-gradient estimator uses a term proportional to $\nabla\log\pi_\theta(a\mid s)$ times a return or advantage estimate.

Worked example

For a nonterminal transition, current Q=2, reward=1, discount=0.9 and best next Q=4 give target 1+0.9×4=4.6. With alpha=0.5, updated Q is 2+0.5×(4.6−2)=3.3. Reusing the same target for a terminal transition would be incorrect.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import isclose
q, reward, discount, next_q, rate = 2, 1, .9, 4, .5
target = reward + discount*next_q
assert isclose(target, 4.6)
assert isclose(q+rate*(target-q), 3.3)
terminal_target = reward
assert q+rate*(terminal_target-q) == 1.5

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Prediction is not intervention. Randomly splitting overlapping time windows can leak future information. An environment reward is a modeling choice and may omit important real-world consequences.