Policy gradients and actor-critic, PPO

The big idea

Q-learning learns values and derives the policy from them: take the action with the highest Q. That works when you can list the actions and take a max. It breaks down for a robot arm choosing a continuous torque, and it gives only deterministic policies, where sometimes randomness is the best strategy (think of rock-paper-scissors).

Policy gradientReinforcement learning methods that represent the policy directly as a parameterized distribution over actions and adjust its parameters by gradient ascent on expected return. methods skip the detour. The policy is a network that outputs a probability for each action (or the mean and spread of a continuous action), and training nudges its weights to make actions that led to high return more likely. Think of a coach watching recordings of games. Plays that ended well get practised more, plays that ended badly get practised less, with no attempt to assign every position a precise score. Two refinements make this practical: compare each outcome with what was expected (a baseline, which becomes a learned critic), and never change the policy too much in one go (PPO's clip).

How it works

Step 1: REINFORCE

Write the policy as ${\pi }_{\theta }(a\mid s)$. The policy gradient theorem gives the direction that increases expected return $J(\theta )$:

REINFORCEThe basic Monte Carlo policy gradient algorithm (Williams, 1992): run an episode, then raise the log-probability of each action in proportion to the return that followed it. estimates this from whole episodes: play one, compute the return $G_t$ after each step, and move $\theta$ along $\nabla \log \pi (a_t\mid s_t)G_t$. $\nabla \log \pi$ points in the direction that makes the chosen action more likely; $G_t$ scales the step. Good outcomes pull their actions up, bad ones push them down. No model of the environment and no max over actions are needed, so continuous actions are fine.

The weakness is noise. Returns vary a lot from episode to episode, often for reasons unrelated to the action (a lucky slip, a random opponent move). If every return is positive, every action taken gets reinforced, just some more than others, and it takes many episodes for the average to sort out which actions were actually good.

Step 2: Baselines and advantages

Subtract a baselineA reference value subtracted from the return in a policy gradient update; any baseline that does not depend on the action leaves the gradient unbiased while reducing its variance. $b(s_t)$ from the return. As long as the baseline does not depend on the action, the expected gradient is unchanged, but its variance drops. The natural baseline is the state's value $V(s_t)$, which turns the weight into the advantageHow much better an action turned out than the state's expected value: A(s, a) = Q(s, a) - V(s), estimated in practice from returns or TD errors.:

Now the question for each action is not "was the outcome good?" but "was it better than usual from here?". A return of 7 from a state that usually yields 9 pushes the action down, not up.

Step 3: Actor-critic

Actor-criticA reinforcement learning architecture with two learned parts: an actor (the policy) that chooses actions and a critic (a value function) whose estimates supply the advantage used to update the actor. methods learn the baseline as a second model, the critic, and often use its one-step TD error as the advantage instead of waiting for the full return:

That gives lower variance (one reward instead of a whole noisy future), updates every step instead of at the end of an episode, and some bias, because the critic's estimates are imperfect. Practical methods mix the two, using several steps of real rewards before trusting the critic.

Step 4: PPO

Policy gradients have a sharp edge: one step too large and the policy can collapse, after which the data it collects is useless for recovering. You would also like to reuse each batch of experience for several gradient steps, but the plain objective is only valid near the policy that collected the data. The PPO paper reports that optimizing it for several steps on the same data often leads to destructively large policy updates.

PPOProximal policy optimization: an actor-critic method that reuses each batch for several epochs while clipping the probability ratio between new and old policy, so no single update moves the policy far. (Schulman and colleagues, 2017) keeps the new policy close to the one that collected the data. With the probability ratio $r_t(\theta )={\pi }_{\theta }(a_t\mid s_t)/{\pi }_{\text{old}}(a_t\mid s_t)$, it maximizes

with $ϵ$ around 0.2.

For a good action (positive advantage), the objective grows as the action becomes more likely, but flattens once the ratio passes $1+ϵ$: there is no gain from pushing further, so the gradient is zero. For a bad action it stops at $1-ϵ$. The min makes it pessimistic: when the ratio moves the wrong way, the unclipped term applies in full. That gives PPO much of the safety of the older trust-region method (TRPO) with plain first-order optimization, and lets it run several epochs of minibatch updates on each batch. In the paper's comparison over seven simulated robot tasks, clipping with $ϵ=0.2$ scored best (0.82 on their normalized scale), while the same objective with no clipping or penalty scored -0.39, worse than a random policy on average.

A full PPO agent is an actor-critic: a policy network trained with $L^{\text{CLIP}}$, a value network trained on returns, often with an entropy bonus to keep exploring.

Going deeper

Code: REINFORCE, a baseline, actor-critic and the clip

import numpy as np

# The corridor from ML-13.01: cells 0..6, start in 3, exits worth +1 (cell 0) and +10 (cell 6), moves
# slip with probability 0.1, gamma 0.9. The policy is a softmax over two logits per cell.
N, START, GAMMA = 7, 3, 0.9
pi = lambda th, s: np.exp(th[s] - th[s].max()) / np.exp(th[s] - th[s].max()).sum()

def episode(th, rng):
    s, traj = START, []
    while s not in (0, N - 1):
        a = rng.choice(2, p=pi(th, s)); s2 = s + (1 if a else -1) * (1 if rng.random() < 0.9 else -1)
        traj.append((s, a, {0: 1.0, N - 1: 10.0}.get(s2, 0.0), s2)); s = s2
    return traj

def value(th):                                          # exact V(start) by iterating the Bellman equation
    V = np.zeros(N)
    for _ in range(300):
        for s in range(1, N - 1):
            p = pi(th, s); q = [sum(pr * ({0: 1, N - 1: 10}.get(s2, 0) + GAMMA * V[s2] * (0 < s2 < N - 1))
                                    for pr, s2 in [(0.9, s + d), (0.1, s - d)]) for d in (-1, 1)]
            V[s] = p @ q
    return V[START]

def grad_logpi(th, s, a): g = np.zeros_like(th); g[s] = -pi(th, s); g[s, a] += 1; return g

def train(method, episodes=40, lr=0.05, seed=0):
    rng, th, V = np.random.default_rng(seed), np.zeros((N, 2)), np.zeros(N)
    for _ in range(episodes):
        traj, G = episode(th, rng), 0.0
        for t in reversed(range(len(traj))):
            s, a, r, s2 = traj[t]; G = r + GAMMA * G
            if method == "REINFORCE": adv = G                                    # raw return
            elif method == "REINFORCE + baseline": adv = G - V[s]; V[s] += 0.1 * adv
            else:                                                                # actor-critic: TD error
                adv = r + GAMMA * V[s2] * (0 < s2 < N - 1) - V[s]; V[s] += 0.1 * adv
            th += lr * GAMMA ** t * adv * grad_logpi(th, s, a)
    return value(th)

print(f"V(start): random policy {value(np.zeros((N, 2))):.2f}, always right {value(np.array([[0, 50.]] * N)):.2f}")
for m in ["REINFORCE", "REINFORCE + baseline", "actor-critic"]:
    v = [train(m, seed=sd) for sd in range(20)]
    print(f"{m:21s} after 40 episodes: V(start) mean {np.mean(v):.2f}, spread {np.std(v):.2f}, worst of 20 seeds {np.min(v):.2f}")

# PPO's clip: 30 gradient epochs on ONE batch of 10 episodes, with and without clipping the ratio.
rng = np.random.default_rng(1); th0 = np.zeros((N, 2)); batch = [x for _ in range(10) for x in episode(th0, rng)]
for clip in [None, 0.2]:
    th = th0.copy()
    for _ in range(30):
        g = np.zeros_like(th)
        for s, a, r, s2 in batch:
            ratio, adv = pi(th, s)[a] / pi(th0, s)[a], r + GAMMA * 5.0 * (0 < s2 < N - 1) - 5.0   # crude critic V = 5
            if clip is None or (adv > 0 and ratio < 1 + clip) or (adv < 0 and ratio > 1 - clip):
                g += ratio * adv * grad_logpi(th, s, a)                          # gradient of ratio x advantage
        th += 0.5 * g / len(batch)
    r_max = max(pi(th, s)[a] / pi(th0, s)[a] for s, a, _, _ in batch)
    print(f"{'no clip' if clip is None else 'clip 0.2':8s} largest ratio pi_new / pi_old {r_max:7.2f}")

V(start): random policy 2.84, always right 7.56
REINFORCE             after 40 episodes: V(start) mean 6.27, spread 0.79, worst of 20 seeds 3.74
REINFORCE + baseline  after 40 episodes: V(start) mean 6.64, spread 0.18, worst of 20 seeds 6.34
actor-critic          after 40 episodes: V(start) mean 6.60, spread 0.19, worst of 20 seeds 6.32
no clip  largest ratio pi_new / pi_old    1.86
clip 0.2 largest ratio pi_new / pi_old    1.24

The corridor from RL fundamentals and MDPs , now starting in cell 3, with a softmax policy (two logits per cell). A random policy is worth 2.84 from the start; always moving right is worth 7.56. Each method trains for 40 episodes, 20 seeds each, with the same learning rate.

• REINFORCE reaches a mean of 6.27, but with a spread of 0.79 across seeds and a worst seed at 3.74: some runs got unlucky early and reinforced the wrong moves.
• Adding a baseline (a running estimate of each cell's value) gives a mean of 6.64, a spread of 0.18, and a worst seed of 6.34. Same expected gradient, much less noise.
• Actor-critic with a one-step TD error is about as good here (6.60, spread 0.19), and it can update before an episode ends.
• The clip in action. Thirty gradient epochs on one batch of 10 episodes: without clipping, the most-changed action becomes 1.86 times as likely as under the policy that collected the data. With $ϵ=0.2$ it stops at 1.24. The objective gives no further reward past 1.2, so it stops there apart from the overshoot of the final step.

The link to RLHF

Fine-tuning a language model with human feedback is the same machinery at scale ( RLHF ). The policy is the language model, an action is the next token, an episode is one response, and the reward comes from a reward model trained on human preference comparisons, usually with a penalty for drifting too far from the original model. PPO was the standard optimizer for this step; its clip and the drift penalty both keep the model from changing so much at once that it starts exploiting the reward model.

Trade-offs

Common mistakes

• Using raw returns without a baseline, then blaming the learning rate for unstable training.
• Letting the baseline depend on the action, which biases the gradient.
• Reusing a batch for many epochs without PPO's clip or a KL penalty, so the policy jumps far from the data.
• Letting the policy's entropy collapse early; once it is near-deterministic it stops exploring.
• Reading the clipped objective as the agent's performance. It is a training signal; track returns.

Where you will see this

PPO and its relatives are a common default for training agents in simulation: robot locomotion, game-playing agents and control tasks with continuous actions. In language models, PPO was the optimizer of the original RLHF pipelines, and newer preference-tuning methods are usually described by how they differ from it ( RLHF ). Wherever actions are continuous or the best policy should stay random, policy gradients are the tool of choice.

Check yourself

Recap