If the current state is $S_t$ and the actions are chosen according to $\pi$, what is the expectation of $R_{t+1}$ in terms of $\pi$ and $p$?

Question

I'm trying to solve exercise 3.11 from the book Sutton and Barto's book (2nd edition)

Exercise 3.11 If the current state is $S_t$ , and actions are selected according to a stochastic policy $\pi$, then what is the expectation of $R_{t+1}$ in terms of $\pi$ and the four-argument function $p$ (3.2)?

Here's my attempt.

For each state $s$, the expected immediate reward when taking action $a$ is given in terms of $p$ by eq 3.5 in the book:

$$r(s,a) = \sum_{r \in R} r \, \sum_{s'\in S} p(s',r \mid s,a) = E[R_t \mid S_{t-1} = s, A_{t-1} = a] \tag{1}\label{1}$$

The policy $\pi(a \mid s)$, on the other hand, gives the probability of taking action $a$ given the state $s$.

Is it possible to express the expectation of the immediate reward over all actions $A$ from the state $s$ using (1) as

$$E[R_t \mid S_{t-1} = s, A] = \sum_{a \in A} \pi(a \mid s) r(a,s) \tag{2}\label{2}$$

?

If this is valid, is this also valid in the next time step

$$E[R_{t+1} \mid S_{t} = s, A] = \sum_{a \in A} \pi(a \mid s) r(a, s) \tag{3}\label{3}$$

?

If (2) and (3) are OK, then

$$E[R_{t+1} \mid S_{t} = s, A] = \sum_{a \in A} \pi(a \mid s) \sum_{r \in R} r \, \sum_{s'\in S} p(s',r \mid s,a)$$

Answer 1

First note that $\mathbb{E}[R_{t+1} |S_t=s] = \sum_{s',r}rm(s',r|s)$ where $m(\cdot)$ is the mass function for the joint distribution of $S_{t+1},R_{t+1}$.

If you are currently in state $S_t$ and we condition on taking action $a$ then the expected reward at time $t+1$ is given as follows:

\begin{align} \mathbb{E}[R_{t+1} | S_t = s, A_t=a] & = \sum_{s',r}rp(s'r|s,a)\;. \end{align}

However, action $A_t$ is taken according to some stochastic policy $\pi$ so we need to marginalise this out of our expectation by using the tower law - i.e. we take

$$\mathbb{E}_{A_t \sim \pi}[\mathbb{E}[R_{t+1} | S_t = s, A_t=a]|S_t = s] = \sum_a \pi(a|s)\sum_{s',r}rp(s'r|s,a) = \mathbb{E}[R_{t+1} | S_t = s]\;.$$

To see why this holds we can re-write using some arbitrary mass functions $f(\cdot),h(\cdot),g(\cdot),m(\cdot)$ as

$$\pi(a|s)p(s'r|s,a) = \frac{f(a,s)}{g(s)} \times \frac{h(s',r,a,s)}{f(a,s)} = m(s',r,a|s)\;,$$ so when we would end up with (after re-arranging the summations)

$$\sum_{s',r}r \sum_{a}m(s',r,a|s) = \sum_{s',r}r m(s',r|s) = \mathbb{E}[R_{t+1}|S_t = s]\;;$$ as required.

NB: what you have done is mostly correct except be careful when going from (2) to (3). They are exactly the same equation except for the time stamps, which means you would have to change the time stamps in your $r(s,a)$. Note that when you are at time step $t$ you take your action $A_t$ from your current $S_t$ to transition into state $S_{t+1}$ and then receive reward $R_{t+1}$ (and the next state).