Chapter 5 · Learning from rewards: MDPs and reinforcement learning

§5.2 Value iteration and policy iteration

0:00/9:27
AI-narrated by Kokoro

§5.1 introduced the value function VπV^\pi and the action-value QπQ^\pi but did not say how to compute them. This section closes that gap for the case the rest of the chapter does not assume — the case where the transition PP and the reward RR are known. That assumption sounds restrictive, and from a robotics standpoint it is: you almost never have a clean tabular PP for a real arm. But the two algorithms that fall out of it — value iteration and policy iteration — are the algorithmic foundation that every later method in this book, from Q-learning in §5.3 to the PPO update of §7.3 to the critic inside SAC, inherits its structure from. Read this section as the place where the shape of all subsequent RL algorithms is set, with the inconvenient assumption stripped away later.

Both algorithms rest on a single equation. We spend the first half of the section understanding that equation, then turn it into code.

The Bellman equations

Fix a policy π\pi and a state ss. The value Vπ(s)V^\pi(s) is the expected return from ss under π\pi. The defining recursive identity is the Bellman expectation equation:

Vπ(s)  =  aπ(as)sP(ss,a)[R(s,a,s)+γVπ(s)].V^\pi(s) \;=\; \sum_{a} \pi(a \mid s) \sum_{s'} P(s' \mid s, a) \,\bigl[\,R(s, a, s') + \gamma V^\pi(s')\,\bigr].

In words: the value of ss under π\pi is the average — over the actions π\pi would pick and the transitions those actions would induce — of the immediate reward plus the discounted value of where you land. The equation is a statement about consistency, not about optimality. It is true for any policy. The corresponding identity for QπQ^\pi is

Qπ(s,a)  =  sP(ss,a)[R(s,a,s)+γaπ(as)Qπ(s,a)],Q^\pi(s, a) \;=\; \sum_{s'} P(s' \mid s, a)\,\bigl[\,R(s, a, s') + \gamma \sum_{a'} \pi(a' \mid s') Q^\pi(s', a')\,\bigr],

which differs only in committing to action aa at the current step before falling back on π\pi.

The Bellman optimality equation is the same identity, written for the optimal value V(s)=maxπVπ(s)V^\star(s) = \max_\pi V^\pi(s):

V(s)  =  maxasP(ss,a)[R(s,a,s)+γV(s)].V^\star(s) \;=\; \max_{a}\, \sum_{s'} P(s' \mid s, a)\,\bigl[\,R(s, a, s') + \gamma V^\star(s')\,\bigr].

The aπ(as)\sum_a \pi(a \mid s) has become a maxa\max_a. The optimal value of a state is the best you can do in one step plus the optimal value of where that step lands you. Bellman (1957) proved that this equation has a unique fixed point in VV for γ<1\gamma < 1 and bounded rewards. Both algorithms in this section are ways to find that fixed point.

Value iteration

Define the Bellman optimality operator TT^\star acting on a value function VV:

(TV)(s)  =  maxasP(ss,a)[R(s,a,s)+γV(s)].(T^\star V)(s) \;=\; \max_{a}\, \sum_{s'} P(s' \mid s, a)\,\bigl[\,R(s, a, s') + \gamma V(s')\,\bigr].

A value function is optimal exactly when V=TVV = T^\star V. The algorithm called value iteration picks an initial V0V_0 — usually the all-zeros vector indexed by S\mathcal{S} — and applies TT^\star repeatedly:

Vk+1(s)  =  (TVk)(s).V_{k+1}(s) \;=\; (T^\star V_k)(s).

The key analytical fact is that TT^\star is a γ\gamma-contraction in the \ell_\infty norm: for any two value functions VV and VV', TVTVγVV\|T^\star V - T^\star V'\|_\infty \leq \gamma \|V - V'\|_\infty. Banach’s fixed-point theorem then guarantees that VkVV_k \to V^\star at a geometric rate, and after kk steps the worst-case error is at most γkV0V\gamma^k \|V_0 - V^\star\|_\infty. With γ=0.95\gamma = 0.95, an error tolerance of 10310^{-3} on rewards bounded by 11 takes roughly 150 sweeps — small enough to run on a laptop for tabular problems and small enough to fit in a footnote in every RL textbook.

Once VV^\star is in hand, the optimal policy is the greedy policy with respect to it:

π(s)  =  argmaxasP(ss,a)[R(s,a,s)+γV(s)].\pi^\star(s) \;=\; \arg\max_{a}\, \sum_{s'} P(s' \mid s, a)\,\bigl[\,R(s, a, s') + \gamma V^\star(s')\,\bigr].

Here is value iteration as a 12-line Python loop on the §5.1 gridworld:

import numpy as np

n_states, n_actions = 16, 4   # 4x4 grid; N, S, E, W
P = build_transition_tensor()  # shape (n_states, n_actions, n_states)
R = build_reward_tensor()      # shape (n_states, n_actions, n_states)
gamma, tol = 0.95, 1e-6

V = np.zeros(n_states)
while True:
    Q = (P * (R + gamma * V[None, None, :])).sum(axis=2)  # (S, A)
    V_new = Q.max(axis=1)
    if np.max(np.abs(V_new - V)) < tol:
        break
    V = V_new
pi_star = Q.argmax(axis=1)

Two lines do all the work. Q = (P * (R + gamma * V[None, None, :])).sum(axis=2) is the Bellman operator written as a tensor contraction; V_new = Q.max(axis=1) is the maxa\max_a. On the deterministic gridworld of §5.1, this loop converges in well under a hundred sweeps and recovers the optimal policy of “move toward the goal,” with VV^\star taking values close to 1/(1γ)20-1/(1-\gamma) \approx -20 in the corner cells and 0 at the absorbing goal.

Policy iteration

Howard (1960) proposed an alternative scheme that splits the fixed-point problem into two simpler ones. Start with an arbitrary policy π0\pi_0. Repeat two steps:

Policy evaluation. Compute VπkV^{\pi_k} — the value of the current policy — by solving the Bellman expectation equation. Because πk\pi_k is fixed, this is a linear system in S|\mathcal{S}| unknowns:

Vπk(s)  =  sPπk(ss)[Rπk(s,s)+γVπk(s)],V^{\pi_k}(s) \;=\; \sum_{s'} P_{\pi_k}(s' \mid s)\,\bigl[\,R_{\pi_k}(s, s') + \gamma V^{\pi_k}(s')\,\bigr],

where PπkP_{\pi_k} and RπkR_{\pi_k} marginalize over the policy. One can either solve this linear system directly with a matrix inverse — fine for a 16-state gridworld, painful for a 1000-state model — or run a small inner loop of Bellman expectation updates to convergence.

Policy improvement. Replace πk\pi_k with the greedy policy with respect to VπkV^{\pi_k}:

πk+1(s)  =  argmaxasP(ss,a)[R(s,a,s)+γVπk(s)].\pi_{k+1}(s) \;=\; \arg\max_{a}\, \sum_{s'} P(s' \mid s, a)\,\bigl[\,R(s, a, s') + \gamma V^{\pi_k}(s')\,\bigr].

The policy improvement theorem (Howard 1960; Sutton & Barto 2018 §4.2) guarantees that πk+1\pi_{k+1} is at least as good as πk\pi_k everywhere, with strict improvement in at least one state unless πk\pi_k is already optimal. Because there are only finitely many deterministic policies, the iteration terminates in finitely many steps — usually a handful, even for moderately sized MDPs.

A compact Python sketch:

pi = np.zeros(n_states, dtype=int)    # start with "always N"
while True:
    # 1. Evaluation: solve linear system for V^pi
    P_pi = P[np.arange(n_states), pi, :]            # (S, S)
    R_pi = (P[np.arange(n_states), pi, :]
            * R[np.arange(n_states), pi, :]).sum(axis=1)  # (S,)
    V = np.linalg.solve(np.eye(n_states) - gamma * P_pi, R_pi)
    # 2. Improvement: greedy w.r.t. V
    Q = (P * (R + gamma * V[None, None, :])).sum(axis=2)
    pi_new = Q.argmax(axis=1)
    if np.array_equal(pi_new, pi):
        break
    pi = pi_new

On the gridworld of §5.1, policy iteration converges in three or four outer iterations: an initial “always go north” policy is improved to something sensible in the first sweep, refined to the optimal policy in the second, and the third sweep confirms there is nothing left to improve. Each outer step solves a small linear system; each does much more work per step than a value-iteration sweep, but there are far fewer of them.

Two algorithms, one family

Value iteration and policy iteration look different but interpolate into each other. Modified policy iteration replaces the exact linear solve in policy evaluation with mm Bellman-expectation sweeps; at m=1m = 1 this is approximately value iteration, at m=m = \infty it is exact policy iteration. Asynchronous variants — Bertsekas (2017) §1.3 — update individual states rather than the whole vector at once, and converge under the same contraction guarantee as long as every state is visited infinitely often. Sutton & Barto (2018) §4.5 introduces these under the name generalized policy iteration and makes the point this section will end on: every learning algorithm in the rest of the book is some form of GPI. There is an evaluation step that pushes a value estimate toward consistency with the current policy, and an improvement step that pushes the policy toward greediness with respect to the current value. Q-learning, SAC, PPO, and the critic-and-actor objectives of every VLA paper that involves RL are all this pattern with different approximations and different truncations.

The choice between value and policy iteration in practice is a matter of constant factors. Value iteration is simpler to implement, has no inner loop, and is easy to vectorize. Policy iteration converges in many fewer outer steps but each step is expensive, especially if you do an exact linear solve. For small tabular problems the difference is invisible; for the function-approximation versions of Chapter 7, the two algorithms morph into the value-target and policy-target sides of an actor-critic loop, and choosing which to lean on is what gives TD3, SAC, and PPO their distinct flavors.

What dynamic programming gives, and what it cannot

Both algorithms share three properties that are worth pulling out. They are off-line: they need the full PP and RR in advance, not sampled experience. They are exact under their assumptions: at convergence, the policy is provably optimal, not approximately so. And they are exponential in state-space dimension: a 4×4 gridworld has 16 states, a 20×20×2020 \times 20 \times 20 block-stacking task on a discretized table has 80008\,000, and a robot arm with even a moderate discretization of joint angles has more states than there are atoms in something dramatic. Bellman (1957) coined the phrase “curse of dimensionality” for exactly this — every classical dynamic-programming algorithm scales polynomially in S|\mathcal{S}| and S|\mathcal{S}| scales exponentially in the number of state variables.

Two consequences shape the rest of the book. The first is that sampling replaces enumeration. Q-learning in §5.3 and the deep actor-critic methods of Chapter 7 estimate QQ^\star from sampled transitions, never writing PP down. The second is that function approximation replaces tabular storage: a neural network with a few hundred thousand parameters represents a VV or QQ that is defined over millions of states without ever materializing the table. The elegance is lost; the scaling is not. Every deep RL method we discuss from Chapter 7 onward is a way of running approximate value or policy iteration with neural networks standing in for the table, and the analytical guarantees that fall out of the contraction argument here are weakened or lost in the process. That is the price of leaving the tabular world.

Section 5.3 takes the first step in that direction: it drops the assumption that PP is known, keeps the tabular setting, and shows that a single sampled transition is enough to do approximate value iteration in the form known as Q-learning.

This section has been read times.

References

  1. Bellman (1957). Dynamic Programming. Princeton University Press.
  2. Howard (1960). Dynamic Programming and Markov Processes. MIT Press.
  3. Puterman (1994). Markov Decision Processes — Discrete Stochastic Dynamic Programming, Chapter 6. Wiley.
  4. Sutton & Barto (2018). Reinforcement Learning — An Introduction (2nd ed.), Chapter 4. MIT Press.
  5. Bertsekas (2017). Dynamic Programming and Optimal Control, Vol. I, 4th ed. Athena Scientific.