Dynamic Programming: Policy Iteration

Okay, we've seen the Bellman equations, which give us a recursive relationship for value functions. The Bellman optimality equations, in particular, define the best possible value function ($V^*$ or $Q^*$). But how do we find that optimal value function and its corresponding optimal policy $\pi^*$?

Dynamic Programming (DP) offers a class of algorithms that can compute optimal policies, provided we have a complete model of the environment. This means we need to know the transition probabilities $p(s', r | s, a)$ for all states $s$, actions $a$, next states $s'$, and rewards $r$. This is a strong assumption, often not met in practice (which is why we'll study model-free methods later!), but DP provides a solid theoretical foundation.

One fundamental DP method is Policy Iteration. It finds the optimal policy by alternating between two main steps: evaluating the current policy and then improving it. Let's break it down.

Imagine you start with some arbitrary policy, let's call it $\pi_0$. It might be completely random, or it might be based on some simple heuristic. Policy Iteration aims to iteratively refine this policy until it converges to the optimal one, $\pi^*$. The process looks like this:

$\pi_0 \xrightarrow{\text{Evaluate}} V^{\pi_0} \xrightarrow{\text{Improve}} \pi_1 \xrightarrow{\text{Evaluate}} V^{\pi_1} \xrightarrow{\text{Improve}} \pi_2 \xrightarrow{\text{Evaluate}} \dots \xrightarrow{\text{Improve}} \pi^* \xrightarrow{\text{Evaluate}} V^*$

This involves two core components repeated in a loop:

1. Policy Evaluation: Given a policy $\pi$, calculate its state-value function $V^\pi$.
2. Policy Improvement: Using the calculated $V^\pi$, generate a new, potentially better policy $\pi'$ by acting greedily with respect to $V^\pi$.

This cycle continues until the policy no longer changes between iterations, indicating that we've found the optimal policy.

Policy Evaluation: How good is the current policy?

In this step, we take the current policy $\pi$ and compute the state-value function $V^\pi(s)$ for all states $s$. Remember the Bellman expectation equation for $V^\pi$:

$$V^\pi(s) = \sum_a \pi(a|s) \sum_{s', r} p(s', r | s, a) [r + \gamma V^\pi(s')]$$

For a finite MDP, this equation gives us a system of linear equations, one for each state. If there are $|S|$ states, we have $|S|$ equations and $|S|$ unknowns ($V^\pi(s)$ for each $s$). We could solve this system directly.

However, an often simpler and more generally applicable approach, especially within the iterative nature of policy iteration, is to compute $V^\pi$ iteratively. We start with an arbitrary value function $V_0$ (e.g., all zeros) and repeatedly apply the Bellman expectation equation as an update rule:

$$V_{k+1}(s) \leftarrow \sum_a \pi(a|s) \sum_{s', r} p(s', r | s, a) [r + \gamma V_k(s')] \quad \text{for all } s \in S$$

We sweep through all states, updating their values based on the values from the previous iteration ($V_k$). This process is guaranteed to converge to the true $V^\pi$ as $k \to \infty$. In practice, we stop when the changes in the value function become very small between iterations (i.e., when $\|V_{k+1} - V_k\|_{\infty} < \theta$ for some small threshold $\theta$). This iterative approach is essentially applying the Bellman operator repeatedly.

Policy Improvement: Can we do better?

Once we have a reasonably accurate estimate of $V^\pi$ for our current policy $\pi$, we ask: "Can we improve this policy?" The policy improvement theorem gives us a way to do this.

For a given state $s$, we can look ahead one step and consider the expected return if we choose an action $a$ (possibly different from $\pi(a|s)$) and then follow the existing policy $\pi$ thereafter. This is precisely the definition of the action-value function $Q^\pi(s, a)$:

$$Q^\pi(s, a) = \sum_{s', r} p(s', r | s, a) [r + \gamma V^\pi(s')]$$

Now, suppose for some state $s$, we find an action $a$ such that $Q^\pi(s, a) > V^\pi(s)$. This means that taking action $a$ once in state $s$ and then following $\pi$ is better than just following $\pi$ from the start. The policy improvement theorem states that if we construct a new policy $\pi'$ that chooses this "better" action $a$ in state $s$ (and follows $\pi$ elsewhere, or better yet, acts greedily everywhere), the new policy $\pi'$ will be strictly better than or equal to $\pi$.

The most common way to improve the policy is to make it greedy with respect to $V^\pi$. That is, for each state $s$, the new policy $\pi'(s)$ deterministically chooses the action $a$ that maximizes $Q^\pi(s, a)$:

$$\pi'(s) \leftarrow \underset{a}{\mathrm{argmax}} \, Q^\pi(s, a) = \underset{a}{\mathrm{argmax}} \sum_{s', r} p(s', r | s, a) [r + \gamma V^\pi(s')]$$

This new greedy policy $\pi'$ is guaranteed to be at least as good as $\pi$. If $V^{\pi'} = V^\pi$, then we have reached the optimal value function $V^*$ and the optimal policy $\pi^*$, because we satisfy the Bellman optimality equation. If $V^{\pi'} \neq V^\pi$, then $\pi'$ is strictly better, and we loop back to evaluate this improved policy.

The Policy Iteration Algorithm

Putting it together, the Policy Iteration algorithm proceeds as follows:

1. Initialization:
   • Choose an arbitrary initial policy $\pi(s)$ for all $s \in S$.
   • Choose an arbitrary initial value function $V(s)$ for all $s \in S$ (e.g., $V(s)=0$).
2. Policy Evaluation Loop:
   • Repeat until $V$ converges:
     • $\Delta \leftarrow 0$
     • For each state $s \in S$:
       • $v \leftarrow V(s)$
       • $V(s) \leftarrow \sum_a \pi(a|s) \sum_{s', r} p(s', r | s, a) [r + \gamma V(s')]$ (using the previous iteration's $V$ values on the right side, or the latest available values if updating in-place)
       • $\Delta \leftarrow \max(\Delta, |v - V(s)|)$
     • If $\Delta < \theta$ (a small positive threshold): break the evaluation loop.
3. Policy Improvement:
   • policy_stable $\leftarrow$ true
   • For each state $s \in S$:
     • old_action $\leftarrow \pi(s)$ (assuming deterministic policy for simplicity)
     • $\pi(s) \leftarrow \underset{a}{\mathrm{argmax}} \sum_{s', r} p(s', r | s, a) [r + \gamma V(s')]$
     • If old_action $\neq \pi(s)$: policy_stable $\leftarrow$ false
4. Repeat or Terminate:
   • If policy_stable is true: Stop. The current policy $\pi$ and value function $V$ are optimal ($\pi^*$ and $V^*$).
   • Otherwise: Go back to Step 2 (Policy Evaluation).

The cycle of Policy Evaluation and Policy Improvement in Policy Iteration. The process continues until the policy no longer changes after the improvement step.

Policy Iteration is guaranteed to converge to the optimal policy $\pi^*$ in a finite number of iterations for a finite MDP. This is because there are only a finite number of possible policies, and each policy improvement step results in a strictly better policy (unless it's already optimal). Since the value function associated with each policy increases or stays the same, and values are bounded, the process must terminate.

The main drawback? Each step, especially policy evaluation, can be computationally expensive, requiring potentially many sweeps through the entire state space. Furthermore, it critically depends on having the complete MDP model ($p(s', r | s, a)$), which is often unavailable in real-world problems. This limitation motivates the model-free methods we will explore in subsequent chapters.