Policy iteration

Policy evaluation for a given $π (s)$ :

V_{k}^{π} (s) \leftarrow \sum_{s^{'}} P (s^{'} | s, π (s)) (R (s, π (s), s^{'}) + γ V_{k - 1}^{π} (s^{'}))

At convergence:

\forall s V^{π} (s) \leftarrow \sum_{s^{'}} P (s^{'} | s, π (s)) (R (s, π (s), s^{'}) + γ V^{π} (s^{'}))

$V^{π} (s)$ is the expected discounted return from $s$ when following the same policy $π$ at every step. This evaluates the given policy; its value need not be optimal.

$π (s)$ is the action chosen by the policy in the current state.
$P (s^{'} | s, π (s))$ is the probability of reaching each possible next state.
$R (s, π (s), s^{'})$ is the immediate reward for that transition.
$γ V^{π} (s^{'})$ is the discounted return from the next state, continuing with the same policy.

The sum is a probability-weighted average of reward now plus discounted future return. There is no $max$ because $π (s)$ already chooses the action.

At convergence, successive rounds approach the same $V^{π}$ , so the iteration index $k$ disappears and $\leftarrow$ can be read as equality. There is one equation per state, and the unknown values are linked through the possible next states.

Exercise: Stochastic Policy Evaluation

Consider a stochastic policy $π (a | s)$ , where $π (a | s)$ is the probability of taking action a when in state s. Which of the following is the correct update to perform policy evaluation for this stochastic policy?

Incorrect

$V_{k + 1}^{π} (s) \leftarrow max_{a} \sum_{s^{'}} P (s^{'} | s, a) (R (s, a, s^{'}) + γ V_{k}^{π} (s^{'}))$

max over a chooses the action with the largest expected return. It replaces the current action choice and ignores the given policy’s action probabilities. Policy evaluation must average actions according to π(a|s).
Correct

$V_{k + 1}^{π} (s) \leftarrow \sum_{s^{'}} \sum_{a} π (a | s) P (s^{'} | s, a) (R (s, a, s^{'}) + γ V_{k}^{π} (s^{'}))$

There are two sources of randomness: the policy chooses a with probability π(a|s), and the environment then reaches s′ with probability P(s′|s,a). Their product is the probability of this combined outcome. The two sums include every outcome and compute its probability-weighted reward plus discounted future value.

For finite sums, the order of the two summations can be swapped without changing the answer.
Incorrect

$V_{k + 1}^{π} (s) \leftarrow \sum_{a} π (a | s) max_{s^{'}} P (s^{'} | s, a) (R (s, a, s^{'}) + γ V_{k}^{π} (s^{'}))$

This averages over actions correctly, but max over s′ keeps only the largest single probability-weighted next-state contribution. The other outcomes are left out. An expectation must add the weighted contributions of all possible next states.

For example, suppose there is one action and two possible next states, each with probability 0.5 and reward plus discounted future value of 10. The correct expectation is 0.5 × 10 + 0.5 × 10 = 10. Option 3 takes max(5, 5) = 5.

One iteration of policy iteration:

1. Policy evaluation for current policy $π_{k}$ :

Keep the policy fixed and iterate until convergence:

V_{i + 1}^{π_{k}} (s) \leftarrow \sum_{s^{'}} P (s^{'} | s, π_{k} (s)) [R (s, π_{k} (s), s^{'}) + γ V_{i}^{π_{k}} (s^{'})]

Here, k counts policy changes and i counts value updates while evaluating the same policy. This step gives $V^{π_{k}}$ : the expected return from each state when continuing with the current policy.

2. Policy improvement: find the best action according to one-step look-ahead

π_{k + 1} (s) \leftarrow \underset{a}{arg max} \sum_{s^{'}} P (s^{'} | s, a) [\underset{\begin{matrix} ↑ \\ Reward now \end{matrix}}{R (s, a, s^{'})} + γ \underset{\begin{matrix} ↑ \\ Then follow π_{k} \end{matrix}}{V^{π_{k}} (s^{'})}]

At state s, compare every available action a. For each candidate, ask: “If I take this action now, then follow the old policy from the next state onward, what return should I expect?”

$R (s, a, s^{'})$ is the reward for the candidate action now.
$γ V^{π_{k}} (s^{'})$ includes all later rewards when following the old policy, discounted by one step. These values come from evaluation and stay fixed while comparing actions.
The sum with $P (s^{'} | s, a)$ averages over all possible next states.
arg max returns the action with the largest expected return. Store that action as $π_{k + 1} (s)$ . Do this for every state to build the new policy.

“One-step” means only the next transition is expanded explicitly. All later steps are already included in $V^{π_{k}} (s^{'})$ . The candidate score is $Q^{π_{k}} (s, a)$ ; it uses the current policy’s future return, which may still be below the optimal return.

Example

Suppose the current policy chooses Left at s, the transitions are deterministic, and γ = 0.9. Evaluation gives a value of 5 for the state reached by Left and 8 for the state reached by Right:

Action now	Reward now	Next-state value under the old policy	Candidate score
Left	1	5	1 + 0.9 × 5 = 5.5
Right	0	8	0 + 0.9 × 8 = 7.2

arg max chooses Right, so the new policy chooses Right at s. Its lower immediate reward is outweighed by the larger future return. The score 7.2 assumes the old policy after this first action; evaluate the new policy in the next round to find its value.

Repeat until the policy converges

Alternate evaluation and improvement. The old action is among the candidates, so the chosen action’s score is at least the current value. For a finite discounted MDP with exact evaluation, making these greedy choices in every state gives a new policy whose value is at least as high in every state. One improvement step need not produce the optimal policy.

If improvement leaves the policy unchanged, it is optimal. Keep the current action when it ties for the best score, so ties do not cause unnecessary switches. Policy iteration can converge in fewer outer iterations than value iteration under some conditions; each outer iteration also includes policy evaluation.