Value Iteration on a Small MDP

Problem Statement

Run value iteration to convergence on a finite MDP and return both the optimal value function and the optimal deterministic policy.

Background

Policy iteration alternates a full evaluation with a greedy improvement. The evaluation is the expensive half, and value iteration collapses it away: it turns out you can truncate policy evaluation after a single sweep, provided that sweep uses the max instead of the policy average.

$v_{k+1}(s) \leftarrow \max_a \Big[ r(s,a) + \gamma \sum_{s'} p(s' \mid s,a)\, v_k(s') \Big]$

That is one update, doing evaluation and improvement at the same time. It is the Bellman optimality operator, a $\gamma$-contraction in the max norm, so it converges to the unique fixed point $v_*$ from any start.

The policy is not maintained during the loop. It is extracted once at the end, greedily with respect to the converged values. This is the sharpest illustration of generalized policy iteration: evaluation and improvement do not have to complete, or even be separate steps, for the pair to converge — they only have to keep pushing toward consistency and greediness.

A practical note: the values converge much faster than the greedy policy changes, so the extracted policy is usually optimal long before delta gets small. Stopping early therefore costs you accurate values, rarely a good policy.

Your Task

Implement:

def value_iteration(P, R, gamma, theta=1e-10, max_sweeps=100000):
    ...

• Initialise v to all zeros.
• Each sweep is synchronous: compute the entire new table from the old one, then delta = max(abs(new - old)); stop when delta < theta.
• After convergence, extract the greedy policy, breaking ties toward the lowest action index.

Return a tuple (v, policy) — a list of floats and a list of ints.

Input / Output Format

Nested lists of floats in. A (list_of_floats, list_of_ints) tuple out; the grader rounds the values to 4 decimals.

Sample

P = [[[0.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 1.0]]]
R = [[1.0, 0.0], [0.0, 0.0]]
v, policy = value_iteration(P, R, 0.9)
print([round(x, 4) for x in v], policy)

Output:

[1.0, 0.0] [0, 0]

State 1 is absorbing with zero reward, so v(1) = 0. In state 0, action 0 earns 1 and moves to state 1 while action 1 earns 0 and stays put, so v(0) = 1 and the optimal action is 0.