# Deriving 'State-Action marginal' in Reinforcement Learning

I see the following equation1 in "In Reinforcement Learning Course CS294". I want to prove "1 equation is same 2 equation."

I tried but failed :( What's wrong..?

$$p_\theta(\tau) = p_\theta(s_0, (s_1, a_1), ..., (s_t, a_t))$$ is the distribution over all trajectories, which could also be seen as the joint distribution over all states and actions under the policy. Thus, $$\text{E}_{\tau \sim p_\theta(\tau)}[r(s_t, a_t)] = \text{E}_{(s_t, a_t) \sim p_\theta(\tau)}[r(s_t, a_t)] = \text{E}_{(s_t, a_t) \sim p_\theta(s_t, a_t)}[r(s_t, a_t)],$$ where $$(s_t, a_t) \sim p_\theta(\tau)$$ could be seen as the probability of $$(s_t, a_t)$$ occurring, no matter which trajectory was taken to get there since it's not relevant here.

The expression $$\text{E}_{(s_t, a_t) \sim p_\theta(s_t, a_t \mid s_{t-1}, a_{t-a} )}[r(s_t, a_t)]$$ is the conditional expectation of the reward at time $$t$$ given the state and action in time $$t-1$$, which is a different entity. This would be interesting when asking questions such as what is the expected return of the next time-step if the current state and action is $$(s_t, a_t)$$. A more rigorous way of looking at it is $$\text{E}_{\tau \sim p_\theta}\left[\sum_{t=1}^T r(s_t, a_t)\right] = \sum_{\tau}p_\theta(\tau)\sum_{t=1}^T r(s_t, a_t) = \sum_{t=1}^T\sum_{\tau} p_\theta(\tau)r(s_t, a_t) = \sum_{t=1}^T\sum_{u=1}^{t-1}r(s_t, a_t)p_\theta\left(s_t, a_t \mid s_0, \{s_k, a_k\}_{k=1}^u\right)p_\theta\left(s_0, \{s_k, a_k\}_{k=1}^u\right) = \sum_{t=1}^T\sum_{u=1}^{t-1}p_\theta\left(s_0, \{s_k, a_k\}_{k=1}^u\right)\text{E}[r(s_t, a_t) \mid \pi_\theta, s_0, \{s_k, a_k\}_{k=1}^u] = \sum_{t=1}^T\text{E}_{(s_t, a_t)\sim p_\theta(s_t, a_t)}[r(s_t, a_t)].$$

This answer about the linearity of expectation for dependent variables in general could help shining some light on the matter.

• Does anybody verify emarcus' answer to be correct? I think he/she missed a sum in the second line that states equations and and the initial state is $\mathbf{s}_1$ rather than $\mathbf{s}_0$.
– dim
Jan 13, 2022 at 10:49

I'm trying to rephrase the second and more rigorous part of emarcus' answer in more detail. We start by assuming that probability over a trajectory $$\tau$$ factorizes like $$p_\theta(\tau)=p_\theta(s_1,a_1,\ldots,s_T,a_T)= p(s_1)\pi_\theta(a_1|s_1)\prod_{t=1}^Tp(s_{t+1}|s_t,a_t)\pi_\theta(a_{t+1}|s_{t+1})=\\ p_\theta(s_1,a_1)\prod_{t=1}^Tp_\theta(s_{t+1},a_{t+1}|s_t,a_t)$$ assuming the graphical model of a Markov decision process and following the notation from Sergey Levine's deepRL course CS285 (although there it for some reason says that $$p_\theta(\tau)=p(s_1)\prod_{t=1}^Tp(s_{t+1}|s_t,a_t)\pi_\theta(a_{t}|s_{t})$$ so $$\tau$$ theoretically ends with $$s_1,a_1\ldots s_T,a_T,s_{T+1}$$). Now we can expand the expectation assuming continuous random variables $$E_{\tau\sim p_\theta(\tau)}[\sum_{t=1}^Tr(s_t,a_t) ]= \int p_\theta(\tau) \sum_{t=1}^Tr(s_t,a_t)d\tau=\sum_{t=1}^T\int p_\theta(\tau)r(s_t,a_t)d\tau$$ where we used the linearity property of the expectation operator. Inserting the result above leads to $$=\sum_{t=1}^T\int p_\theta(s_1,a_1)\prod_{k=1}^Tp_\theta(s_{k+1},a_{k+1}|s_k,a_k)r(s_t,a_t)d\tau$$ Since each summand depends only on terms from $$1$$ to $$t$$, we can further simplify $$=\sum_{t=1}^T\int p_\theta(s_1,a_1)\prod_{k=2}^{t}p_\theta(s_{k},a_{k}|s_{k-1},a_{k-1})r(s_t,a_t)d\{s_k,a_k\}_{k=1}^t$$ We now realize that $$p_\theta(s_1,a_1)\prod_{k=2}^{t}p_\theta(s_{k},a_{k}|s_{k-1},a_{k-1})=p(\{s_k,a_k\}_{k=1}^t)$$ so $$=\sum_{t=1}^T\int p(\{s_k,a_k\}_{k=1}^t)r(s_t,a_t)d\{s_k,a_k\}_{k=1}^t$$ and use $$E_{(x,y)\sim p(x,y)}[x]=\int p(x,y)xdxdy=\int x\int p(x,y)dy dx=\int xp(x)dx=E_{x\sim p(x)}[x]$$, we arrive at $$=\sum_{t=1}^T\int r(s_t,a_t) \int p(\{s_k,a_k\}_{k=1}^t)d\{s_k,a_k\}_{k=1}^{t-1}ds_t da_t\\=\sum_{t=1}^T\int p(s_t,a_t)r(s_t,a_t)ds_t da_t= \sum_{t=1}^TE_{(s_t,a_t)\sim p(s_t,a_t)}[r(s_t,a_t)]$$

• Thank you for the detailed derivation. This is really helpful! Jun 1, 2023 at 14:57