Policy Iteration

Policy iteration repeats two steps, fully score the policy then greedily fix it, and stops with the optimal policy once no action changes anymore.

Why Does This Exist?

A known environment holds $|A|^{|S|}$ deterministic policies, far too many to enumerate. Fifty states with four moves each means $4^{50}$ candidates. Scoring one policy is easy, it is linear, but jumping straight to the best is not, because each state's value depends on choices made elsewhere.

Policy iteration breaks the knot into a loop. Policy evaluation scores the current policy exactly, then policy improvement greedily repairs it. Ronald Howard introduced the loop in 1960, and the improvement theorem guarantees every cycle helps or the policy is already optimal. Its thrifty cousin value iteration truncates the scoring step; the family portrait lives under generalized policy iteration.

Think of It Like This

Proofreading a manual with two passes

An editor polishes a manual by alternating two passes. First she reads the whole draft and scores every chapter as it stands. Then she rewrites only the chapters where a better version is now obvious from the scores.

She never rewrites blind, and never scores a draft she just changed without acting on it. When a full scoring pass suggests zero rewrites, the manual is done. Few drafts need many rounds, because each rewrite round fixes every visible flaw at once.

The analogy stops at cost: scoring a whole manual is cheap, while scoring a huge state space can dominate runtime. Real implementations therefore cut the scoring pass short once the ranking of actions is clear.

How It Actually Works

The loop runs $\pi_0 \to v_{\pi_0} \to \pi_1 \to v_{\pi_1} \to \dots$ until the policy stops moving.

Evaluation solves the Bellman expectation equation for the current policy, the same backup as its own page, run to convergence. Improvement then rebuilds the policy greedily:

\pi_{k+1}(s) = \arg\max_a \sum_{s'} P(s' \mid s, a) \, [R(s, a, s') + \gamma v_{\pi_k}(s')]

The improvement theorem gives $v_{\pi_{k+1}} \ge v_{\pi_k}$ everywhere, strictly somewhere unless the policy already satisfies Bellman optimality. Finite policies plus no repeats means finite termination at $\pi^*$ . In practice the climb takes a handful of cycles even when the bound $|A|^{|S|}$ looks terrifying.

Worked example

Three rooms with discount $\gamma = 0.5$ : room 1 can wait (loop, reward 0) or step right (reward 1 into room 2), room 2 can step back (reward 0) or exit (reward 10 into terminal room 3 worth 0).

Cycle 0 starts with waiting everywhere, $\pi_0 = [L, L]$ . Evaluation: room 1 loops on reward 0, so $v = 0.0$ ; room 2 steps back into 0.0, so $v = 0.0$ . Improvement scores $q(s_1, R) = 1 + 0.5 \times 0.0 = 1.0$ against $q(s_1, L) = 0.0$ , and $q(s_2, R) = 10.0$ against $q(s_2, L) = 0.0$ . Both flip: $\pi_1 = [R, R]$ . Not stable, so continue.

Cycle 1 evaluates the new policy: $v(s_2) = 10 + 0.5 \times 0 = 10.0$ , then $v(s_1) = 1 + 0.5 \times 10.0 = 6.0$ . Improvement rechecks: $q(s_1, L) = 0.5 \times 6.0 = 3.0 < 6.0$ and $q(s_2, L) = 0.5 \times 6.0 = 3.0 < 10.0$ . Nothing flips, so $\pi_2 = \pi_1$ and the loop stops after 2 cycles at $\pi^* = [R, R]$ with $v^* = [6.0, 10.0]$ .

Code

This snippet runs the same two room loop with plain sweeps and greedy extraction.

1def evaluate(policy, transitions, gamma=0.5, theta=1e-9):2    """Sweep Bellman expectation backups until values settle."""3    values = {s: 0.0 for s in transitions}4    while True:5        delta = 0.06        for s in transitions:7            target = sum(p * (r + gamma * values[n]) for p, n, r in transitions[s][policy[s]])8            delta = max(delta, abs(target - values[s]))9            values[s] = target10        if delta < theta:11            break12    return values13
14
15def improve(values, transitions, gamma=0.5):16    """Return the greedy policy for the given values."""17    return {s: max(acts, key=lambda a: sum(p * (r + gamma * values[n]) for p, n, r in acts[a])) for s, acts in transitions.items()}18
19
20transitions = {21    "s1": {"L": [(1.0, "s1", 0.0)], "R": [(1.0, "s2", 1.0)]},22    "s2": {"L": [(1.0, "s1", 0.0)], "R": [(1.0, "s3", 10.0)]},23    "s3": {"L": [(1.0, "s3", 0.0)], "R": [(1.0, "s3", 0.0)]},24}25policy = {"s1": "L", "s2": "L", "s3": "L"}26for cycle in (1, 2, 3):27    values = evaluate(policy, transitions)28    nxt = improve(values, transitions)29    print(cycle, nxt, round(values["s1"], 1), round(values["s2"], 1))30    if nxt == policy:31        break32    policy = nxt

Watch Out For

Polishing early policies to full precision

Converging evaluation to tiny tolerance on a policy that is about to flip wastes nearly the whole budget. Early cycles only need the ranking of actions, not their exact values.

Symptom: long silent evaluation phases with big policy changes right after. Fix: cap early evaluation at a few sweeps and tighten the threshold as flips get rare. One sweep per cycle is already a named algorithm: value iteration.

The Quick Version

Policy iteration alternates full evaluation with greedy improvement until the policy freezes.
Each cycle improves every state or certifies optimality through Bellman optimality.
Finite policies plus no repeats gives finite termination at the optimal policy.
Here it took 2 cycles: values 0.0 became 6.0 and 10.0, then no action flipped.
Few outer cycles at high inner cost is the signature tradeoff against value iteration.