Iterative Policy Evaluation

mediumFREEDynamic ProgrammingReinforcement LearningRL 1: Foundations

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.

Example

Input

policy_evaluation([[1.0, 0.0], [1.0, 0.0]], [[[0.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 1.0]]], [[1.0, 0.0], [0.0, 0.0]], 0.9)

Output

[1.0, 0.0]

Reasoning

State 1 is absorbing with zero reward so its value stays 0. State 0 takes action 0, collects reward 1 and moves to state 1, so v(0) = 1 + 0.9*0 = 1. The sweeps converge to this after two passes.

Constraints

1 <= n_states <= 200, 1 <= n_actions <= 20
Each pi[s] sums to 1; each P[s][a] sums to 1.
0.0 <= gamma <= 1.0. gamma = 1.0 is only used with MDPs that reach an absorbing zero-reward state.
Updates must be synchronous (new table computed entirely from the old one).
theta and max_sweeps are keyword arguments with the defaults shown.
Do not round inside the function.

Learn this in

Markov Decision ProcessesReinforcement Learning Dynamic ProgrammingReinforcement Learning

More from RL 1: Foundations

GAE Over a Full Trajectoryhard PPO Clipped Objective (Single Sample)medium TD Residuals for GAEmedium Return Standardization Baselinemedium REINFORCE Loss Termmedium

Concept primer

1. Background Knowledge

Policy Evaluation is a fundamental procedure in Reinforcement Learning that determines the value function $v_\pi$ for a given policy $\pi$ . The value function represents the expected cumulative discounted reward an agent will receive starting from a specific state and following the policy thereafter. In a finite Markov Decision Process (MDP), this relationship is defined by the Bellman Expectation Equation. While this equation can be solved directly as a system of linear equations, doing so requires matrix inversion, which has a computational cost of $O(n^3)$ where $n$ is the number of states. This becomes prohibitively expensive for large state spaces.

Instead, we use Iterative Policy Evaluation, which treats the Bellman equation as an update rule. By repeatedly applying the Bellman Backup operation, we generate a sequence of value functions that converge to the true $v_\pi$ . The update rule for a state $s$ at iteration $k+1$ is:

$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]$

This process relies on the property that the Bellman operator is a $\gamma$ -contraction mapping. This means that each iteration brings the value estimates closer to the true values, guaranteeing convergence to the unique fixed point $v_\pi$ regardless of the initial values (typically initialized to zero). The discount factor $\gamma$ ( $0 \le \gamma < 1$ ) ensures that future rewards are weighted less, stabilizing the convergence.

2. Algorithm Approach

The core approach is Synchronous Iterative Update. Unlike asynchronous methods (like TD-learning) that update values in-place or randomly, synchronous evaluation computes a completely new value table $v_{k+1}$ based entirely on the previous table $v_k$ . This ensures that no information from the current sweep "leaks" forward prematurely, maintaining the mathematical guarantee of the contraction mapping.

The algorithm follows a Fixed-Point Iteration pattern:

Initialize: Create a value array v of size $n$ (number of states), filled with zeros.
Iterate: Repeat the following steps until convergence or a maximum iteration limit is reached:

Create a copy of the current values, v_new.
For every state $s$ , compute the expected return using the policy $\pi$ , transition probabilities $P$ , and rewards $R$ .
Store the result in v_new[s].
Calculate the maximum absolute difference $\Delta$ between v_new and v.
Update v to v_new.

Terminate: Stop when $\Delta < \theta$ (a small threshold indicating stability) or when max_sweeps is exceeded.

This approach is essentially Gauss-Seidel style iteration but strictly separated into read-only and write-only phases per sweep to ensure synchronous behavior.

3. Step-by-Step Strategy

Determine State Space: Identify the number of states $n$ from the input dimensions (e.g., len(P)). Initialize v as a list of $n$ zeros.
Loop Control: Start a while loop that continues as long as sweep_count < max_sweeps.
Prepare New Values: Inside the loop, create v_new as a copy of v (or a new list of zeros) to store the updated values for the current sweep. Initialize delta = 0.
Compute Bellman Backup:

Iterate through each state $s$ from $0$ to $n-1$ .
Initialize expected_value = 0 for state $s$ .
Iterate through each action $a$ available in the MDP.
Retrieve the policy probability $\pi(a|s)$ . If it is zero, you can skip this action for efficiency.
Calculate the immediate reward $R[s][a]$ .
Calculate the expected future value: Sum over all next states $s'$ , multiplying the transition probability $P[s][a][s']$ by the current value $v[s']$ (from the previous sweep).
Combine these: term = pi[a] * (R[s][a] + gamma * sum(P[s][a][s'] * v[s'])).
Add term to expected_value.
Set v_new[s] = expected_value.

Track Convergence:

Compute the absolute difference abs(v_new[s] - v[s]).
Update delta = max(delta, abs_diff).

Update and Check:

Set v = v_new.
If delta < theta, break the loop.
Increment sweep_count.

Return: Return the final list v.

4. Common Pitfalls

In-Place Updates: A critical error is updating v[s] directly while iterating. If you update v and then use it to calculate v in the same sweep, you are performing an asynchronous update. This changes the convergence properties and may not yield the correct synchronous result expected by the problem. Always use a separate v_new array.
Indexing Errors: The inputs are nested lists. Ensure you correctly access P[s][a][s']. Confusing the order of indices (e.g., P[a][s][s']) will lead to incorrect calculations or index out-of-bounds errors.
Ignoring Zero Probabilities: While not strictly an error, iterating over actions with $\pi(a|s) = 0$ wastes computation. However, for correctness, ensure you handle the summation correctly even if probabilities are non-zero but small.
Floating Point Precision: The threshold $\theta$ is small ( $10^{-8}$ ). Ensure you are using standard floating-point arithmetic. Do not round intermediate values, as this can prevent convergence or introduce significant error. Only round the final output if required by the grader.
Infinite Loops: If theta is too small or gamma is very close to 1, convergence might be slow. Always respect the max_sweeps limit to prevent Time Limit Exceeded (TLE) errors.

5. Time & Space Complexity

Time Complexity:
Let $n$ be the number of states and $|\mathcal{A}|$ be the number of actions.
One sweep involves iterating over all $n$ states and, for each state, all $|\mathcal{A}|$ actions. For each action, we sum over $n$ possible next states.
Cost per sweep: $O(n \cdot |\mathcal{A}| \cdot n) = O(n^2 |\mathcal{A}|)$ .
The number of sweeps depends on the convergence rate, which is governed by $\gamma$ . Typically, it takes $O(\log(1/\theta))$ sweeps to converge.
Total Time: $O(n^2 |\mathcal{A}| \cdot \log(1/\theta))$ . In the worst case (many sweeps), this is bounded by max_sweeps.
Space Complexity:
We store the value function v and v_new, each of size $n$ .
The input structures $P$ , $R$ , and $\pi$ are given.
Auxiliary space is $O(n)$ for the value arrays.
Total Space: $O(n)$ (excluding input storage).

solution.py

Python 3.13

Test results