Decision Making I: MDPs and Value Iteration

A plan is a line through space; a policy is an answer for every state. This chapter derives the Bellman equation, proves value iteration converges, and shows why noise makes the optimal route longer.

The plan that slides off the floor

Here is the experiment, run twice on the same world, with the same solver and the same random numbers. On the left, Rusty commits to the optimal action sequence and replays it open-loop. On the right, Rusty consults the same solver's answer at every step.

At zero slip the two pictures are identical, and that identity is the whole reason deterministic planning ever worked: when execution is perfect, the sequence of actions along the optimal path is the optimal policy restricted to that path. Turn the slip up and the pictures diverge immediately, and they diverge in a specific, diagnostic way.

The plan does not fail gradually. It fails at the first divergence and then accumulates. Once the robot is one cell off the planned trajectory, the remaining actions in the sequence were computed for somebody else's position. Step seventeen of a plan is a claim about where you will be at step seventeen, and after a few unlucky draws it is a claim about a place you have never been.

The policy is unbothered. It never notices the "error", because the notion of error requires a reference trajectory and there isn't one. It reads the current cell and answers. The realized path is longer than the nominal one — noise costs something, and the cost is real — but the arrival rate barely moves. Run that world four thousand times with matched seeds and the split is brutal:

A five percent chance of veering — which the next section prices at about $13.5°$ of heading uncertainty, an ordinary afternoon for a localizer — already sinks the plan four times out of five. The policy pays for the same noise in steps rather than in failures, and steps are a currency you can budget.

The two objects are different types. This is worth being pedantic about, because the type is the lesson:

A plan is indexed by time. A policy is indexed by state. Everything else in this chapter follows from taking that difference seriously.

The obvious objection is that replanning fixes this: run the planner again from wherever you actually are. It does, and it is what most deployed systems do. But notice what replanning converges to as the replanning rate goes up — an answer for every state you might visit, computed on demand. A policy is what you get when you precompute all of them, and it is worth building once so you can see what the on-demand version is approximating. (When the state space is too big to precompute, you go back to on-demand: that is Chapter 23.)

Building intuition

Paint a reward, get a behavior

The following widget is the chapter in one object. A gridworld is laid over the Apartment floorplan; you paint payoff onto cells; value iteration re-converges live and the arrow field is read off whatever the value function currently is. Nothing is precomputed, and Rusty walks the current policy — including while it is still wrong.

Three things repay a minute of watching, and each one is a theorem later in the chapter.

Value propagates backward from the payoff, roughly a ring per sweep. Immediately after you paint a goal, only its neighbours know about it. That is why value iteration feels like a wave, and why the sweep count scales with the diameter of the map rather than its area. Look closely and the wave is lopsided — it runs further in one direction per sweep than the other — because the backups are done in place, so a cell late in the sweep already sees the cells updated before it. Flip the synchronous sweeps toggle and the lopsidedness disappears, along with some of the speed.

The arrows straighten out long before the numbers do. Watch the residual readout: the arrow field stops changing while $\lVert \Delta V \rVert_\infty$ is still large. That is not a rendering artifact, and it is not luck either — the loss bound explains why a value function nowhere near converged still yields a usable policy, and the measurement below shows this scene's policy going exactly optimal after 45 of its 72 sweeps. argmax only cares about the ordering of the $Q$-values, and orderings settle first.

Slip bends the field away from the hazard. At $s = 0$ the arrows run right past the spill cell — why not, you never veer. Push $s$ toward $0.3$ and a visible berth opens up. Nobody programmed a safety margin; the margin is what the expectation computes when some of the futures in it involve veering into the thing you are hugging. Measured on the default scene, the optimal policy's mean journey grows from $37.0$ steps at $s = 0$ to $45.5$ at $s = 0.2$ to $56.9$ at $s = 0.4$. The detour margin is the noise, priced.

The mathematics

Notation

Thrun et al. write the payoff as $c(x')$, collected on entering a state, and put the discount outside the integral. We fold the expected entry payoff into $r(x, u)$, which is the textbook form and changes not a single number:

The book's grid compiler does exactly this, which is why goal payoff shows up in the neighbours of the goal rather than in the goal cell itself.

The five pieces, and where each one came from

A Markov decision process is the tuple

with three standing assumptions, each of which is load-bearing and each of which is false somewhere.

1. Markov. $p(x' \mid x, u)$ depends on the past only through $x$. This is Chapter 5's assumption, and it is what licenses a policy to be a function of the current state alone. Drop it and the policy needs history.
2. Full observability. The robot knows $x$ exactly at decision time. This is the assumption this chapter buys its tractability with, and it is the one Chapter 22 pays back.
3. Stationarity. $p$ and $r$ do not depend on $t$. Together with an infinite horizon, this is what makes the optimal policy itself stationary — one lookup table, not one per timestep.

Only the first two pieces are new here. $\mathcal{X}$ is Chapter 13's occupancy grid with the walls inflated by the robot radius, exactly the configuration-space trick of Chapter 20. $p(x' \mid x, u)$ is Chapter 9's velocity motion model, discretized. The reward is the only genuinely new modelling decision, and it is where all the honesty lives: an MDP does not tell you what to want, it tells you what to do given what you want.

Where the slip parameter comes from

Gridworld papers announce a slip probability. We are going to earn ours, because the book's promise is that every probability has a pedigree.

The robot commands "move to the neighbouring cell in direction $d$". It drives for $\Delta t = \ell / v$ seconds, where $\ell$ is the cell size. The cell it lands in is decided by the bearing of its net displacement: inside the $\pm \pi/8$ sector around $d$ and it landed on the intended neighbour; outside, and it landed on one of the two flanking cells. So

where $(\Delta x, \Delta y)$ is drawn by the Chapter 9 sampler: perturb the commanded $(v, \omega)$ by $\hat{v} = v + \varepsilon_{\alpha_1 v^2}$ and $\hat{\omega} = \omega + \varepsilon_{\alpha_3 v^2}$, then integrate the arc for $\Delta t$, starting from a heading that is itself uncertain by $\sigma_\theta$ — because the robot aims using the heading it believes it has.

The one place the two rows disagree is instructive. At $\sigma_\theta = 5°$ the sampler reports $0.0005$ where the formula says $1.3 \times 10^{-5}$ — a factor of forty. The excess is not noise; it is $\Pr[\hat v < 0] = \Phi(-1/\sqrt{\alpha_1}) = 7.8 \times 10^{-4}$, the fraction of draws in which the Gaussian velocity model of Chapter 9 hands back a negative speed and Rusty reverses into the cell behind him. Below about six degrees of heading uncertainty, the dominant source of "slip" in this model is an artifact of describing a non-negative quantity with a Gaussian. Every noise model has a regime where its tails stop being a harmless approximation, and this is where the velocity model's is.

That table is the most important thing in the section, and it says something the gridworld literature usually hides. A well-localized robot barely slips. At five degrees of heading uncertainty, honest discretized slip is $1.3 \times 10^{-5}$. The slip parameter is not a property of the floor; it is overwhelmingly a property of the localizer. Discretized slip is a state-estimation error wearing a motion-model costume — which is a broad hint that the honest version of this chapter's problem is Chapter 22's.

The widgets let you set $s$ up to $0.4$ anyway. Read those settings as "a robot with a bad localizer, a wet floor, or 0.3 m cells that are too coarse for its dynamics" — all three are real, and the middle one is why the parameter is worth a slider.

Return, and why $\gamma$ buys convergence

Fix a policy $\pi$ and run it forever. The return is the discounted sum of payoffs,

Discounting does two jobs, and it is worth separating them because only one is mathematical.

The mathematical job: if $|r| \le r_{\max}$ then $|V^\pi| \le r_{\max}/(1-\gamma)$, so an infinite sum of infinitely many terms is a finite number and the whole theory has something to be about. Without it, "reach the dock eventually" and "reach the dock in nine steps" both score $-\infty$ and $\arg\max$ has nothing to chew on.

The modelling job: $\gamma$ is a statement about how far ahead the robot cares. The effective horizon is $1/(1-\gamma)$ steps, in the sense that payoff beyond it is discounted below $1/e$:

For a 0.3 m grid, $\gamma = 0.98$ means the robot can see about 15 metres of consequence. Set $\gamma = 0.9$ and the charging dock becomes invisible from the far bedroom — the value function goes flat out there, the arrows go arbitrary, and the robot wanders. That failure is not a bug in the solver; it is the solver correctly reporting that under your stated preferences, the dock is not worth walking to. You can watch it happen with the Policy Painter's $\gamma$ slider.

The Bellman optimality equation

Define the optimal value function $V^*(x) = \max_\pi V^\pi(x)$ and the state–action value

— the value of committing to $u$ once and behaving optimally forever after. Then the whole of dynamic programming is the observation that $V^*$ is the pointwise max of $Q^*$:

Read it as a sentence: the best you can do from here is the best over first moves of (what that move pays now) plus (discounted, averaged over where it might land you, the best you can do from there). The colors match the widget below: green is the immediate payoff, blue is the discounted future, purple is the answer.

That widget exists because the equation above looks like linear algebra and is not. A backup is clerical work: read a handful of neighbours, weight them by how likely you are to reach them, scale by $\gamma$, add the payoff, keep the biggest. Step it a few times and watch the value wave crawl backwards from the goal and bend through the doorway — the same wave Chapter 20's wave-front planner produced, which is not a coincidence and is proved below.

Value iteration converges, and here is the stopping rule

Write the Bellman equation as an operator on value functions:

Value iteration is $V_{k+1} = \mathcal{T} V_k$, starting from $V_0 = 0$. The theorem that makes it an algorithm rather than a hope is that $\mathcal{T}$ is a contraction.

A half-converged value function still gives a good policy

Greedy extraction is the reason value iteration is useful at all:

Two readings of that bound, and they point in opposite directions.

Pessimistic: the amplification factor $2\gamma/(1-\gamma)$ is $98$ at $\gamma = 0.98$. A value function good to $0.1$ certifies only a policy within $9.8$ of optimal — nearly vacuous.

Optimistic, and what actually happens: the bound is worst-case over adversarial MDPs. In a gridworld, $\arg\max$ depends only on the order of the $Q$-values at each state, and orderings stabilize long before magnitudes do. Measured on the Apartment scene, the greedy policy becomes exactly optimal — identical to $\pi^*$ in all 925 free cells — after 45 of the 72 sweeps value iteration eventually takes. The residual at that moment is $2.07$, so the theorem certifies only $\lVert V^\pi - V^* \rVert_\infty \le 101$, on a problem whose optimal value at the start is $7.9$. The bound is not wrong; it is simply not the thing that happens.

That gap is what real-time dynamic programming (Barto, Bradtke and Singh, 1995) is built to exploit: you may act on a value function you have no right to trust, as long as you keep backing it up along the states you actually visit. It is also why the Policy Painter can run the robot while the wave is still crossing the map.

Policy iteration: solve, improve, repeat

Value iteration nudges every state a little on every sweep. Policy iteration takes the other extreme: pick a policy, evaluate it exactly, then improve it everywhere at once.

The evaluation step is where the linear algebra finally shows up. Fix $\pi$; there is no max any more, so the Bellman equation for $V^\pi$ is a linear system:

with $P_\pi$ the $|\mathcal{X}| \times |\mathcal{X}|$ row-stochastic matrix $[P_\pi]_{x x'} = p(x' \mid x, \pi(x))$ — sparse, with at most three nonzeros per row in our gridworld. That is the same shape of solve as the sparse Cholesky in Chapter 15, and the Rust implementation hands it to the same crate.

Stochastic shortest paths: what navigation actually is

Discounting is a strange thing to want from a delivery robot. Nobody prefers a package delivered now over the same package delivered in a minute by a factor of $0.98^{60}$; we just want it delivered, quickly, and the discount was a mathematical convenience. The formulation that says what we mean is the stochastic shortest path (SSP):

• $\gamma = 1$ — no discounting;
• one or more absorbing goal states with zero payoff;
• $r(x, u) = -1$ per step (or $-\text{cost}(u)$), so every action strictly hurts.

Then $-V^*(x)$ is literally the minimum expected number of steps to the goal, which is a quantity you can hold in your head. But the contraction argument is gone: $\gamma = 1$ makes the modulus $1$, and there is no Banach theorem to invoke. What replaces it is a condition on the problem rather than on the discount.

Both conditions are checkable, and both fail in ways you will meet. Condition (i) fails if the inflated map has a walled-off room: the goal is unreachable, $V^* = -\infty$ there, and a solver that does not detect it will happily return $-10^{12}$ and an arrow field of noise. Condition (ii) fails if you give the robot a free action — a zero-cost stay — because then "stand still forever" is an improper policy with finite cost $0$, and it is optimal. This is not a hypothetical; it is the single most common bug in hand-rolled SSP solvers, and it is why the library's stay action still charges stepCost.

The wave-front planner, unmasked

Now set the slip to zero. Each action has exactly one successor, $x' = f(x, u)$, so the sum over $x'$ collapses to a single term and the SSP Bellman equation becomes

That is the Bellman–Ford relaxation. Sweep it synchronously and you have Bellman–Ford. Sweep it in-place in order of increasing $-V^*$ — always expanding the unfinished state closest to the goal — and you have Dijkstra. Fill it outward from the goal one ring at a time and you have Chapter 20's wave-front planner, exactly.

So the wave-front planner is not analogous to value iteration; it is value iteration, on a deterministic MDP with unit costs, with a particularly clever sweep order. Everything this chapter adds is what happens when the arrow out of a cell is no longer a promise. And the "potential field with no local minima" that Chapter 20 wanted is now definitional: $V^*$ has no spurious local optimum because it is not a field somebody sculpted, it is the expected cost-to-go, and a state whose neighbours are all worse than it would be violating its own Bellman equation.

How much detour is noise worth?

We have said twice now that noise makes the optimal route longer. Here is the quantitative version, in the smallest world that can carry it: two corridors from start to goal, one short and flanked by a drop, one long and flanked by walls.

Before touching the widget, commit to a guess. The ledge is $6$ cells; the detour is $16$ — nearly three times as long. A veer on the ledge costs a penalty of $6$ and dumps the robot back at the start. At what slip $s$ does the optimal policy abandon the ledge?

Write your number down, then read the derivation.

Six and a bit percent. Most readers guess somewhere between $0.15$ and $0.3$, because the detour is so much longer. The values say otherwise:

The ledge is exponentially fragile and the detour is merely linearly annoying, so they cross early. And notice what $s^* = 0.067$ corresponds to in the units of the previous section: about $15°$ of heading uncertainty. The decision of whether to take the short route through the narrow gap is being made, in effect, by the localizer.

The algorithms

This is Thrun et al.'s Table 15.1 with two additions: the stopping rule from the contraction proof (line 9), and the explicit policy extraction (line 10). One subtlety hides in line 6. If $\hat V$ is a single array that you write into as you go, later states in the sweep see the updated values of earlier ones — that is the Gauss–Seidel or asynchronous variant, and it is what the library does by default. If you write into a fresh array, it is the synchronous Jacobi variant. Thrun's draft is explicit that this does not affect whether value iteration converges, only how fast. On the Apartment, Gauss–Seidel needs 72 sweeps where Jacobi needs 84 — and if you order the sweep by breadth-first distance from the goal, so that every backup reads neighbours that were updated moments ago, it drops to 33. Same algorithm, same fixed point, less than half the work, purely from the order of a for loop.

The priority key on line 6 is an upper bound on how much this backup could possibly move that predecessor — so the heap orders states by potential change without paying for a trial backup. Moore and Atkeson (1993) introduced this for learned models; it is just as useful for known ones.

A worked example you can check by hand

Four states in a line: $A$, $B$, $C$, and the goal $G$, which absorbs. Rusty has two actions.

• roll — nudge forward. Advances one cell with probability $0.8$; with probability $0.2$ the wheels spin and nothing happens. Costs $1$.
• lunge — dump enough current into the motors to guarantee the cell change. Costs $2$.

Undiscounted, $\gamma = 1$: a stochastic shortest path. Every number below is reproduced by a Rust unit test and by self-check 1 in lib/decision/__checks_ch21__.ts.

The fixed point, by hand

Work backwards from the goal. At $C$, roll gives $-1 + 0.8 \cdot V(G) + 0.2 \cdot V(C) = -1 + 0.2\,V(C)$, and if roll is optimal there, that equals $V(C)$:

which is just $1/0.8$: the expected number of attempts to make one cell of progress. Each cell costs the same, so

Is roll really optimal? Check the alternative at each state: lunge from $C$ scores $-2$ against $-1.25$; from $B$, $-2 + (-1.25) = -3.25$ against $-2.5$; from $A$, $-4.5$ against $-3.75$. Buying determinism at a price of $2$ is a bad deal when the stochastic option costs $1.25$ in expectation. The max gate has something to do at every state, and it always makes the same choice.

The first three sweeps

Now run the algorithm synchronously from $V_0 = 0$ and watch it get there. Every entry is one line of arithmetic:

Check sweep 3 at $B$ yourself: $-1 + 0.8 \cdot V_2(C) + 0.2 \cdot V_2(B) = -1 + 0.8(-1.2) + 0.2(-2) = -2.36$. And notice the shape of the convergence. After three sweeps $C$ is 99% of the way to its answer and $A$ is only 80%; after four, $C$ is done to three decimals and $A$ still is not. Information flows backward from the goal at one cell per synchronous sweep, which is the whole reason sweep order matters.

Value iteration needs 24 sweeps to reach $10^{-12}$ here. Policy iteration needs one improvement: evaluate the initial all-roll policy exactly, discover it is already greedy, stop.

Implementation in Rust

The module is crates/ch21_mdp. The design constraint that shapes everything: a gridworld MDP has tens of thousands of states and at most three successors per state-action, so the transition model must be sparse and the inner loop must not allocate.

The model

use nalgebra::DVector;

/// One successor and its probability. `p` is f32 because a transition row is
/// read far more often than it is written, and halving the row halves the cache
/// misses in the inner loop — the only place this crate spends time.
#[derive(Clone, Copy, Debug)]
pub struct Transition {
    pub s: u32,
    pub p: f32,
}

/// A sparse row of p(· | x, u). Probabilities sum to 1 up to f32 epsilon.
pub type SparseDist = Vec<Transition>;

/// A finite MDP with a compile-time action count.
///
/// `A` is a const generic because the action set of a gridworld is fixed by its
/// connectivity (4, 8, or 9 with `stay`), and pinning it lets the per-state
/// arrays live inline instead of behind a second indirection.
pub struct Mdp<const A: usize> {
    pub n_states: usize,
    /// Indexed [state][action].
    pub trans: Vec<[SparseDist; A]>,
    /// r(x, u): the expected immediate payoff, with entry payoff folded in.
    pub reward: Vec<[f64; A]>,
    /// γ ∈ (0, 1]. γ = 1 is legal only when `absorbing` is non-empty (SSP).
    pub gamma: f64,
    /// V(x) ≡ 0 here, and the episode stops.
    pub absorbing: Vec<bool>,
    pub action_labels: [&'static str; A],
}

impl<const A: usize> Mdp<A> {
    /// Q(x, u) = r(x, u) + γ Σ_{x'} p(x' | x, u) V(x').
    #[inline]
    pub fn q(&self, v: &[f64], x: usize, u: usize) -> f64 {
        if self.absorbing[x] {
            return 0.0;
        }
        let mut acc = 0.0;
        for t in &self.trans[x][u] {
            acc += f64::from(t.p) * v[t.s as usize];
        }
        self.reward[x][u] + self.gamma * acc
    }

    /// The max gate: (TV)(x) and the action that attains it.
    #[inline]
    pub fn backup(&self, v: &[f64], x: usize) -> (f64, u8) {
        if self.absorbing[x] {
            return (0.0, 0);
        }
        let mut best = f64::NEG_INFINITY;
        let mut arg = 0u8;
        for u in 0..A {
            let q = self.q(v, x, u);
            if q > best {
                best = q;
                arg = u as u8;
            }
        }
        (best, arg)
    }

    /// ‖TV − V‖∞ over non-absorbing states. Zero exactly at the fixed point,
    /// which makes it the only honest convergence certificate.
    pub fn max_residual(&self, v: &DVector<f64>) -> f64 {
        (0..self.n_states)
            .filter(|&x| !self.absorbing[x])
            .map(|x| (self.backup(v.as_slice(), x).0 - v[x]).abs())
            .fold(0.0, f64::max)
    }
}

Value iteration, with the stopping rule that the proof earned

use nalgebra::DVector;
use crate::mdp::Mdp;

pub struct ViResult {
    pub v: DVector<f64>,
    pub policy: Vec<u8>,
    pub sweeps: usize,
    /// ‖V_k − V_{k−1}‖∞ after each sweep — the widgets' convergence wave.
    pub residuals: Vec<f64>,
    pub converged: bool,
}

/// Stop when one sweep moves V by less than this, and ‖V − V*‖∞ ≤ ε is
/// guaranteed. From ‖V_k − V*‖ ≤ γ/(1−γ)·‖V_k − V_{k−1}‖.
///
/// γ = 1 (a stochastic shortest path) has no such bound: the contraction
/// modulus is 1 in the sup norm, so we fall back to the raw residual and the
/// caller is on the hook for a proper policy existing.
pub fn stopping_threshold(gamma: f64, eps: f64) -> f64 {
    if gamma >= 1.0 { eps } else { eps * (1.0 - gamma) / gamma }
}

/// One in-place Gauss–Seidel sweep. Returns ‖ΔV‖∞.
///
/// In place on purpose: a backup late in the sweep sees the states updated
/// earlier in it, so with a good ordering information crosses the whole map in
/// a single pass instead of one cell per sweep.
pub fn sweep_in_place<const A: usize>(mdp: &Mdp<A>, v: &mut DVector<f64>) -> f64 {
    let mut residual: f64 = 0.0;
    for x in 0..mdp.n_states {
        if mdp.absorbing[x] {
            v[x] = 0.0;
            continue;
        }
        let next = mdp.backup(v.as_slice(), x).0;
        residual = residual.max((next - v[x]).abs());
        v[x] = next;
    }
    residual
}

/// Thrun et al., Table 15.1 — `MDP_value_iteration`.
pub fn value_iteration<const A: usize>(mdp: &Mdp<A>, eps: f64, max_sweeps: usize) -> ViResult {
    let mut v = DVector::zeros(mdp.n_states);
    let threshold = stopping_threshold(mdp.gamma, eps);
    let mut residuals = Vec::new();

    let converged = loop {
        if residuals.len() >= max_sweeps {
            break false;
        }
        let r = sweep_in_place(mdp, &mut v);
        residuals.push(r);
        if r < threshold {
            break true;
        }
    };

    ViResult { policy: greedy_policy(mdp, &v), sweeps: residuals.len(), v, residuals, converged }
}

/// π(x) = argmax_u Q(x, u). Optimal once V = V*; within 2γε/(1−γ) before that.
pub fn greedy_policy<const A: usize>(mdp: &Mdp<A>, v: &DVector<f64>) -> Vec<u8> {
    (0..mdp.n_states).map(|x| mdp.backup(v.as_slice(), x).1).collect()
}

Policy evaluation as a sparse solve

The exact evaluation step is where faer earns its place. The matrix $I - \gamma P_\pi$ is sparse, nonsymmetric, and — for $\gamma < 1$ — strictly diagonally dominant by rows, which means it is nonsingular and an LU with partial pivoting is stable without any reordering heroics.

use faer::sparse::{SparseColMat, Triplet};
use faer::prelude::*;
use nalgebra::DVector;
use crate::mdp::Mdp;
use crate::vi::greedy_policy;

/// Solve (I − γ P_π) V^π = r_π exactly.
///
/// Same shape of problem as the normal equations in Chapter 15, and the same
/// crate solves it — the difference is that here the matrix is nonsymmetric, so
/// it is LU rather than Cholesky.
pub fn policy_evaluation<const A: usize>(mdp: &Mdp<A>, pi: &[u8]) -> DVector<f64> {
    let n = mdp.n_states;
    let mut triplets: Vec<Triplet<usize, usize, f64>> = Vec::with_capacity(4 * n);
    let mut rhs = Mat::<f64>::zeros(n, 1);

    for x in 0..n {
        if mdp.absorbing[x] {
            // V(x) ≡ 0: a 1×1 identity row keeps the system square and the
            // absorbing convention explicit rather than implied.
            triplets.push(Triplet::new(x, x, 1.0));
            continue;
        }
        let u = pi[x] as usize;
        triplets.push(Triplet::new(x, x, 1.0));
        for t in &mdp.trans[x][u] {
            // Accumulating duplicates is exactly right here: a self-loop lands
            // on the diagonal and must be subtracted from the identity.
            triplets.push(Triplet::new(x, t.s as usize, -mdp.gamma * f64::from(t.p)));
        }
        rhs[(x, 0)] = mdp.reward[x][u];
    }

    let a = SparseColMat::try_new_from_triplets(n, n, &triplets)
        .expect("policy transition matrix is well formed by construction");
    let v = a.sp_lu().expect("I − γP_π is nonsingular for γ < 1").solve(&rhs);
    DVector::from_iterator(n, (0..n).map(|i| v[(i, 0)]))
}

/// Howard's policy iteration: evaluate exactly, improve greedily, repeat.
pub fn policy_iteration<const A: usize>(mdp: &Mdp<A>, max_iter: usize) -> (DVector<f64>, Vec<u8>, usize) {
    let mut pi = vec![0u8; mdp.n_states];
    let mut v = DVector::zeros(mdp.n_states);

    for it in 0..max_iter {
        v = policy_evaluation(mdp, &pi);
        let next = greedy_policy(mdp, &v);
        // Switch only on a strict improvement. Ties between equally good
        // actions would otherwise make the loop oscillate forever.
        let mut changed = 0usize;
        let mut merged = pi.clone();
        for x in 0..mdp.n_states {
            let (a, b) = (next[x] as usize, pi[x] as usize);
            if a != b && mdp.q(v.as_slice(), x, a) > mdp.q(v.as_slice(), x, b) + 1e-12 {
                merged[x] = next[x];
                changed += 1;
            }
        }
        pi = merged;
        if changed == 0 {
            return (v, pi, it + 1);
        }
    }
    (v, pi, max_iter)
}