Greedy Policy Improvement

Problem Statement

Given a value function and the model, produce the greedy policy with respect to it, and report whether the policy changed.

Background

The policy improvement theorem is the engine of every control algorithm here. If you have $v_\pi$ and you build a new deterministic policy that acts greedily with respect to it,

$\pi'(s) = \arg\max_a \Big[ r(s,a) + \gamma \sum_{s'} p(s' \mid s, a)\, v_\pi(s') \Big]$

then $v_{\pi'}(s) \ge v_\pi(s)$ for every state. Not on average, not eventually — every state, guaranteed. One greedy step can never make the policy worse.

The second half of the theorem is what gives you a stopping rule: if the greedy policy is identical to the one you started from, then $v_\pi$ already satisfies the Bellman optimality equation, so $\pi$ is optimal and you are done. That equality check is the policy_stable flag in policy iteration.

Because floating-point ties do occur (symmetric actions in a gridworld, for instance), you need a deterministic tie-break or the "did it change?" test will flip-flop forever between equally good policies. Here: break ties toward the lowest action index.

Your Task

Implement:

def policy_improvement(policy, v, P, R, gamma):
    ...

• policy — a list of n_states ints, the current deterministic policy.
• v — a list of n_states floats.
• P[s][a][s2], R[s][a], gamma as usual.

Return a tuple (new_policy, stable) where new_policy is a list of ints and stable is a bool that is True exactly when new_policy == policy.

Input / Output Format

Lists in; a (list, bool) tuple out.

Sample

P = [[[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]
R = [[0.0, 1.0], [2.0, 0.0]]
v = [1.0, 5.0]
new_policy, stable = policy_improvement([0, 1], v, P, R, 0.9)
print(new_policy, stable)

Output:

[0, 1] True

State 0: action 0 gives 0.0 + 0.9*5.0 = 4.5 and action 1 gives 1.0 + 0.9*1.0 = 1.9, so action 0 wins. State 1: action 0 gives 2.0 + 0.9*1.0 = 2.9 and action 1 gives 0.0 + 0.9*5.0 = 4.5, so action 1 wins. The greedy policy [0, 1] equals the policy passed in, so stable is True and policy iteration would halt here.