What is the solution to the exercise 3.11 in the RL book by Sutton and Barto

Question

I'm trying to solve the exercise 3.11 from the second-edition from the Sutton and Barto's book

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| s,a) = E[R_t | S_{t-1} = s, A_{t-1} = a]$ (1)

The policy $\pi(a | s)$ on the other hand gives the probability of taking action $a$ given 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 | S_{t-1} = s, A] = \sum_{a \in A} \pi(a|s) r(a,s) \ \ \ \ \ \ \ \ (2) ?$$

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

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

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

$$E[R_{t+1} | S_{t} = s, A] = \sum_{a \in A} \pi(a|s) \sum_{r \in R} r \, \sum_{s'\in S} p(s',r| 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).