Could you explain these 2 steps of the derivation of the Bellman equation as a recursive equation in Sutton & Barto?

Question

I am reading the Sutton & Barto (2018) RL textbook.

On page 59, it derives the recursive property of the state-value function as below. Could you explain the steps of third and fourth equality? Here is what I have. \begin{aligned} E_\pi[R_{t+1}+\gamma G_{t+1}|S_t=s] &= E_\pi[R_{t+1}|S_t=s]+E_\pi[\gamma G_{t+1}|S_t=s] \\ &= \sum_a \pi(a|s) \sum_{s',r} p(s',r|s,a) \times r + ??? \end{aligned}

Answer 1

To expand $\mathbb{E}_\pi[\gamma G_{t+1}|S_t=s]$, you can take the same expectation over next state and reward as for $R_{t+1}$ (in fact this is normally shown without separating the two terms as you have done, but as it is the expansion of this part where you want help, we can do it separately).

The key thing is to move forward one time step - choosing the action using the policy, and the reward and next state using the state transition function - and express the expectation as a sum of the probabilities for the next step (also noticing that this changes the condition from $S_t=s$ to $S_{t+1}=s'$):

$$\mathbb{E}_\pi[\gamma G_{t+1}|S_t=s] = \gamma \sum_{a}\pi(a|s) \sum_{r,s'} p(r,s'|s,a) \mathbb{E}_\pi[\gamma G_{t+1}|S_{t+1}=s']$$

Then, we can notice that $\mathbb{E}_\pi[\gamma G_{t+1}|S_{t+1}=s']$ is $v_{\pi}(s')$ and get

$$\mathbb{E}_\pi[\gamma G_{t+1}|S_t=s] = \gamma \sum_{a}\pi(a|s)\sum_{r,s'} p(r,s'|s,a) v_{\pi}(s')$$

Now we can recombine this with the other expression for expected immediate reward that you have already resolved, because the summation is the same in both parts, and get result 3.14 from the book.

It is more common to resolve both parts of the expectation the same way though, and not split them up to later recombine. The tricky part is perhaps realising that you cannot resolve the expected return fully when looking ahead one time step, only express it as a sum of expected returns from all possible next states.