Value Iteration

Find the optimal values and a greedy optimal policy of a known MDP with value iteration.

Policy evaluation scores one fixed policy. Value iteration searches for the best one. It replaces the average over actions with a maximum, which gives the Bellman optimality update:

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

Start from V = 0 and repeat until the largest change in one sweep is below tol. Then read off the policy as argmax_a Q(s, a) in each state, with ties going to the lower action index.

P has shape (S, A, S) and R has shape (S, A), as in policy evaluation. Return (V, policy), where policy is an integer array of actions.

Function signature

def value_iteration(P, R, 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]])
value_iteration(P, R, gamma=0.9)

Output

(array([19., 20.]), array([1, 0]))

Reasoning

State 1 can stay and earn 2 forever, worth 20. From state 0, staying earns nothing, while switching earns 1 and then reaches the state worth 20, for a total of 19. So the best policy switches in state 0 and stays in state 1.

Constraints and notes

• 0 <= gamma < 1.
• Compute the greedy policy from the final V.