Policy Improvement

Policy improvement upgrades a policy by switching each state to the action with the best action value, with a theorem guaranteeing no state gets worse.

Why Does This Exist?

Policy evaluation measures a policy but never changes it. Something must turn those measurements into better decisions. The naive move, tweaking actions by gut feel, can backfire because one changed action reroutes every trajectory that passes through its state.

The policy improvement theorem removes the gamble. It says: if an action looks better under your current values, adopting it everywhere can only help, from every state. That guarantee turns policy search into a steady climb and powers policy iteration, which alternates this step with evaluation until nothing changes.

Think of It Like This

Replacing one bus route after a full audit

A bus company audits its whole network under the current timetable and learns exactly what each stop earns. One audit line stands out: swapping diesel for electric on route 7 earns more even if every other route stays exactly as it is.

So they swap route 7 permanently. Buses that never touch route 7 still earn more, because breakdowns and detours occasionally land them on it. One local swap, verified against the full audit, lifts the entire network. That is greedification: change what the numbers favor, keep the proof that nothing else suffers.

The analogy stops at stationarity: it assumes traffic never adapts to the new buses. The theorem needs the same, a fixed environment that does not fight back.

How It Actually Works

Start from evaluated values $v_\pi(s)$ , the expected return from $s$ under current policy $\pi$ . For any candidate action $a$ , the action value $q_\pi(s, a)$ is what you would earn by taking $a$ once and following $\pi$ after:

q_\pi(s, a) = \sum_{s', r} p(s', r \mid s, a) \, [r + \gamma v_\pi(s')]

Here $p(s', r \mid s, a)$ is the environment model and $\gamma$ the discount factor. The new policy picks the winner in each state:

\pi'(s) = \arg\max_a q_\pi(s, a)

The theorem: if $q_\pi(s, \pi'(s)) \ge v_\pi(s)$ everywhere, then $v_{\pi'}(s) \ge v_\pi(s)$ everywhere. Proof sketch: unroll the inequality one step at a time, each successor satisfying the same inequality, until the discounted tail vanishes and only $\pi'$ returns remain. If nothing changes anywhere, the values satisfy the Bellman optimality equation and the policy is optimal.

Worked example

Two depots, discount $\gamma = 0.9$ . Depot 1 offers cautious $a_1$ or bold $a_2$ ; depot 2 has one fixed action. Under baseline $\pi$ with $\pi(s_1) = a_1$ , solving the Bellman system gives $v_\pi(s_1) = 18.18$ and $v_\pi(s_2) = 22.73$ .

Score the candidates in depot 1 against those values. Staying cautious reproduces the baseline: $q_\pi(s_1, a_1) = 18.18$ . Going bold gives:

q_\pi(s_1, a_2) = 0.9 \times (0.9 \times 22.73) + 0.1 \times (-1.0 + 0.9 \times 18.18) = 19.95

Since $19.95 > 18.18$ , greedify: $\pi'(s_1) = a_2$ . Re-evaluating under the new policy gives $v_{\pi'}(s_1) = 23.71$ and $v_{\pi'}(s_2) = 26.76$ . Depot 2 kept its action yet gained $+4.03$ , because detours keep landing its buses on the improved depot 1. That ripple is the diagram's right panel.

Code

This snippet replays the greedification with plain arithmetic over the two depots.

1def q_value(transitions, values, gamma=0.9):2    """One-step lookahead: sum over p * (r + gamma * V)."""3    return sum(p * (r + gamma * values[s]) for p, s, r in transitions)4
5
6v = {"s1": 18.18, "s2": 22.73}7a1 = [(0.8, "s1", 1.0), (0.2, "s2", 1.0)]8a2 = [(0.9, "s2", 0.0), (0.1, "s1", -1.0)]9q1 = q_value(a1, v)10q2 = q_value(a2, v)11print(round(q1, 2))12print(round(q2, 2))13print("a2" if q2 > q1 else "a1")

Watch Out For

Greedifying on half baked values

The guarantee assumes exact $v_\pi$ . Act on noisy or half converged estimates and the "better" action may be worse, so the policy degrades instead of climbing.

Symptom: values fall after an improvement step that the numbers promised would help. Fix: converge evaluation tightly before greedifying in exact settings, and use conservative or trust region updates once values are approximated from samples.

Breaking ties at random

When two actions tie, flipping between them each round stops the outer loop from ever settling, even with fully converged values.

Symptom: the policy oscillates forever while values sit still. Fix: keep the current action on ties within tolerance, and order any remaining ties by a fixed rule. Stability of the policy, not just the values, is the stopping signal.

The Quick Version

Policy improvement sets $\pi'(s) = \arg\max_a q_\pi(s, a)$ using current values.
The theorem guarantees $v_{\pi'}(s) \ge v_\pi(s)$ from every state, never just on average.
Gains ripple: here depot 2 rose from 22.73 to 26.76 without changing its own action.
No change anywhere means Bellman optimality, so the policy is optimal.
Exact values and deterministic ties are the two conditions that keep the guarantee honest.