Iterative Policy Evaluation

Problem Statement

Compute $v_\pi$ by sweeping the Bellman expectation backup until the values stop moving.

Background

Solving $(I - \gamma P^\pi) v = R^\pi$ directly costs $O(n^3)$. Iterative policy evaluation instead turns the Bellman expectation equation into an assignment and applies it over and over:

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

Each sweep costs $O(n^2 |\mathcal{A}|)$ and the operator is a $\gamma$-contraction, so $v_k \to v_\pi$ from any initialisation. Starting from all zeros is conventional.

When to stop is the part that needs a decision. The standard rule tracks the largest change any state made during a sweep,

$\Delta = \max_s |v_{k+1}(s) - v_k(s)|$

and halts once $\Delta < \theta$. Note what this is: a measure of how much the estimate is still moving, not of how far it is from the truth. The two are related — the remaining error is bounded by roughly $\gamma\Delta/(1-\gamma)$ — but as gamma approaches 1 that bound blows up, and a small $\Delta$ stops meaning a small error. This problem uses a synchronous sweep: every new value is computed from the previous sweep's table.

Your Task

Implement:

def policy_evaluation(pi, P, R, gamma, theta=1e-8, max_sweeps=100000):
    ...

• pi[s][a], P[s][a][s2], R[s][a] as in the earlier problems.
• Initialise v to all zeros.
• Each sweep computes a whole new table from the old one, then delta = max(abs(new - old)).
• Stop as soon as delta < theta, or after max_sweeps sweeps.

Return the final v as a list of floats.

Input / Output Format

Nested lists of floats in, a list of floats out, rounded to 4 decimals by the grader.

Sample

pi = [[1.0, 0.0], [1.0, 0.0]]
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]]
print([round(x, 4) for x in policy_evaluation(pi, P, R, 0.9)])

Output:

[1.0, 0.0]

The policy always takes action 0. From state 0 it collects reward 1 and moves to state 1, which is absorbing with zero reward, so v(1) = 0 and v(0) = 1 + 0.9*0 = 1.