Policies for which the policy improvement theorem holds

Question

According to Reinforcement Learning (2nd Edition) by Sutton and Barto, the policy improvement theorem states that for any pair of deterministic policies $\pi'$ and $\pi$, if $q_\pi(s,\pi'(s)) \geq v_\pi(s)$ $\forall s \in \mathcal{S}$, then $v_{\pi'}(s) \geq v_\pi(s)$ $\forall s \in \mathcal{S}$.

The proof of this theorem seems to rely on $\pi$ and $\pi'$ being identical for all states except $s$. To my best understanding, this is what allows us to write the expectation $\mathbb{E}[R_{t+1} + \gamma v_\pi(S_{t+1})|S_t = a, A_t = \pi'(s)]$ as $\mathbb{E}_{\pi'}[R_{t+1} + \gamma v_\pi(S_{t+1})|S_t = a]$ in line 2, which is central to the proof (re-produced from the book below).

\begin{aligned} v_{\pi}(s) & \leq q_{\pi}\left(s, \pi^{\prime}(s)\right) \\ &=\mathbb{E}\left[R_{t+1}+\gamma v_{\pi}\left(S_{t+1}\right) \mid S_{t}=s, A_{t}=\pi^{\prime}(s)\right] \\ &=\mathbb{E}_{\pi^{\prime}}\left[R_{t+1}+\gamma v_{\pi}\left(S_{t+1}\right) \mid S_{t}=s\right] \\ & \leq \mathbb{E}_{\pi^{\prime}}\left[R_{t+1}+\gamma q_{\pi}\left(S_{t+1}, \pi^{\prime}\left(S_{t+1}\right)\right) \mid S_{t}=s\right] \\ &=\mathbb{E}_{\pi^{\prime}}\left[R_{t+1}+\gamma \mathbb{E}_{\pi^{\prime}}\left[R_{t+2}+\gamma v_{\pi}\left(S_{t+2}\right) \mid S_{t+1}, A_{t+1}=\pi^{\prime}\left(S_{t+1}\right)\right] \mid S_{t}=s\right] \\ &=\mathbb{E}_{\pi^{\prime}}\left[R_{t+1}+\gamma R_{t+2}+\gamma^{2} v_{\pi}\left(S_{t+2}\right) \mid S_{t}=s\right] \\ & \leq \mathbb{E}_{\pi^{\prime}}\left[R_{t+1}+\gamma R_{t+2}+\gamma^{2} R_{t+3}+\gamma^{3} v_{\pi}\left(S_{t+3}\right) \mid S_{t}=s\right] \\ & \vdots \\ & \leq \mathbb{E}_{\pi^{\prime}}\left[R_{t+1}+\gamma R_{t+2}+\gamma^{2} R_{t+3}+\gamma^{3} R_{t+4}+\cdots \mid S_{t}=s\right] \\ &=v_{\pi^{\prime}}(s) \end{aligned}

Does this mean that the proof is merely proving the special case of the policy improvement theorem for when the policies are identical except at $s$? I am having trouble seeing why the proof holds for the more general case of the two policies being potentially different for all states. In that case, line 2 would not hold and the theorem would not hold for all states as it claims to do.

Answer 1

There are two policies in the derivation, $\pi$ and $\pi'$. When you have an expectation in the form

$$\mathbb{E}_{\pi_A}[R_{t+1} + \gamma v_{\pi_B}(S_{t+1}) | S_t = s]$$

, it means that from your current state $s$, you take an action according to $\pi_A$. Let's say this action is $a$, then the reward function from the environment will transition to the next state, represented by the random variable $S_{t+1}$ and also gives you a reward $R_{t+1}$ which can depend on the current state $s$, the action taken $a$ and the next state $S_{t+1}$. That's where your reliance on $\pi_A$ ends because $v_{\pi_B}$ is an expectation itself:

$$v_{\pi_B}(s) = \mathbb{E}_{\pi_B}[G_{t+1}|S_{t} = s]$$

So going from line 2 to line 3 in your derivation just means that you are taking the next action according to $\pi'$, but nothing inside the expectation changes, it still uses the same policy $\pi$. The next state, $S_{t+1}$, comes from the environment and depends on $\pi'$. The same is true for the observed reward $R_{t+1}$, but nothing inside $v_\pi$ needs the policy $\pi'$.

The proof then replaces $v_\pi(S_{t+1})$ with $q_\pi(S_{t+1}, \pi'(S_{t+1}))$ but the equality changes to a less than or equal to sign because the whole discussion in the book starts with the condition $q_\pi(s, \pi'(s)) \ge v_\pi(s)$.

It does not matter in how many states the policies differ. If the two policies are the same in states $s_i$, then $q_\pi(s_i, \pi'(s_i)) = v_\pi(s_i)$. If policy $\pi'$ is better than policy $\pi$ in states $s_j$, then $q_\pi(s_j, \pi'(s_j)) > v_\pi(s_j)$. The initial condition $q_\pi(s, \pi'(s)) \ge v_\pi(s)$ combines these two statements. If you know in which states the policy $\pi'$ is better, then you could think of replacing the $\le$ signs with $<$ signs in the derivation, and the rest of the states with $=$. You can't really do that because these are random variables, so even if you knew in which states your new policy performs better, you often don't know the environment dynamics that determine the next state, $S_{t+1}$.