Policy Evaluation

Compute the value of every state under a fixed policy in a known MDP.

A finite MDP has S states and A actions. You are given:

• P, shape (S, A, S): P[s, a, s2] is the probability of moving to s2 after taking action a in state s
• R, shape (S, A): the expected reward for taking action a in state s
• policy, shape (S, A): policy[s, a] is the probability of choosing a in state s

The value of a state is the expected discounted return when you start there and follow the policy. It satisfies the Bellman expectation equation:

V(s) = sum_a policy[s, a] * (R[s, a] + gamma * sum_s2 P[s, a, s2] * V(s2))

Start from V = 0 and apply this update to all states at once, again and again, until the largest change in one sweep is below tol. You may instead solve the linear system (I - gamma * P_pi) V = R_pi directly. Both reach the same answer.

Function signature

def evaluate_policy(P, R, policy, gamma, tol=1e-10):

Example

Input

# two states; action 0 stays, action 1 switches
P = np.array([[[1.0, 0.0], [0.0, 1.0]],
              [[0.0, 1.0], [1.0, 0.0]]])
R = np.array([[0.0, 1.0],
              [2.0, 0.0]])
policy = np.array([[0.0, 1.0],    # state 0 switches
                   [1.0, 0.0]])   # state 1 stays
evaluate_policy(P, R, policy, gamma=0.9)

Output

array([19., 20.])

Reasoning

State 1 stays forever and earns 2 each step, so V(1) = 2 / (1 - 0.9) = 20. State 0 earns 1 once and then moves to state 1, so V(0) = 1 + 0.9 * 20 = 19.

Constraints and notes

• 0 <= gamma < 1.
• Rows of policy and each P[s, a] sum to 1.
• Return an array of shape (S,).