Solve the Bellman Expectation Equation Exactly

Problem Statement

Given a policy already folded into the MDP, compute $v_\pi$ exactly by solving a linear system — no iteration.

Background

Fix a policy and the MDP collapses into a Markov reward process: a transition matrix $P^\pi$ where $P^\pi_{s,s'} = \sum_a \pi(a\mid s)\, p(s'\mid s,a)$, and a reward vector $R^\pi_s = \sum_a \pi(a \mid s)\, r(s,a)$. The Bellman expectation equation is then one vector equation:

$v_\pi = R^\pi + \gamma P^\pi v_\pi$

This is linear in $v_\pi$, which people often miss because the recursive statement looks like something you have to unroll. Rearranged:

$(I - \gamma P^\pi)\, v_\pi = R^\pi \quad\Longrightarrow\quad v_\pi = (I - \gamma P^\pi)^{-1} R^\pi$

For gamma < 1 the matrix $I - \gamma P^\pi$ is always invertible, so the solution exists and is unique. The catch is cost: a direct solve is $O(n^3)$ in the number of states, which is why iterative policy evaluation exists at all. Knowing the closed form still matters — it is the ground truth you check your iterative code against.

Your Task

Implement:

def solve_v_pi(P_pi, R_pi, gamma):
    ...

• P_pi — an n x n nested list, P_pi[s][s2] is the probability of moving from s to s2 under the policy.
• R_pi — a list of n floats, the expected immediate reward in each state under the policy.
• gamma — discount factor, 0.0 <= gamma < 1.0.

Return a list of n floats.

Input / Output Format

Nested lists of floats in, a list of floats out. The grader rounds to 4 decimals.

Sample

P = [[0.5, 0.5, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 1.0]]
R = [1.0, 2.0, 0.0]
print([round(x, 4) for x in solve_v_pi(P, R, 0.9)])

Output:

[3.4545, 2.0, 0.0]

State 2 is absorbing with zero reward so v(2) = 0. Then v(1) = 2 + 0.9*0 = 2. Finally v(0) = 1 + 0.9(0.5v(0) + 0.52)** gives **0.55v(0) = 1.9**, i.e. v(0) = 3.4545....