Policy Iteration

A finite Markov decision process (MDP) is specified by a state set $\mathcal{S}$, an action set $\mathcal{A}$, transition probabilities $p(s' \mid s, a)$, a reward function $r(s, a, s')$, and a discount factor $\gamma \in [0, 1)$. A policy $\pi$ maps states to (distributions over) actions. The state-value function $v_\pi(s)$ is the expected discounted return starting from $s$ and following $\pi$ thereafter, satisfying the Bellman Optimality Equation's cousin, the Bellman expectation equation:

Policy evaluation is the task of computing $v_\pi$ for a fixed policy $\pi$. Assuming Iterative Policy Evaluation as a known primitive, this is solved by repeatedly applying the Bellman expectation backup as an update rule until the value function stops changing (or changes by less than a tolerance $\theta$).

Policy improvement takes a value function $v_\pi$ and produces a new, deterministic policy $\pi'$ by acting greedily:

Policy iteration is the composition of these two steps, repeated: evaluate $\pi_0$ to get $v_{\pi_0}$, improve to get $\pi_1$, evaluate $\pi_1$ to get $v_{\pi_1}$, improve to get $\pi_2$, and so on. The sequence of policies is monotonically non-decreasing in value at every state, and because the number of deterministic policies over a finite state–action space is finite, the sequence must reach a fixed point — a policy that is greedy with respect to its own value function. Such a fixed point satisfies the Bellman optimality equation exactly, so it is an optimal policy.

The algorithm has an outer loop (policy improvement) wrapping an inner loop (policy evaluation). The inner loop can be run to full convergence (classical policy iteration) or truncated after a fixed number of sweeps (a variant that shades continuously into value iteration). ``` Algorithm: Policy Iteration Input: MDP (S, A, p, r, γ), evaluation tolerance θ > 0 Output: optimal policy π, optimal value function v 1. Initialize: V(s) ← 0 for all s in S # or any initial guess π(s) ← arbitrary action in A(s) # for all s in S 2. repeat (outer loop: "policy iteration")

# ---- Step A: Policy Evaluation ---- repeat Δ ← 0 for each s in S: v_old ← V(s) V(s) ← Σ_{s'} p(s' | s, π(s)) · [ r(s, π(s), s') + γ · V(s') ] Δ ← max(Δ, |v_old − V(s)|) until Δ < θ

# ---- Step B: Policy Improvement ---- policy_stable ← true for each s in S: old_action ← π(s) π(s) ← argmax_a Σ_{s'} p(s' | s, a) · [ r(s, a, s') + γ · V(s') ] if old_action ≠ π(s): policy_stable ← false 3. until policy_stable = true 4. return π, V ``` Two implementation notes matter in practice. First, ties in the `argmax` should be broken consistently (e.g. keep the old action if it is still optimal) — this avoids policies oscillating between equally-good actions without ever satisfying `policy_stable`. Second, evaluation does not need to restart from zero each outer iteration: warm-starting $V$ from the previous policy's values converges the inner loop faster, since $\pi_{k+1}$ is typically close to $\pi_k$.

Per-iteration cost. Let $n = |\mathcal{S}|$ and $k = |\mathcal{A}|$ (assume a bounded action set per state). One sweep of policy evaluation costs $O(n \cdot b)$, where $b$ is the branching factor of the transition function (number of possible next states per state–action pair, often treated as $O(n)$ in the worst case for a dense MDP, but typically much smaller for structured problems). A single evaluation phase runs some number of sweeps $E$ until convergence to tolerance $\theta$; empirically $E$ is small (tens of sweeps) and is bounded theoretically because the Bellman expectation backup is a $\gamma$-contraction in the max-norm, giving $E = O\!\left(\frac{\log(1/\theta)}{\log(1/\gamma)}\right)$. Policy improvement costs $O(n \cdot k \cdot b)$, one $Q$-value computation per action per state.

Outer loop cost. The number of outer iterations is bounded by the number of deterministic policies, $k^n$, since no policy can repeat (proved below) — but this bound is astronomically loose. In practice the number of outer iterations observed is typically small, often logarithmic-feeling in problem size, though no tight polynomial bound is known in general for arbitrary discount factors; for fixed discount factor $\gamma$, policy iteration is known to run in strongly polynomial time in $n$, $k$, and $1/(1-\gamma)$ (a result following from connections to linear programming and simplex-type analyses of MDPs). This is a much sharper guarantee than the naive $k^n$ bound suggests.

Space complexity. $O(n)$ for the value function, $O(n)$ for the policy, plus the model itself which may require $O(n^2 k)$ to store dense transition probabilities (often sparse in practice).

Correctness — the Policy Improvement Theorem. This is the crux of why the algorithm works, so it deserves a full argument rather than an assertion. Claim: If $\pi'$ is greedy with respect to $v_\pi$, i.e. $q_\pi(s, \pi'(s)) \geq v_\pi(s)$ for all $s$, then $v_{\pi'}(s) \geq v_\pi(s)$ for all $s$. Proof sketch. Start from the hypothesis and unroll one step at a time:

Now substitute the same inequality for $v_\pi(S_{t+1}) \le q_\pi(S_{t+1}, \pi'(S_{t+1}))$ inside the expectation, and repeat this substitution indefinitely:

Each substitution is valid because the bound $q_\pi(s,a) \ge v_\pi(s)$ holds at every state, not just the starting one, so it can be chained along the entire trajectory generated by following $\pi'$. Because rewards are bounded and $\gamma < 1$, this infinite telescoping sum converges and the interchange of limit and expectation is justified. $\blacksquare$

Why iteration terminates at optimality. The theorem shows the value never decreases across an improvement step. Suppose the loop reaches a policy $\pi$ that is not improved, meaning $\pi$ is already greedy with respect to its own $v_\pi$:

This is exactly the Bellman Optimality Equation. Since that equation has a unique solution $v_*$ (a consequence of the Bellman optimality operator being a $\gamma$-contraction in max-norm, hence having a unique fixed point by the Banach fixed-point theorem), $v_\pi = v_*$ and $\pi$ is optimal. Finally, termination in finitely many steps follows because: (1) each improvement step either leaves the policy unchanged (in which case we have just shown it is optimal and we stop) or strictly increases $v(s)$ at some state without decreasing it anywhere; (2) there are only $k^n$ distinct deterministic policies; (3) a policy that strictly improves can never recur, since its value is strictly larger than at any earlier visit. Hence the sequence of distinct policies visited is finite, and the loop must halt at an optimal policy.

Tradeoffs. Policy iteration requires a full model — it is not directly applicable when $p$ and $r$ are unknown (that setting motivates temporal-difference and Monte Carlo control methods). Running policy evaluation to full convergence at every outer step is often wasteful: most of the benefit of evaluation is captured in the first few sweeps, which is precisely the insight exploited by Value Iteration, where evaluation is truncated to a single sweep.

Consider a tiny 4-state corridor MDP: states $\{1, 2, 3, 4\}$ laid out in a line, with state 4 being terminal (goal, reward 0 on entering, episode ends). From non-terminal states, two actions are available: Right (move toward state 4) and Left (move toward state 1, or stay at state 1 if already there). Every non-terminal transition gives reward $-1$. Discount $\gamma = 1$ for this episodic example (standard for finite-horizon-like gridworld/corridor tasks). Transitions are deterministic: Right moves $s \to s+1$ (clamped at 4, terminal), Left moves $s \to \max(s-1, 1)$.

Initialization. $\pi_0(s) = \text{Left}$ for all $s \in \{1,2,3\}$, $V_0(s) = 0$ for all $s$.

Iteration 1 — Evaluate $\pi_0$ (all-Left). Under all-Left, state 1 loops on itself forever incurring $-1$ each step and never reaches the terminal — its true value is $-\infty$ (or, with the iterative sweeps used in practice, a very large negative number after enough sweeps; conceptually this is a bad policy). States 2 and 3 eventually get pushed to state 1 and get stuck there too. Sweeps of iterative policy evaluation drive all three values toward $-\infty$ (or, truncated at some large negative cutoff for illustration), confirming numerically that this policy is bad everywhere.

Improve. For each state, compare $q(s,\text{Left})$ against $q(s,\text{Right})$ using the (very negative) $V$ just computed.

Because state 4 is terminal with $V(4) = 0$ fixed, $q(3, \text{Right}) = -1$ is finite and strictly better than looping indefinitely, so Right wins at every state. New policy $\pi_1(s) = \text{Right}$ for all $s$.

Iteration 2 — Evaluate $\pi_1$ (all-Right). This policy is now deterministic and acyclic (every state marches straight to the terminal), so exact values are computable directly:

Improve again.

`policy_stable = true` — no action changed. The loop halts. $\pi_1$ is optimal, with $v_* (1,2,3) = (-3,-2,-1)$, matching the obvious ground truth: the shortest path to the goal from state $s$ takes $4-s$ steps. This example is deliberately non-symmetric: the starting policy (all-Left) is close to worst-case, forcing a real improvement step rather than a trivial no-op, and the terminal is off-center enough that the value gradient is easy to sanity-check by hand.

Policy iteration is the backbone of classical operations-research tools for inventory control, queueing systems, and maintenance scheduling, where transition models are known from physical or statistical modeling and exact solutions are required. In robotics, it underlies motion-planning solvers that treat navigation as an MDP with a known map, computing provably shortest-expected-cost policies before deployment. In finance, portfolio-rebalancing and options-pricing models formulated as MDPs use policy-iteration-style solvers when transition dynamics can be estimated from historical data. Conceptually, its evaluate-then-improve template — termed generalized policy iteration by Sutton and Barto — is the schema instantiated by nearly every practical reinforcement-learning algorithm: SARSA and Q-learning approximate the evaluation and improvement steps with sampled transitions instead of a known model; actor-critic architectures make the split explicit, with the critic performing (approximate) policy evaluation and the actor performing (approximate) policy improvement; and even large-scale deep-RL systems like AlphaGo Zero's self-play loop can be read as policy iteration where evaluation is done via Monte Carlo tree search and improvement is distilled back into a neural network policy. Anyone building on Value Iteration, actor-critic methods, or model-based RL is, at bottom, building on this same alternation.

Ask the tutor: How does modified policy iteration (truncating evaluation to $k$ sweeps) formally interpolate between policy iteration and value iteration?