> ## Documentation Index
> Fetch the complete documentation index at: https://aegean.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Policy Improvement

> The policy improvement step in dynamic programming.

In the policy improvement step we are given the value function and simply apply the greedy heuristic to it.

$\pi^\prime = \mathtt{greedy}(v_\pi)$

It can be shown that this heuristic results into a policy that is better than the one the prediction step started ($\pi^\prime \ge \pi$) and this extends into multiple iterations. We can therefore converge into an optimal policy - the interested reader can follow [this](https://www.youtube.com/watch?v=Nd1-UUMVfz4\&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ\&index=3) lecture for a justification.

The step itself is a simple one-pass greedification: given $V_\pi$, compute $Q_\pi(s, a)$ for every action and take the argmax.

```python theme={null}
import numpy as np

def greedify(V, P, R, gamma):
    """One-step policy improvement: π'(s) = argmax_a [ R(s,a) + γ Σ P(s'|s,a) V(s') ].

    Args:
        V: value function, shape (nS,)
        P: transition probabilities, shape (nS, nA, nS)
        R: expected immediate rewards, shape (nS, nA)
        gamma: discount factor

    Returns:
        pi: deterministic greedy policy, shape (nS,), values in [0, nA)
    """
    Q = R + gamma * P @ V              # shape (nS, nA)
    return np.argmax(Q, axis=1)
```

The full *policy iteration* algorithm alternates policy evaluation with this step until the policy is stable; see [Policy Iteration](/aiml-common/lectures/mdp/dynamic-programming-algorithms/policy-iteration) for the outer loop and a gridworld demo.

**Key references**: (Mansour & Singh, 2013; Ahmed et al., 2018; Schulman et al., 2015; Huang et al., 2022)

## References

* Ahmed, Z., Le Roux, N., Norouzi, M., Schuurmans, D. (2018). *Understanding the impact of entropy on policy optimization*.
* Huang, S., Kanervisto, A., Raffin, A., Wang, W., Ontañón, S., et al. (2022). *A2C is a special case of PPO*.
* Mansour, Y., Singh, S. (2013). *On the Complexity of Policy Iteration*.
* Schulman, J., Levine, S., Moritz, P., Jordan, M., Abbeel, P. (2015). *Trust Region Policy Optimization*.

***

<Callout icon="pen-to-square" iconType="regular">
  [Edit this page on GitHub](https://github.com/aegean-ai/eaia/edit/main/src/aiml-common/lectures/mdp/policy-improvement/index.mdx) or [file an issue](https://github.com/aegean-ai/eaia/issues/new/choose).
</Callout>