Policy Improvement Step

Problem

Perform one policy improvement step given the current value function.

For every state compute Q(s,a) = sum of p * (r + gamma * V[next]) for all actions and return the greedy action per state. Break ties by choosing the smallest action index.

Hint

This is policy evaluation's one-step backup followed by an argmax over actions.

Test cases

case 1

in: [[[[[1,0]],[[1,1]]],[[[1,1]],[[1,0]]]],[[1,0],[0,2]],0.9,[0,0]]

expected: [0,1]

case 2

in: [[[[[1,0]],[[1,1]]],[[[1,1]],[[1,0]]]],[[1,0],[0,2]],0.9,[0,5]]

expected: [1,0]

case 3

in: [[[[[1,0]],[[1,1]]],[[[1,1]],[[1,0]]]],[[1,0],[0,2]],1,[1,1]]

expected: [0,1]

case 4

in: [[[[[1,0]],[[1,0]]]],[[2,2]],0.9,[3]]

expected: [0]

Starter codePython

def policy_improvement(transitions, rewards, gamma, V):
    # Your code here
    pass

4 test cases · run your code against them in the editor