Policy Iteration One Round

Problem

Run one round of policy iteration: first evaluate the current policy, then improve it.

Evaluation starts from V = 0 and performs sweeps synchronous updates V_new[s] = sum of p * (r + gamma * V[next]) under policy[s]. Then compute the greedy policy from that V (smallest action index on ties) and return [V, new_policy].

Hint

Evaluate exactly `sweeps` times, then reuse the same one-step backup for the greedy policy.

Test cases

case 1

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

expected: [[8.784233454094307,0],[0,1]]

case 2

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

expected: [[8.321905377563029,9.246561530625588],[0,1]]

case 3

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

expected: [[3.875],[0]]

Starter codePython

def policy_iteration_round(transitions, rewards, gamma, policy, sweeps):
    # Your code here
    pass

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