Greedy Policy Improvement

easyFREEDynamic ProgrammingReinforcement LearningRL 1: Foundations

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.

Example

Input

policy_improvement([0, 1], [1.0, 5.0], [[[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]], [[0.0, 1.0], [2.0, 0.0]], 0.9)

Output

[0, 1] True

Reasoning

State 0: q(0,0)=0.0+0.95.0=4.5 beats q(0,1)=1.0+0.91.0=1.9, so the greedy action is 0. State 1: q(1,0)=2.0+0.91.0=2.9 versus q(1,1)=0.0+0.95.0=4.5, so the greedy action is 1. The greedy policy [0, 1] matches the input policy, so stable is True.

Constraints

1 <= n_states <= 200, 1 <= n_actions <= 20
Ties for the maximum must be broken toward the lowest action index.
stable must be True only when the new policy equals the input policy in every state.
0.0 <= gamma <= 1.0
Return (new_policy, stable) in that order.

Learn this in

Dynamic ProgrammingReinforcement Learning

More from RL 1: Foundations

GAE Over a Full Trajectoryhard PPO Clipped Objective (Single Sample)medium TD Residuals for GAEmedium Return Standardization Baselinemedium REINFORCE Loss Termmedium

Concept primer

1. Background Knowledge

Policy Iteration is a fundamental algorithm in Reinforcement Learning used to find the optimal policy for a Markov Decision Process (MDP). It consists of two alternating steps: Policy Evaluation and Policy Improvement. This problem focuses exclusively on the improvement step. The core theoretical guarantee is the Policy Improvement Theorem, which states that if a new policy $\pi'$ is greedy with respect to the value function $v_\pi$ of an existing policy $\pi$ , then $\pi'$ is guaranteed to be at least as good as $\pi$ for all states. Specifically, $v_{\pi'}(s) \ge v_\pi(s)$ for all $s$ .

The "greedy" policy selects the action that maximizes the expected return from each state. The expected return for taking action $a$ in state $s$ is defined by the Bellman Expectation Equation for the action-value function $q_\pi(s, a)$ :

q_\pi(s, a) = r(s, a) + \gamma \sum_{s'} p(s' \mid s, a) v_\pi(s')

Here, $r(s, a)$ is the immediate reward, $\gamma$ is the discount factor, $p(s' \mid s, a)$ is the transition probability to next state $s'$ , and $v_\pi(s')$ is the estimated value of that next state. By choosing $a = \arg\max_a q_\pi(s, a)$ , we construct a policy that locally maximizes the expected future reward based on the current value estimates.

A critical aspect of policy iteration is the stopping criterion. If the greedy policy $\pi'$ is identical to the current policy $\pi$ , then $\pi$ is already optimal, and $v_\pi$ satisfies the Bellman Optimality Equation. This equality check is what determines the stable flag. In practice, floating-point arithmetic can lead to ties between actions with nearly identical Q-values. To ensure deterministic behavior and prevent infinite loops between equally good policies, a strict tie-breaking rule is required: always choose the action with the lowest index when Q-values are equal.

2. Algorithm Approach

The approach is a direct implementation of the Policy Improvement step. For each state in the MDP, you must calculate the Q-value for every possible action using the provided value function $v$ , transition probabilities $P$ , rewards $R$ , and discount factor $\gamma$ . Once all Q-values for a state are computed, select the action that yields the maximum Q-value. If multiple actions yield the same maximum Q-value, select the one with the smallest index.

The algorithm iterates through each state $s$ from $0$ to $n\_states - 1$ . For each state, it iterates through all available actions $a$ . It computes the expected return for each action by summing the immediate reward and the discounted value of all possible next states, weighted by their transition probabilities. After evaluating all actions for a state, it identifies the best action and stores it in the new_policy list. Finally, it compares the new_policy with the input policy to determine stability.

3. Step-by-Step Strategy

Initialize new_policy: Create an empty list or a list of zeros with length equal to the number of states. This will store the greedy action for each state.
Iterate over States: Loop through each state index $s$ from $0$ to $n\_states - 1$ .
Compute Q-values for Actions: For the current state $s$ , initialize variables to track the best_action and max_q_value. Set max_q_value to negative infinity to ensure any valid Q-value will replace it.
Iterate over Actions: Loop through each action index $a$ available in state $s$ .

Calculate the immediate reward $r = R[s][a]$ .
Calculate the expected future value: iterate through all possible next states $s'$ , multiplying the transition probability $P[s][a][s']$ by the value $v[s']$ , and sum these products.
Compute the Q-value: $q = r + \gamma \times \text{expected\_future\_value}$ .

Select Best Action: Compare the computed $q$ with max_q_value.

If $q > \text{max\_q\_value}$ , update max_q_value to $q$ and set best_action to $a$ .
If $q == \text{max\_q\_value}$ , do not update. This ensures that if a later action has the same value, the earlier (lower index) action is retained, satisfying the tie-breaking rule.

Store Result: Assign best_action to new_policy[s].
Check Stability: After processing all states, compare new_policy with the input policy. If they are identical, set stable = True; otherwise, stable = False.
Return: Return the tuple (new_policy, stable).

4. Common Pitfalls

Incorrect Tie-Breaking: The most common error is updating the best action when $q \ge \text{max\_q\_value}$ . This would select the last action with the maximum value, violating the requirement to break ties toward the lowest action index. You must use strict inequality ( $>$ ) for updates.
Floating-Point Precision: While the problem specifies a deterministic tie-break, be aware that floating-point comparisons can be tricky. However, since the tie-break rule is explicit (lowest index), you should rely on the strict inequality logic rather than epsilon-based comparisons for equality, unless the problem specifically requires handling near-equal values as ties. In this context, exact equality in the comparison logic combined with the order of iteration handles the tie-break correctly.
Indexing Errors: Ensure you correctly index the 3D array $P$ . The structure is $P[s][a][s']$ . Confusing the order of indices (e.g., $P[a][s][s']$ ) will lead to incorrect Q-value calculations.
Ignoring Transition Probabilities: Do not simply take the reward of the most likely next state. You must sum over all possible next states $s'$ , weighted by their probabilities $p(s' \mid s, a)$ . Even states with low probability contribute to the expected value.
Modifying Input Policy: Do not modify the input policy list in place. Create a new list for new_policy to ensure the stability check compares the original policy against the newly computed one.

5. Time & Space Complexity

Time Complexity: Let $N$ be the number of states and $A$ be the maximum number of actions per state. For each state, we iterate through all actions. For each action, we iterate through all possible next states (which is at most $N$ ). Thus, the complexity is $O(N \cdot A \cdot N) = O(N^2 A)$ . In many grid-world scenarios, the number of next states is small and constant, making it effectively $O(N A)$ .
Space Complexity: The algorithm requires $O(N)$ space to store the new_policy list. The input arrays $P$ , $R$ , and $v$ are read-only. No additional significant memory is allocated, so the auxiliary space complexity is $O(N)$ .

solution.py

Python 3.13

Test results