Markov and Partially Observable Decision Processes › Implementation: Planning a Gridworld with Value Iteration
World Models HandbookAll chapters Markov and Partially Observable Decision Processes › Implementation: Planning a Gridworld with Value IterationOne section · Sections in this chapter8 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We now make the first half of the chapter executable. Our running example is a office gridworld: the delivery robot starts at the bottom-left corner, a charging goal sits at the top-right, and an open stairwell (a pit) sits at the center. Each move succeeds with probability 0.8 and slips to a perpendicular direction with probability 0.1 each, simulating the slippery floor. Reaching the goal earns 20, falling into the pit costs 20, and every other step costs 1, encouraging short routes.
First, the imports and the environment model. We represent the transition kernel and reward function as dense arrays, indexed as T[action, state, next_state].
import matplotlib.pyplot as plt
import numpy as np
n = 5 # 5x5 grid
S = n * n
idx = lambda r, c: r * n + c
goal, pit, start = idx(0, 4), idx(2, 2), idx(4, 0)
moves = [(-1, 0), (1, 0), (0, -1), (0, 1)] # up, down, left, right
perp = {0: [2, 3], 1: [2, 3], 2: [0, 1], 3: [0, 1]}
slip = 0.1 # probability of each perpendicular slip
T = np.zeros((4, S, S)) # T[a, s, s']
Rmat = np.full((4, S, S), -1.0) # Rmat[a, s, s'], step cost by default
for a, (dr, dc) in enumerate(moves):
for r in range(n):
for c in range(n):
s = idx(r, c)
if s in (goal, pit): # terminals absorb with no further reward
T[a, s, s] = 1.0
Rmat[a, s, s] = 0.0
continue
for b, (br, bc) in enumerate(moves):
if b == a:
p = 1.0 - 2 * slip
elif b in perp[a]:
p = slip
else:
continue
nr = min(max(r + br, 0), n - 1) # bumping a wall: stay put
nc = min(max(c + bc, 0), n - 1)
ns = idx(nr, nc)
T[a, s, ns] += p
Rmat[a, s, ns] = (
20.0 if ns == goal else (-20.0 if ns == pit else -1.0)
)
assert np.allclose(T.sum(axis=2), 1.0)
for terminal in (goal, pit):
assert np.allclose(T[:, terminal, terminal], 1.0)
assert np.allclose(Rmat[:, terminal, terminal], 0.0)The arrays encode everything the planner needs: dynamics, slip stochasticity, wall clamping, and rewards. Note that the reward is attached to transitions, so stepping into the goal collects 20 exactly once, and the absorbing self-loops at terminals carry zero reward to keep values finite. The += accumulation matters in the corners: an attempted move off the grid clamps back to the same cell, so several probability masses can land on a single successor and must be summed rather than overwritten. At the top-left corner under a leftward action, for example, the intended move contributes to the clamped self-transition and the upward perpendicular slip contributes another , for total self-transition mass .
Next, value iteration. We record the maximum change per sweep so we can inspect the contraction behavior predicted by the theory.
gamma = 0.9
V = np.zeros(S)
errors = []
for k in range(500):
Q = (T * Rmat).sum(axis=2) + gamma * (T @ V) # Q[a, s]
V_new = Q.max(axis=0)
errors.append(np.max(np.abs(V_new - V)))
V = V_new
if errors[-1] < 1e-12:
break
Q = (T * Rmat).sum(axis=2) + gamma * (T @ V)
greedy_policy = Q.argmax(axis=0)Converged in 53 sweeps, final change 5.87e-13 V*(start) = 1.440 V*(goal) = 0.000, V*(pit) = 0.000
The contraction analysis supplies a conservative sweep bound: with , successive Bellman-backup differences can shrink by at most a factor of 0.9. The first recorded change is , attained by a goal-adjacent state whose best action enters the goal, so the envelope first crosses the threshold at the zero-based envelope index , corresponding to the 290th recorded sweep, or roughly three hundred sweeps. Executing the cell above reaches the stopping tolerance in 53 sweeps. The gap is expected: is a worst-case contraction factor, while this grid's transition structure yields faster realized shrinkage. There is a compact way to see the backup arithmetic in the code itself: the line computing Q first averages immediate expected rewards over successors, then adds the discounted value of successors, so a single matrix product performs one full Bellman backup for every state and action simultaneously. Let us compare the observed decline with the theoretical envelope.

The observed curve equals the reference envelope at the initial sweep and stays strictly below it afterward, reaching the tolerance much sooner than the worst-case bound requires. Its near-straight descent on the logarithmic axis indicates geometric shrinkage, while its steeper slope shows that the realized rate is faster than . The Bellman operator's sup-norm Lipschitz constant is , so each successive difference contracts by at most that factor; an individual sweep need not attain the bound. This distinction matters: contraction theory guarantees convergence for the exact discounted Bellman operator, but it does not say every problem attains the worst-case rate. If a practical planner fails to converge, inspect whether its update is actually the same contraction mapping—for example, approximation, changing models, or optimization error can invalidate that premise.
Now the interesting output: the optimal value surface and the policy it induces.

Read the figure as a decision-making document. Color encodes : bright near the goal, dark near the pit, with a visible valley of low value surrounding the hazard. The arrows use only up and right in this seeded grid, but their routing bends around the pit rather than taking the most direct diagonal corridor, because the greedy policy weighs the 0.2 total slip probability against a 20-unit penalty. This is exactly the kind of behavior that emerges from the Bellman backup without ever being programmed in: risk-aware routing is a consequence of expected-value arithmetic, not a hand-coded rule.
Finally, let us close the loop between theory and simulation. If really is the expected return of the greedy policy, then Monte Carlo rollouts of that policy should reproduce up to sampling noise. This is the chapter's first taste of validation practice: never trust a planned value until simulated experience confirms it.
def rollout(policy_fn, rng, max_steps=100):
s, G, disc = start, 0.0, 1.0
for _ in range(max_steps):
a = policy_fn(s)
s_next = rng.choice(S, p=T[a, s])
G += disc * Rmat[a, s, s_next]
disc *= gamma
s = s_next
if s in (goal, pit):
break
return G
rng = np.random.default_rng(0)
mc_returns = np.array(
[rollout(lambda s: int(greedy_policy[s]), rng) for _ in range(3000)]
)Monte Carlo mean return : 1.425 +/- 0.064 (standard error) Planned value V*(start) : 1.440
The two numbers agree closely, which is the entire point of the exercise. The value function computed by iterating the Bellman operator is not an abstract score; it is a verifiable prediction about closed-loop performance. Discrepancies here would indicate a bug, a modeling mismatch between rollout and planner, or statistical noise, and telling those apart is a skill in its own right. When a learned world model's planned values diverge from realized returns in later chapters, this gridworld check is the minimal diagnostic you should reach for.
Beyond matching means, the shape of the return distribution itself carries diagnostic information. The next plot separates the rollouts into those that reached the goal and those that fell into the pit.
def rollout_with_outcome(policy_fn, rng, max_steps=100):
s, G, disc = start, 0.0, 1.0
for _ in range(max_steps):
a = policy_fn(s)
s_next = rng.choice(S, p=T[a, s])
G += disc * Rmat[a, s, s_next]
disc *= gamma
s = s_next
if s in (goal, pit):
break
return G, s
rng_dist = np.random.default_rng(42) # fresh seed so this run is reproducible
returns_g, returns_p, returns_timeout = [], [], []
for _ in range(4000):
G, final = rollout_with_outcome(lambda s: int(greedy_policy[s]), rng_dist)
if final == goal:
returns_g.append(G)
elif final == pit:
returns_p.append(G)
else:
returns_timeout.append(G)
planned_start_value = V[start]
assert len(returns_g) > 0 and len(returns_p) > 0
assert len(returns_timeout) == 0
assert len(returns_g) + len(returns_p) == 4000
assert np.isclose(planned_start_value, 1.440, atol=5e-4)
Two distinct lobes appear instead of a symmetric spread around the mean: a tall lobe of successful deliveries whose returns spread across the planned value, with unobstructed routes landing above it and early slip-induced detours landing below, and a thin but costly lobe of pit outcomes far below zero. All 4,000 seeded rollouts reached one of the two terminals within the 100-step cap; the code asserts that no timeout was silently discarded. The empirical mean differs slightly from the planned value because of finite-sample noise. The value-iteration approximation contributes at most about here, negligible beside the Monte Carlo standard error. Within the two-lobe empirical decomposition, increasing the pre-change pit probability from to while holding the conditional return distributions fixed would shift the population mean by times the gap between their conditional means. That is why the tail should be inspected separately from the average.