# Understanding why the expectation is over the new policy $\pi'$ in the proof of the Policy Improvement Theorem

In reinforcement learning, policy improvement is a part of an algorithm called policy iteration, which attempts to find approximate solutions to the Bellman optimality equations.

Pages 84 and 85 in Sutton and Barto's book on RL mentions the following theorem:

Policy Improvement Theorem

Given two deterministic policies $$\pi$$ and $$\pi'$$, then $$v_\pi(s) \leq q_\pi(s, \pi'(s)), \forall s \in S.$$

where $$S$$ is the set of all states.

In the the right-hand side of the inequality, the agent acts according to policy $$\pi'$$ (given that $$\pi'(s)$$ is used in the inequality), in the current state $$s$$, and for all subsequent states acts according to policy $$\pi$$.

In the left-hand side of the inequality, the agent acts according to policy $$\pi$$ (hence the subscript $$_\pi$$ of $$v_\pi(s)$$), starting from the current state $$s$$.

The claim is the following

$$v_\pi(s) \leq v_{\pi'}(s), \forall s \in S$$

In other words, $$\pi'$$ is is an improvement over $$\pi$$.

However, I have a difficulty in understanding the proof. This is discussed below.

Proof

$$v_\pi(s) \leq q_\pi(s, \pi'(s)) = \mathbb{E}_{\pi'}[R_{t+1} + \gamma v_\pi(S_{t+1}) \mid S_t = s]$$

I am stuck here. The q-function is evaluated over the policy $$\pi$$ (note the subscript $$_\pi$$ in $$q_\pi(s, \pi'(s))$$). That being the case, how is the expectation over the policy $$\pi'$$?

My guess is the following. In the proof given in Sutton and Barto, the expectation is unrolled in time. At each time step, the agent follows the policy $$\pi'$$ for that particular time step, and then follows $$\pi$$ from then on. In the limit of this process, the policy transforms from $$\pi$$ to $$\pi'$$. As long as the expression for the return inside the expectation is finite, the governing policy should be $$\pi$$; only in the limit of this process does the governing policy transform to $$\pi'$$.

The expectation is over the policy $$\pi'$$ because the action at the state $$S_t = s$$ is taken according to $$\pi'$$, and, for the proof, the book text (2nd edition, paragraph below Equation 4.8) defines $$\pi'$$ to be a policy that is identical to $$\pi$$ except that $$\pi'(s) = a \neq \pi(s)$$, where $$s$$ is one particular state.
So, essential, the book text tries to prove that, for such a changed policy $$\pi'$$, if $$q_\pi(s, a) = q_\pi(s, \pi'(s)) > v_\pi(s)$$, then the changed policy is better than $$\pi$$. Note that $$\pi'(s) = a$$ in this case.
That is why $$q_\pi(s, \pi'(s)) = \mathbb{E}_{\pi'}[R_{t+1} + \gamma v_\pi(S_{t+1}) \mid S_t = s]$$.