Is my proof of equation 0.6 in the book "Reinforcement Learning: Theory and Algorithms" correct?

Question

In Sham Kakade's Reinforcement Learning: Theory and Algorithms, this equation (page 17) is used preceding the proof of performance difference lemma.

I am attempting to prove equation 0.6. Here is my current attempt:

\begin{align*} \mathbb{E}_{\tau \sim \rho^\pi}\left[\sum\limits_{t=0}^\infty \gamma^t f(s_t,a_t)\right] &= \sum\limits_{t=0}^\infty \gamma^t \mathbb{E}_{\tau \sim \rho^\pi} [f(s_t,a_t)]\\ &= \sum\limits_{t=0}^\infty \gamma^t \mathbb{E}_{s_t, a_t} [f(s_t,a_t)]\\ &= \sum\limits_{t=0}^\infty \gamma^t \sum\limits_{s, a} \mathbb{P}(s_t = s, a_t = a) f(s,a)\\ &= \sum\limits_{t=0}^\infty \gamma^t \sum\limits_{s} \mathbb{P}(s_t = s) \sum\limits_{a}\pi(a_t = a|s_t = s) f(s,a)\\ &= \frac{1 - \gamma}{1 - \gamma}\sum\limits_{t=0}^\infty \gamma^t \sum\limits_{s} \mathbb{P}(s_t = s) \mathbb{E}_{a \sim \pi(s)} [f(s,a)]\\ &= \frac{1}{1 - \gamma} \sum\limits_{s} (1-\gamma) \sum\limits_{t=0}^\infty \gamma^t \mathbb{P}(s_t = s) \mathbb{E}_{a \sim \pi(s)} [f(s,a)]\\ &=\frac{1}{(1-\gamma)} \mathbb{E}_{s \sim d^\pi}\left[\mathbb{E}_{a \sim \pi(s)}\left[f(s,a)\right]\right] \\ \end{align*}

Is the swapping of expectation and summation in this way allowed (given that the series converges)?

Note that this is not the proof of the performance difference lemma, but just an attempt to show equation 0.6, which is used but not proved in the book.

Answer 1

The expectation of a sum is equal to the sum of the expectation this just follows from the linearity property of expectations

$$ \begin{aligned} E[\sum_{t} f(s_t,a_t)] &= \sum_{\tau} p(\tau)\left(\sum_t f(s_t,a_t)\right) \\ &= \sum_\tau\sum_{t}p(\tau)f(s_t,a_t) \\ &= \sum_t\sum_\tau p(\tau)f(s_{t,\tau},a_{t,\tau}) \\ &= \sum_tE[f(s_t,a_t)] \end{aligned} $$

note that in the penultimate line I swapped the sums around which can be understood as looking at all trajectories for a single time-step instead of looking at time-steps within a single trajectory (the second expression). I append subscript for variables in the penultimate line for clarity since they also depend on which trajectory they were drawn from.

---

As mentioned by nboro in the comments, the linearity of property only holds for infinite sums if

$$\sum_{i=0}^\infty E[|X|]< \infty \quad or \quad E\left[\sum_{i=0}^{\infty}|X|\right]. <\infty$$

More details can be found here. Using this theory alongside knowledge that since we know that value functions in RL are always bounded (in continuing tasks discounting factors are introduced) then we can say there exists an $M$ such that

$$ |f(s_{t,\tau},a_{t,\tau})| \leq M \qquad \forall t,\tau $$

Then

$$ \begin{aligned}\sum_t\sum_\tau \lambda^t p(\tau)|f(s_{t,\tau},a_{t,\tau})| &\leq \sum_t\sum_\tau \lambda^t p(\tau)M \\ &=M\sum_t\lambda^t\sum p(\tau) \\ &=M\sum\lambda^t \\ &< \infty \end{aligned}$$

As for the rest of the proof is there anything else you need clarity on? it looks good to me. From that step to the next step where you change distribution from $p(\tau)$ to $p(s,a)$ is just because of marginalisation but I think you got that.