Reinforcement learning objective as conditional expectations

Question

In one of his lectures Levine describes the objective of reinforcement learning as: $$J(\tau) = E_{\tau\sim p_\theta(\tau)}[r(\tau)]$$ where $\tau$ refers to a single trajectory and $p_\theta(\tau)$ is the probability of having taken that trajectory so that $p_\theta(\tau) = p(s_1)\prod_{t = I}^T \pi_{\theta}(a_t, s_t)p(s_{t+1}|s_t, a_t))$.

Starting from this definition, he writes the objective as $J(\tau) =\sum_{t=1}^T E_{(s_t, a_t)\sim p_\theta(\tau)}[r(s_t, a_t)]$ and argues that this sum can be decomposed by using conditional expectations, so that it becomes:

$$J(\tau) = E_{s_1 \sim p(s_1)}[E_{a_1 \sim \pi(a_1|s_1)}[r(s_1, a_1) + E_{s_2 \sim p(s_2|s_1, a_1)}[E_{a_2 \sim \pi(a_2|s_2)}[r(s_2, a_2)] + ...|s_2]|s_1,a_1]|s_1]]$$

Can anyone explain this last step? I guess the law of total expectation is involved, but I can not figure out how exactly.

Answer 1

In the YouTube depiction of CS294-112 fall 2017 lecture 3 Reinforcement Learning, Levine, the transition of the finite horizon expected reward to a form where each transition is decoupled from the entire Markov chain of state-action marginals is explained between $t_{video}$ = 44:04 and $t_{video}$ = 45:22.

At t=44:29, the probabilistic expectation where no actions take place is given as a starting point. It's a distribution already known, the initial state distribution.

The statement, "Inside that expectation ...," at $t_{video}$ = 44:39, begins the description of how to build an expectation expression for the first two vertices of the Markov chain. This is the beginning of a process where expectations $E$ are factored recursively. The purpose of the expansions is to decouple the expectation at each vertex from those of other vertices in the Markov chain, clarifying dependencies and exposing the mathematical features of the recursion.

The right-hand factor is the integrated expectation of reward over the achievement of the next state-action marginal, where the action is distributed according to the fixed policy.

How a state and action pair fits into a state-action marginals which can then form a Marcov Chain, provided the policy is fixed and known, was explained at $t_{video}$ = 21:40 in this same lecture. Each vertex in the chain is taken to be the aggregation of a state with the proposed action to be applied to it, leading to a reward $r(s, a)$.

Returning to $t_{video}$ = 44:39, notice the plus sign without the right-hand addend, waiting for another recursion. When the recursion is expressed later in the process, the right-hand factor becomes the integrated expectation of reward over the achievement of the finite sequence represented by the remainder of the Markov chain.

Multiplication is implied because we are now conditioning on both probabilities of the initial state and a transition from it. Notice also that the second vertex in the Marcov Chain is a state-action marginal.

Notice now that rewards are additive because the goal of learning was defined as a sum of products, which is a simplifying assumption.

Also notice that each expectation distributions of reward after the first one can be expressed as a product of single-vertex expectation distributions. Each factor is a localized expectation along the sequence of vertices in the Markov chain of state-action marginals. This is the factoring of expectation distributions.

This factoring is the application of Product distribution, which can be proven using the "law" of total expectations.

At $t_{video}$ = 44:51, the pattern recursion is expressed with a single ellipsis. Recognition of the recursion pattern can be facilitated by re-expressing with a few modifications of mathematical style.

• Use of the original variable to represent the value objective
• Extending the equality
• Using two ellipses instead of one clarify the recursion is a continuation of conditions from right to left as well as continuation of rewards from left to right
• Strategically placing line returns

$$\pi_{\theta}(\tau) = \sum_{t=1}^T \operatorname{E}_{(s_t, a_t)\sim p_\theta(\tau)}[r(s_t, a_t)] \\ = E_{s_1 \sim p(s_1)}[E_{a_1 \sim \pi(a_1|s_1)}[r(s_1, a_1) \\ + E_{s_2 \sim p(s_2|s_1, a_1)}[E_{a_2 \sim \pi(a_2|s_2)}[r(s_2, a_2)] \\ + \; ... \\ ...|s_2]|s_1,a_1]|s_1]]$$

Inverting the transformation, recombining the expansion back to its original form may also be illustrative.

Answer 2

I tried to sketch a mathematical justification of the equality. So we have:

$$J(\tau) = E_{\tau\sim p_\theta(\tau)}[r(\tau)]$$ where $p_\theta(\tau) = p(s_1)\prod_{t = 1}^T \pi_{\theta}(a_t, s_t)p(s_{t+1}|s_t, a_t))$.

Now, $p_\theta(\tau)$ can be re-written in terms of the Markov Chain transition probabilities, namely: $p_\theta(\tau)= \prod_{t = 1}^T \pi_{\theta}(a_t, s_t)p(s_{t}|s_{t-1}, a_{t-1})) = \prod_{t = 1}^T p(s_{t}, a_{t})$.

Here we focus on T = 2 (I guess it is not difficult to prove with a general T, maybe by induction over T). The following holds: $$J(\tau) = E_{\tau\sim p_\theta(\tau)}[r(\tau)] = E_{(s_1, a_1, s_2, a_2)\sim p(s_1,a_1)p(s_2,a_2)}[r(s_1,a_1) + r(s_2,a_2)] \\ = \int_{(s_1,a_1)} \int_{(s_2,a_2)} (r(s_1,a_1) + r(s_2,a_2))p(s_1,a_1)p(s_2,a_2) d(s_1,a_1)d(s_2, a_2) \\ = \int_{(s_1,a_1)} \left( \int_{(s_2,a_2)} (r(s_1,a_1) + r(s_2,a_2))p(s_2,a_2) d(s_2,a_2)\right)p(s_1,a_1)d(s_1, a_1) \\ = E_{(s_1, a_1)\sim p(s_1,a_1)}\left[E_{(s_2, a_2)\sim p(s_2,a_2)}[r(s_1,a_1) + r(s_2,a_2)]\right] \\ = E_{(s_1, a_1)\sim p(s_1,a_1)} \left[E_{(s_2, a_2)\sim p(s_2,a_2)}[r(s_1,a_1)] + E_{(s_2, a_2)\sim p(s_2,a_2)}[r(s_2,a_2)] \right] \\ = E_{(s_1, a_1)\sim p(s_1,a_1)} \left[ r(s_1,a_1) + E_{(s_2, a_2)\sim p(s_2,a_2)}[r(s_2,a_2)] \right] $$

where the second line is due to the definition of the expectation, in the third line we just changed the order of the terms, in the fourth we used the definition of expectation, in the fifth we used the linearity property of the expectation and the last line is because $r(s_1, a_1)$ is constant with respect to the integration over $(s_2, a_2)$. I preferred to move to the integral form of the expectation in line two and three because it was not clear to me how exactly the expectation could factorise.

To conclude the justification, note that we defined $p(s_t, a_t) = \pi_{\theta}(a_t, s_t)p(s_{t}|s_{t-1}, a_{t-1})$, so we have another product here which can be factorised in the same way we did in line two and three of the set of equations above. This brings us the to conditional expectation form of the objective.

In general, every term we have in the original trajectory distribution $p_\theta(\tau)$ corresponds to an expectation and every reward can be moved up to the first outer expectation which it depends on.