The Bellman equations proved: contraction, value iteration and policy iteration

The Bellman equations are the mathematical heart of reinforcement learning: they turn the definition of value, an average over infinitely many possible futures, into an equation that links the value of a state to that of the next states. All algorithms come from them, from exact ones (value and policy iteration) to those that learn from experience (Q-learning, DQN, actor-critic). This lesson derives them, proves that they have one and only one solution and that the algorithms converge, and applies them to the gridworld of the previous lesson with Python code.

The Bellman equation for a policy

Theorem. For every policy and every state :

Proof. Start from the recursive relation of the return, , and take the average conditioned on :

The first term is the average reward: is chosen with probability and on average is received. In the second one also conditions on the next state: by the Markov property, once is known, the future does not depend on how it was reached, so . Weighting with the probabilities gives the equation. ∎

It is a linear system of equations in unknowns: in vector form , with solution (the matrix is invertible for , because has norm less than 1). For the gridworld with the random policy this gives the values −14, −20, −22 of the previous lesson.

The optimality equation

For the optimal policy the average over actions becomes a maximum:

It is no longer linear (because of the maximum), so it is not solved by inverting a matrix: it is solved by iteration. The result that makes this possible is the following.

The Bellman contraction

Define the optimal Bellman operator , which turns a function into another: . The optimality equation says that is a fixed point: .

Theorem. is a contraction of factor in the max norm: for every , .

Proof. An elementary inequality is used: for two functions of , (if the maximum of is at , , and symmetrically). Then, for every :

because the probabilities sum to 1. ∎

Consequences (Banach fixed-point theorem):

• the fixed point exists and is unique: is well defined;
• starting from any , the sequence converges to , and the error shrinks at least by the factor at each step: .

This is value iteration. On the gridworld with and a deliberately wrong initial estimate (50 in every cell) the maximum error goes from 46.7 to 41.3, 36.5, 32.1, …: ratios between 0.87 and 0.88, always below 0.9 as guaranteed by the theorem.

Policy improvement

Theorem (policy improvement). Given a policy , let be the greedy policy with respect to : . Then for every .

Proof. By construction . Expand taking at the first step, and apply the same inequality to :

repeating to infinity (the residual terms tend to zero). ∎

Alternating evaluation (computing with the linear equation) and improvement (making the policy greedy) gives policy iteration: every step improves or leaves the policy unchanged, deterministic policies are finite in number, so in a finite number of steps one reaches a policy that no longer improves, i.e. that satisfies the optimality equation. This also proves that a deterministic optimal policy always exists. The “evaluate, then improve” scheme returns in all modern algorithms: in actor-critic the critic evaluates and the actor improves.

The code

import numpy as np

N = 4                                         # 4x4 gridworld, exits at 0 and 15, -1 per step
MOVES = [(-1, 0), (1, 0), (0, -1), (0, 1)]    # up, down, left, right
EXITS = {0, 15}

def step(s, a):
    if s in EXITS:
        return s, 0.0
    r, c = divmod(s, N)
    r = min(max(r + MOVES[a][0], 0), N - 1)
    c = min(max(c + MOVES[a][1], 0), N - 1)
    return r * N + c, -1.0

def bellman_optimal(V, gamma):
    """Optimal Bellman operator T: (TV)(s) = max_a [ r + gamma V(s') ]."""
    return np.array([0.0 if s in EXITS else
                     max(r + gamma * V[s2] for s2, r in (step(s, a) for a in range(4)))
                     for s in range(N * N)])

# value iteration with gamma = 1
V = np.zeros(N * N)
for k in range(100):
    V_new = bellman_optimal(V, 1.0)
    if np.max(abs(V_new - V)) < 1e-10:
        break
    V = V_new
print("V* after", k, "sweeps:\n", V.reshape(N, N))

# evaluation of the random policy: average over the four actions instead of the maximum
Vr = np.zeros(N * N)
for _ in range(5000):
    Vr = np.array([0.0 if s in EXITS else
                   np.mean([r + Vr[s2] for s2, r in (step(s, a) for a in range(4))])
                   for s in range(N * N)])
print("random policy:\n", np.round(Vr.reshape(N, N), 1))

# greedy policy with respect to V*: the best arrow in each cell
ARROWS = "↑↓←→"
pol = [" " if s in EXITS else ARROWS[int(np.argmax([r + V[s2] for s2, r in (step(s, a) for a in range(4))]))]
       for s in range(N * N)]
print("\n".join(" ".join(pol[i * N:(i + 1) * N]) for i in range(N)))

# contraction: with gamma = 0.9 the error shrinks at least by 0.9 per sweep
Vs = np.zeros(N * N)
for _ in range(200):
    Vs = bellman_optimal(Vs, 0.9)
U = np.where(np.isin(np.arange(N * N), list(EXITS)), 0.0, 50.0)
errors = []
for _ in range(6):
    U = bellman_optimal(U, 0.9)
    errors.append(np.max(abs(U - Vs)))
print("errors:", np.round(errors, 2), "ratios:", np.round(np.array(errors[1:]) / np.array(errors[:-1]), 3))

Running it prints after 3 sweeps plus the checking one (the values 0, −1, −2, −3), the values of the random policy (−14, −18, −20, −22), the grid of arrows towards the nearest exit and the errors 46.71, 41.31, 36.45, … with ratios 0.884, 0.882, 0.880, … On large problems the table does not fit in memory and is unknown: the next lessons replace the sum over with experience (Q-learning) and the table with a neural network (DQN).