Solve the Bellman Expectation Equation Exactly

mediumFREEMarkov Decision Processes and Bellman EquationsReinforcement LearningRL 1: Foundations

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....

Example

Input

solve_v_pi([[0.5, 0.5, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 1.0]], [1.0, 2.0, 0.0], 0.9)

Output

[3.4545, 2.0, 0.0]

Reasoning

Identify the linear system: The Bellman equation $v_\pi = R^\pi + \gamma P^\pi v_\pi$ is rearranged to $(I - \gamma P^\pi)v_\pi = R^\pi$ . We construct the coefficient matrix $A = I - 0.9 P^\pi$ and the reward vector $R^\pi = [1.0, 2.0, 0.0]^T$ .
Compute matrix $A$ : Subtract $0.9$ times the transition matrix from the identity matrix:
$A = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \end{bmatrix} - 0.9 \begin{bmatrix} 0.5 & 0.5 & 0.0 \\ 0.0 & 0.0 & 1.0 \\ 0.0 & 0.0 & 1.0 \end{bmatrix} = \begin{bmatrix} 0.55 & -0.45 & 0.0 \\ 0.0 & 1.0 & -0.9 \\ 0.0 & 0.0 & 0.1 \end{bmatrix}$
Solve for $v_2$ : The third row of the system corresponds to the absorbing state 2. The equation is $0.1 v_2 = 0.0$ , which yields $v_2 = 0.0$ .
Solve for $v_1$ : Substitute $v_2 = 0.0$ into the second row equation $1.0 v_1 - 0.9 v_2 = 2.0$ . This simplifies to $v_1 = 2.0$ .
Solve for $v_0$ : Substitute $v_1 = 2.0$ into the first row equation $0.55 v_0 - 0.45 v_1 = 1.0$ . This gives $0.55 v_0 - 0.9 = 1.0 \implies 0.55 v_0 = 1.9$ , so $v_0 = 1.9 / 0.55 \approx 3.454545...$
Final Output: Rounding the values to 4 decimal places results in the list [3.4545, 2.0, 0.0].

Constraints

1 <= n <= 200
Each row of P_pi sums to 1.
0.0 <= gamma < 1.0, so I - gamma*P_pi is invertible.
Solve the system directly (numpy.linalg.solve is available); do not iterate to convergence.
Do not round inside the function.

Learn this in

Markov Decision ProcessesReinforcement 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

In Reinforcement Learning, the value of a state under a fixed policy $\pi$ , denoted $v_\pi(s)$ , represents the expected cumulative discounted reward starting from that state. When the policy is fixed, the MDP reduces to a Markov Reward Process (MRP). The relationship between state values is governed by the Bellman Expectation Equation:

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

Here, $R^\pi$ is the vector of expected immediate rewards, $P^\pi$ is the transition probability matrix under policy $\pi$ , and $\gamma$ is the discount factor. This equation is recursive, but crucially, it is linear in $v_\pi$ . While iterative methods like Policy Evaluation approximate this solution, we can solve it exactly by treating it as a system of linear equations.

Rearranging the terms to isolate $v_\pi$ :

$v_\pi - \gamma P^\pi v_\pi = R^\pi$ $(I - \gamma P^\pi) v_\pi = R^\pi$

Where $I$ is the identity matrix of size $n \times n$ . Since $\gamma < 1$ , the matrix $(I - \gamma P^\pi)$ is guaranteed to be invertible (non-singular). Therefore, the unique exact solution is:

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

This closed-form solution provides the "ground truth" against which iterative algorithms are validated. It relies on linear algebra concepts such as matrix inversion and linear system solving.

2. Algorithm Approach

The core approach is to transform the Bellman equation into a standard linear algebra problem $Ax = b$ and solve for $x$ .

Construct Matrix A: Compute $A = I - \gamma P^\pi$ . This involves creating an identity matrix of the same dimension as $P^\pi$ , scaling $P^\pi$ by $\gamma$ , and subtracting the result from the identity matrix.
Define Vector b: This is simply the reward vector $R^\pi$ .
Solve Linear System: Instead of explicitly computing the inverse matrix (which is computationally expensive and numerically unstable), use a robust linear solver to find $v_\pi$ such that $A v_\pi = b$ . In Python, numpy.linalg.solve is the standard tool for this.

This approach leverages the fact that for small to medium-sized state spaces ( $n \le 1000$ ), direct linear solvers are faster and more accurate than iterative methods for finding the exact solution.

3. Step-by-Step Strategy

Import NumPy: You will need numpy for efficient matrix operations.
Convert Inputs to Arrays: Convert the nested list P_pi and list R_pi into NumPy arrays. This allows for vectorized operations.
Determine Dimension: Get the number of states $n$ from the shape of P_pi.
Create Identity Matrix: Use np.eye(n) to create the $n \times n$ identity matrix $I$ .
Compute Coefficient Matrix: Calculate $A = I - \gamma \cdot P^\pi$ . Ensure you perform element-wise multiplication for $\gamma \cdot P^\pi$ .
Solve for Values: Use np.linalg.solve(A, R_pi) to compute $v_\pi$ . This function solves the linear system $Ax = b$ directly.
Return Result: Convert the resulting NumPy array back to a standard Python list of floats if required by the signature, though often NumPy arrays are acceptable. Ensure the output format matches the expected list of floats.

import numpy as np

def solve_v_pi(P_pi, R_pi, gamma):
 # Convert to numpy arrays
 P = np.array(P_pi)
 R = np.array(R_pi)
 
 # Get number of states
 n = P.shape
 
 # Create Identity matrix
 I = np.eye(n)
 
 # Compute A = I - gamma * P
 A = I - gamma * P
 
 # Solve A * v = R for v
 v_pi = np.linalg.solve(A, R)
 
 return v_pi.tolist()

4. Common Pitfalls

Explicit Inversion: Avoid using np.linalg.inv(A) @ R. Explicitly inverting a matrix is numerically less stable and slower than solving the linear system directly. Always prefer np.linalg.solve.
Data Types: Ensure that P_pi and R_pi are converted to floating-point arrays. Integer arrays can lead to truncation errors during subtraction or multiplication.
Matrix Dimensions: Verify that P_pi is square ( $n \times n$ ) and R_pi has length $n$ . Mismatched dimensions will cause np.linalg.solve to raise a LinAlgError.
Gamma Constraint: The problem guarantees $\gamma < 1$ , which ensures invertibility. If $\gamma \ge 1$ , the matrix might be singular, and the solution would not exist or be unique. Do not assume invertibility without this constraint.
Output Format: The problem asks for a list of floats. np.linalg.solve returns a NumPy array. Use .tolist() to convert it to a standard Python list to match the expected output type strictly.

5. Time & Space Complexity

Time Complexity: The dominant operation is solving the linear system $Ax = b$ . Using LU decomposition (which np.linalg.solve uses internally), the complexity is $O(n^3)$ , where $n$ is the number of states. Constructing the matrix $A$ takes $O(n^2)$ . Thus, the overall time complexity is $O(n^3)$ .
Space Complexity: We store the matrix $A$ ( $n \times n$ ) and the vectors $R$ and $v_\pi$ ( $n$ ). The space required is dominated by the matrix storage, which is $O(n^2)$ .

This complexity is acceptable for small to medium $n$ (e.g., $n < 1000$ ). For very large state spaces, iterative methods like Value Iteration or Policy Iteration are preferred because they have lower per-iteration costs and can converge to an approximate solution without storing the full $n \times n$ matrix.

solution.py

Python 3.13

Test results