Policy Iteration

Why Does This Exist?

In a known Markov Decision Process (MDP), the goal is to discover an optimal policy $\pi^*$ that maximizes expected cumulative discounted returns. However, searching through all possible policies is a combinatorial nightmare: for a discrete MDP with state space $S$ and action space $A$, there are $|A|^{|S|}$ distinct deterministic policies. In a simple grid with only 50 states and 4 movement directions, brute-force evaluation must inspect $4^{50} \approx 1.27 \times 10^{30}$ policies—a computational impossibility.

Evaluating an existing policy is a tractable linear problem, but finding the globally optimal policy directly is non-linear because choices made in one state dictate the values of upstream states. Without a systematic mechanism to navigate the space of policies, algorithms either stall in local heuristics or resort to intractable brute-force enumeration.

Policy iteration, introduced by Ronald Howard in 1960, solves this by decomposing control into two alternating sub-problems:

1. Policy Evaluation: Compute the exact long-term value function $v_\pi$ for the agent's current policy $\pi$.
2. Policy Improvement: Construct an improved policy $\pi'$ by acting greedily with respect to the freshly calculated value landscape $v_\pi$.

The mathematical engine powering this loop is the Policy Improvement Theorem. It guarantees that each greedy update produces a policy that is strictly better than the previous one across all states, or identical if the policy is already globally optimal. Because finite MDPs possess a finite number of policies and no policy can ever repeat, policy iteration provably terminates at the exact globally optimal policy $\pi^*$ in a finite number of iterations.

Think of It Like This

How It Actually Works

The Evaluation-Improvement Loop and Convergence Guarantees

Policy iteration alternates between two foundational operators in a closed loop:

$\pi_0 \xrightarrow{\text{Evaluation}} v_{\pi_0} \xrightarrow{\text{Improvement}} \pi_1 \xrightarrow{\text{Evaluation}} v_{\pi_1} \xrightarrow{\text{Improvement}} \dots \xrightarrow{\text{Stable}} \pi^*$

1. Policy Evaluation ($E$)

Given a policy $\pi_k$, policy evaluation calculates the expected return $v_{\pi_k}(s)$ for every state $s \in S$. The value function satisfies the linear Bellman expectation equation:

$v_{\pi_k}(s) = \sum_{a \in A} \pi_k(a \mid s) \sum_{s' \in S} P(s' \mid s, a) \left[ R(s, a, s') + \gamma v_{\pi_k}(s') \right]$

For deterministic policies where $\pi_k(s)$ selects a single action:

$v_{\pi_k}(s) = \sum_{s' \in S} P(s' \mid s, \pi_k(s)) \left[ R(s, \pi_k(s), s') + \gamma v_{\pi_k}(s') \right]$

This forms a system of $|S|$ linear equations with $|S|$ unknowns. It can be solved either directly in matrix form via $(I - \gamma P_{\pi_k}) v_{\pi_k} = r_{\pi_k}$, or iteratively through successive relaxation sweeps:

$v^{(m+1)}(s) \leftarrow \sum_{s'} P(s' \mid s, \pi_k(s)) \left[ R(s, \pi_k(s), s') + \gamma v^{(m)}(s') \right]$

The sweeps repeat until $\max_{s \in S} |v^{(m+1)}(s) - v^{(m)}(s)| < \theta$, where $\theta$ is a small positive convergence threshold.

2. Policy Improvement ($I$)

Once the exact value function $v_{\pi_k}$ is known, the agent evaluates the action-value function $q_{\pi_k}(s, a)$ for every possible action $a \in A$:

$q_{\pi_k}(s, a) = \sum_{s' \in S} P(s' \mid s, a) \left[ R(s, a, s') + \gamma v_{\pi_k}(s') \right]$

The new policy $\pi_{k+1}$ is extracted greedily by choosing the action that maximizes this one-step lookahead:

$\pi_{k+1}(s) = \operatorname{argmax}_{a \in A} q_{\pi_k}(s, a)$

If multiple actions tie for the maximum, the agent may pick any maximizing action deterministically or split selection probability equally.

3. The Policy Improvement Theorem and Finite Convergence

The theoretical backbone of the algorithm is the Policy Improvement Theorem: If $q_{\pi_k}(s, \pi_{k+1}(s)) \ge v_{\pi_k}(s)$ for all $s \in S$, then:

$v_{\pi_{k+1}}(s) \ge v_{\pi_k}(s) \quad \forall s \in S$

If the inequality is strict in at least one state, the entire policy is strictly superior. If no action change occurs—meaning $\pi_{k+1}(s) = \pi_k(s)$ across all states—the policy is stable. At this point, the greedy action selection satisfies:

$v_{\pi_k}(s) = \max_{a \in A} \sum_{s'} P(s' \mid s, a) \left[ R(s, a, s') + \gamma v_{\pi_k}(s') \right]$

This is the Bellman Optimality Equation. Therefore, $v_{\pi_k} = v^*$ and $\pi_k = \pi^*$.

Because an MDP with finite state space $S$ and finite action space $A$ has at most $|A|^{|S|}$ unique deterministic policies, and each iteration strictly increases the value function unless already optimal, no policy can ever be evaluated twice. Consequently, policy iteration is guaranteed to terminate in at most $|A|^{|S|}$ outer iterations. In practice, it converges in a tiny fraction of that bound (often fewer than 10 iterations even on complex problems).

Worked numerical example

Consider a 3-state chain MDP:

• States: $S = \{s_1, s_2, s_3\}$. State $s_3$ is a terminal absorbing state with fixed value $v(s_3) = 0$.
• Actions: $A = \{\text{L (Left)}, \text{R (Right)}\}$.
• Discount factor: $\gamma = 0.5$.
• Transition Dynamics & Rewards:
  • In $s_1$: action $\text{L}$ loops back to $s_1$ with reward $R = 0$. Action $\text{R}$ moves to $s_2$ with reward $R = 1$.
  • In $s_2$: action $\text{L}$ moves to $s_1$ with reward $R = 0$. Action $\text{R}$ transitions to terminal $s_3$ with reward $R = 10$.
  • In terminal $s_3$: absorbing state, all actions yield reward 0.

Cycle 0: Initial Policy $\pi_0$

We initialize an arbitrary policy choosing $\text{L}$ everywhere: $\pi_0(s_1) = \text{L}, \quad \pi_0(s_2) = \text{L}$

Step 0.1: Policy Evaluation of $\pi_0$ Set up the Bellman expectation equations:

• $v(s_1) = R(s_1, \text{L}) + \gamma v(s_1) = 0 + 0.5 \cdot v(s_1) \implies 0.5 v(s_1) = 0 \implies v_{\pi_0}(s_1) = 0.0$
• $v(s_2) = R(s_2, \text{L}) + \gamma v(s_1) = 0 + 0.5 \cdot (0.0) = 0.0 \implies v_{\pi_0}(s_2) = 0.0$

$v_{\pi_0} = [v(s_1) = 0.0, \; v(s_2) = 0.0]$

Step 0.2: Policy Improvement on $\pi_0$ Evaluate action-values $q(s, a) = R + \gamma v(s')$:

• For state $s_1$: $q(s_1, \text{L}) = 0 + 0.5 \cdot v(s_1) = 0 + 0.5(0.0) = 0.0$ $q(s_1, \text{R}) = 1 + 0.5 \cdot v(s_2) = 1 + 0.5(0.0) = 1.0$ $\pi_1(s_1) = \operatorname{argmax}_a \{0.0, 1.0\} = \text{R} \quad (\text{switched from L})$

• For state $s_2$: $q(s_2, \text{L}) = 0 + 0.5 \cdot v(s_1) = 0 + 0.5(0.0) = 0.0$ $q(s_2, \text{R}) = 10 + 0.5 \cdot v(s_3) = 10 + 0.5(0.0) = 10.0$ $\pi_1(s_2) = \operatorname{argmax}_a \{0.0, 10.0\} = \text{R} \quad (\text{switched from L})$

Policy changed from $[\text{L}, \text{L}]$ to $[\text{R}, \text{R}]$. Policy is not stable. Proceed to Cycle 1.

Cycle 1: Policy $\pi_1 = [\text{R}, \text{R}]$

Step 1.1: Policy Evaluation of $\pi_1$ Evaluate values under the newly proposed policy:

• In $s_2$: action $\text{R}$ transitions to terminal $s_3$: $v_{\pi_1}(s_2) = 10 + 0.5 \cdot v(s_3) = 10 + 0.5(0.0) = 10.0$
• In $s_1$: action $\text{R}$ transitions to $s_2$: $v_{\pi_1}(s_1) = 1 + 0.5 \cdot v_{\pi_1}(s_2) = 1 + 0.5(10.0) = 1 + 5.0 = 6.0$

$v_{\pi_1} = [v(s_1) = 6.0, \; v(s_2) = 10.0]$

Notice the monotonic value surge: $v_{\pi_1}(s_1) = 6.0 > v_{\pi_0}(s_1) = 0.0$ and $v_{\pi_1}(s_2) = 10.0 > v_{\pi_0}(s_2) = 0.0$.

Step 1.2: Policy Improvement on $\pi_1$ Re-compute action-values using the new value landscape:

• For state $s_1$: $q(s_1, \text{L}) = 0 + 0.5 \cdot v(s_1) = 0 + 0.5(6.0) = 3.0$ $q(s_1, \text{R}) = 1 + 0.5 \cdot v(s_2) = 1 + 0.5(10.0) = 6.0$ $\pi_2(s_1) = \operatorname{argmax}_a \{3.0, 6.0\} = \text{R} \quad (\text{unchanged})$

• For state $s_2$: $q(s_2, \text{L}) = 0 + 0.5 \cdot v(s_1) = 0 + 0.5(6.0) = 3.0$ $q(s_2, \text{R}) = 10 + 0.5 \cdot v(s_3) = 10 + 0.5(0.0) = 10.0$ $\pi_2(s_2) = \operatorname{argmax}_a \{3.0, 10.0\} = \text{R} \quad (\text{unchanged})$

Updated policy $\pi_2 = [\text{R}, \text{R}]$. Because $\pi_2(s) = \pi_1(s)$ for all states, the policy is stable. The algorithm terminates.

Final Optimal Solution:

• Optimal Policy: $\pi^*(s_1) = \text{R}, \; \pi^*(s_2) = \text{R}$
• Optimal Values: $v^*(s_1) = 6.0, \; v^*(s_2) = 10.0$

Code

1from typing import Dict, List, Tuple2
3# Type aliases representing discrete MDP components4# State -> Action -> list of (transition_probability, next_state, immediate_reward)5MDPTransitions = Dict[str, Dict[str, List[Tuple[float, str, float]]]]6Policy = Dict[str, str]7ValueFunction = Dict[str, float]8
9
10def policy_evaluation(11    policy: Policy,12    transitions: MDPTransitions,13    gamma: float = 0.5,14    theta: float = 1e-9,15) -> ValueFunction:16    """Computes exact state values V_pi for a deterministic policy via Bellman sweeps.17
18    Args:19        policy: Mapping of each state to its selected action.20        transitions: MDP transition dynamics and reward mapping.21        gamma: Discount factor in range [0, 1).22        theta: Convergence threshold for stopping iterative sweeps.23
24    Returns:25        Converged state-value function mapping each state to its expected return.26    """27    states = list(transitions.keys())28    V: ValueFunction = {s: 0.0 for s in states}29
30    while True:31        delta = 0.032        for s in states:33            chosen_action = policy[s]34            # Bellman expectation equation:35            # V(s) = sum_{s'} P(s'|s, a) * [R(s, a, s') + gamma * V(s')]36            expected_val = sum(37                prob * (reward + gamma * V[next_s])38                for prob, next_s, reward in transitions[s][chosen_action]39            )40            delta = max(delta, abs(expected_val - V[s]))41            V[s] = expected_val42
43        if delta < theta:44            break45
46    return V47
48
49def policy_improvement(50    V: ValueFunction,51    transitions: MDPTransitions,52    current_policy: Policy,53    gamma: float = 0.5,54) -> Tuple[Policy, bool]:55    """Greedily extracts a new policy with respect to the value function V.56
57    Args:58        V: Evaluated state values.59        transitions: MDP transition dynamics and reward mapping.60        current_policy: The policy that generated V.61        gamma: Discount factor in range [0, 1).62
63    Returns:64        A tuple of (new_policy, is_stable).65    """66    new_policy: Policy = {}67    policy_stable = True68
69    for s, available_actions in transitions.items():70        q_values: Dict[str, float] = {}71
72        for a, outcomes in available_actions.items():73            # Action-value Q(s, a) lookahead74            q_values[a] = sum(75                prob * (reward + gamma * V[next_s])76                for prob, next_s, reward in outcomes77            )78
79        # Greedily choose action maximizing Q(s, a)80        best_action = max(q_values, key=lambda action: q_values[action])81        new_policy[s] = best_action82
83        if best_action != current_policy[s]:84            policy_stable = False85
86    return new_policy, policy_stable87
88
89def policy_iteration(90    transitions: MDPTransitions,91    gamma: float = 0.5,92    theta: float = 1e-9,93) -> Tuple[Policy, ValueFunction, int]:94    """Runs the full Policy Iteration algorithm until policy convergence.95
96    Returns:97        A tuple of (optimal_policy, optimal_values, total_cycles).98    """99    # Initialize with an arbitrary deterministic policy (first action in each state)100    policy: Policy = {s: list(transitions[s].keys())[0] for s in transitions}101    cycles = 0102
103    while True:104        cycles += 1105        # Step 1: Policy Evaluation106        V = policy_evaluation(policy, transitions, gamma, theta)107
108        # Step 2: Policy Improvement109        new_policy, is_stable = policy_improvement(V, transitions, policy, gamma)110
111        policy = new_policy112        if is_stable:113            break114
115    return policy, V, cycles116
117
118if __name__ == "__main__":119    # 3-state chain MDP from the worked numerical example:120    # State s3 is absorbing (terminal) with 0 return.121    mdp: MDPTransitions = {122        "s1": {123            "L": [(1.0, "s1", 0.0)],124            "R": [(1.0, "s2", 1.0)],125        },126        "s2": {127            "L": [(1.0, "s1", 0.0)],128            "R": [(1.0, "s3", 10.0)],129        },130        "s3": {131            "L": [(1.0, "s3", 0.0)],132            "R": [(1.0, "s3", 0.0)],133        },134    }135
136    optimal_policy, optimal_values, num_cycles = policy_iteration(mdp, gamma=0.5)137
138    print(f"Optimal Policy: s1 -> {optimal_policy['s1']}, s2 -> {optimal_policy['s2']}")139    print(f"Optimal Values: V(s1) = {optimal_values['s1']:.1f}, V(s2) = {optimal_values['s2']:.1f}")140    print(f"Cycles to Convergence: {num_cycles}")141
142    # Expected output:143    # Optimal Policy: s1 -> R, s2 -> R144    # Optimal Values: V(s1) = 6.0, V(s2) = 10.0145    # Cycles to Convergence: 2

Watch Out For

The Quick Version

• Alternating engine: Policy iteration alternates between full policy evaluation ($v_{\pi_k}$) and greedy policy improvement ($\pi_{k+1}(s) = \operatorname{argmax}_a q(s, a)$).
• Monotonic improvement: By the Policy Improvement Theorem, each greedy step guarantees $v_{\pi_{k+1}}(s) \ge v_{\pi_k}(s)$ across every state, eliminating regression.
• Finite convergence: Termination triggers when no action flips ($\pi_{k+1} = \pi_k$), provably reaching the unique optimal policy $\pi^*$ in at most $|A|^{|S|}$ outer iterations.
• Computation tradeoff: Policy iteration makes dramatic policy improvements in very few outer cycles compared to value iteration, but incurs high per-cycle cost from inner evaluation loops.

What to Read Next

• Policy Evaluation
• Policy Improvement
• Value Iteration
• Generalized Policy Iteration (GPI)
• Tabular Solution Methods