Reinforcement Learning Part 13: Policy Gradient Causality Trick and REINFORCE

In the previous post, we showed how we can substitute our usual ε-greedy policy with a parameterized approximation (often a neural network), we then derived the policy gradient theorem required to optimize this approximator function, and demonstrated how it can be estimated using Monte Carlo sampling. At the very end, we pointed out that the policy gradient uses the full trajectory return, G(τ), at each timestep in a given episode. In practice, this term is almost always simplified to the discounted return, G_t, from that timestep looking forward.

By using G_t instead of G(τ), we can reduce the variance and reduce the number of required samples (also known as “sample efficient”). Understanding this substitution requires digging into the mathematical concept of causality. In the rest of the post, we’ll show how that substitution works and then give an overview of the REINFORCE algorithm, one of the foundational policy gradient methods that paved the way for actor-critic algorithms.

The Problem With Full Trajectory Returns

If you recall from our policy gradient estimate, we defined J(θ) to be the expected return under policy π, given a set of parameters θ.

\(J(\boldsymbol{\theta}) = \sum_{\tau} p(\tau; \boldsymbol{\theta}) \cdot G(\tau) = \mathbb{E}_{s_0 \sim p(s_0)} \left[ v_{\pi_{\boldsymbol{\theta}}}(s_0) \right]\) 

In order to perform gradient ascent, which adjusts the parameters (θ) of the neural network (NN) approximator to ideally produce better returns over each episode, we need to find the gradient of J(θ) via backpropagation. We used the policy gradient theorem to find that the gradient ∇J(θ). This is the expected value, over all trajectories the policy could produce, of the sum (across every timestep in that trajectory) of the gradient of the log-probability of the action taken, multiplied by the trajectory’s total return.

This is a mouthful, but the implication is that we can calculate the policy gradient without needing to know the environment dynamics. In other words, it’s model-free. In practice, it’s difficult (if not impossible) to compute this term analytically, but we can use Monte Carlo sampling to estimate its value:

\(\nabla J(\boldsymbol{\theta}) \approx \frac{1}{N} \sum_{i=1}^{N} \left[ \left( \sum_{t} \nabla_{\boldsymbol{\theta}} \log \pi(a_t^{(i)} | s_t^{(i)}; \boldsymbol{\theta}) \right) \cdot G(\tau^{(i)}) \right]\) 

While this works, we can actually reduce the number of samples required to get a useful signal by reducing the variance of this estimation. Imagine working through the calculation, and you’re at t=5 in a long episode.

The formula calls for multiplying every timestep’s probability (as given by the policy) by the full trajectory’s return, G(τ). The issue is that this gives weight to every timestep, including rewards that happened before t=5. As you can probably guess, an action taken at t=5 shouldn’t, on average, have anything to do with rewards received before it (e.g. t=0..4). Those rewards already happened before action a₅ was even chosen, so there’s no way a₅ could be responsible for them.

We can fix this by ignoring rewards from before t and only using the return looking forward from the current timestep, G_t, which covers the current reward and future discounted rewards. This removes the noise those earlier timesteps were adding, without changing what the gradient is actually estimating (which means we need fewer sampled episodes to get a reliable estimate).

I hope this makes intuitive sense, but let’s take a moment to show how this works by demonstrating that the correlation between the gradient for a_t and rewards prior to timestep t is exactly zero.

The Causality Trick

Separate Past and Future Expected Returns

For this derivation, we’ll need to make a simple substitution for our timestep counting. Instead of using 0-based timesteps, we’ll switch to counting from 1. In other words, j = k + 1. This allows us to rewrite G_t as follows:

\(G_{t} = \sum_{j=1}^{T-t} \gamma^{j-1} r_{t+j} \;=\; \sum_{k=0}^{T-t-1} \gamma^{k} r_{t+k+1}\)

Now, take a moment to conceptualize G(τ) from the perspective of some arbitrary timestep, t. We can break the trajectory, τ, into two sections: everything that came before t (which we’ll call the “prefix”) and everything that comes after (up to the terminal state s_T):

\(\tau = \underbrace{(s_0, a_0, r_1, \ldots, r_t, s_t)}_{\text{before t (prefix)}} \;+\; \underbrace{(a_t, r_{t+1}, s_{t+1}, \ldots, s_T)}_{\text{everything after}}\)

Note where we are creating this break: we just observed s_t, and we are just about to take the action at that timestep (a_t). From this, we can break the total episodic return apart into two sections: all the rewards that came before t and all the rewards that will come after.

\(G(\tau) = \underbrace{\sum_{j=1}^{t} \gamma^{j-1} r_j}_{\text{past rewards}} \;+\; \underbrace{\sum_{j=t+1}^{T} \gamma^{j-1} r_j}_{\text{future rewards}}\)

Let’s look at only the second term (the one labeled “future rewards”). We can use the law of exponents, which states that exponents add when you multiply powers with the same base. The sum then becomes the definition of G_t, as we covered in part 4.

\(\begin{align} \sum_{j=t+1}^{T} \gamma^{j-1} r_j &= \sum_{j=t+1}^{T} \gamma^t \cdot \gamma^{j-1-t} \, r_j \\ &= \gamma^t \sum_{j=t+1}^{T} \gamma^{j-1-t} \, r_j \\ &= \gamma^t \cdot G_t \end{align}\)

So, now we’ve established that we can write G(τ) in terms of G_t, where G_t is the sum of discounted rewards that comes after some arbitrary timestep t. 

\(G(\tau) = \left( \sum_{j=1}^{t} \gamma^{j-1} r_j \right) + \gamma^t \cdot G_t\)

With our Monte Carlo estimation of ∇J(θ), let’s look at a single episode. There, we can pull the G(τ) value inside the sum, as it’s just a constant for that episode.

\(\left( \sum_{t} \nabla_\theta \log \pi(a_t|s_t;\theta) \right) \cdot G(\tau) = \sum_{t} \left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G(\tau) \right]\)

We can then substitute in our new definition of G(τ) (in terms of G_t).

\(\sum_{t} \left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G(\tau) \right] = \sum_{t} \left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left( \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) + \gamma^t \cdot G_t \right) \right]\)

We can then distribute the policy gradient term to the past and future reward pieces:

\(\sum_{t} \left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G(\tau) \right] = \sum_{t} \left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) + \gamma^t \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G_t \right]\)

Looking at only the term inside the summation (i.e. one particular timestep t), we can then distribute the policy gradient term to the past and future reward pieces:

\(\nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left( \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) + \gamma^t \cdot G_t \right) = \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) \;+\; \gamma^t \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G_t\)

At this point, we can plug this “past plus future” return back into our definition of the policy gradient:

\(\nabla J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[ \sum_{t} \left( \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) \;+\; \gamma^t \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G_t \right) \right]\)

The Past Does Not Matter

At this point, we want to find a way to simplify the policy gradient formula. It’s almost in a form that we can work with. We’ll show how we can effectively set the “past” portion of the above equation to 0, thus removing a lot of variance.

From probability, we know that the expectation of a sum equals the sum of expectations (known as the linearity of expectation rule). As a result,  we can rewrite the past plus future policy gradient as follows:

\(\nabla J(\theta) = \sum_{t} \left( \mathbb{E}_{\tau \sim \pi_\theta}\left[  \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) \right] \;+\; \mathbb{E}_{\tau \sim \pi_\theta}\left[ \gamma^t \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G_t  \right] \right)\)

Let’s now isolate just that past expected value:

\(\mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) \right]\)

The ∇_θ log π(a_t|s_t; θ) term does not depend on j (the timesteps we’re summing over from the beginning of the episode to the t split point), which means we can rewrite the past expected value as follows:

\(\mathbb{E}_{\tau \sim \pi_\theta}\left[ \sum_{j=1}^{t} \gamma^{j-1} \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \right]\)

We then use the same linearity of expectation rule to rewrite the above equation again with the summation on the outside. We can also pull out the γ ^j-1 term, as it’s a constant value (i.e. not random).

\(\sum_{j=1}^{t} \gamma^{j-1} \cdot \mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \right]\)

Let’s focus on just that inner expected value for a moment, assuming some arbitrary j (with j < t). We can rewrite the expected value as a sum of all possible prefixes (possible combinations of trajectories up to t), where each element is a sum of all possible continuations after that prefix (we’ll call that the “continuation” after t). 

\(\mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \right] = \sum_{\text{prefix}} p(\text{prefix}) \sum_{\text{continuation}} p(\text{continuation} \mid \text{prefix}) \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j\)

That inner (continuation) summation is a sum over outcomes, weighted by their probabilities, which is the definition of an expected value. So, we can rewrite it as follows:

\(\mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \right] = \sum_{\text{prefix}} p(\text{prefix}) \cdot \mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \,\big|\, \text{prefix} \big]\)

The outer summation is also a sum over outcomes, weighted by their probabilities. As a result, this becomes another expected value:

\(\mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \right] = \mathbb{E}_{\text{prefix}}\Big[\, \mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \,\big|\, \text{prefix} \big] \,\Big]\)

Now, let’s just focus on that inner piece:

\(\mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \,\big|\, \text{prefix} \big] = \sum_{\text{continuation}} p(\text{continuation} \mid \text{prefix}) \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j\)

Remember that we’re only looking inside the prefix part of the equation at the moment. As a result, each r_j is already determined. It’s not a random variable in this section. As a result, it can be factored out of this inner sum.

\(\sum_{\text{continuation}} p(\text{continuation} \mid \text{prefix}) \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j = r_j \cdot \sum_{\text{continuation}} p(\text{continuation} \mid \text{prefix}) \cdot \nabla_\theta \log \pi(a_t|s_t;\theta)\)

The term inside the summation is still an expected value, so we can write it as follows:

\(\mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \,\big|\, \text{prefix} \big] = r_j \cdot \mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \,\big|\, \text{prefix} \big]\)

Let’s keep zooming in. That inner expected value term can be written out formally (using the definition of expected value) as a sum over possible gradients, weighted by the probability of the policy taking each action from s_t. Remember that t is our arbitrary split point between the prefix and the continuation of the trajectory.

\(\mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \,\big|\, \text{prefix} \big] = \sum_{a_t} \pi(a_t|s_t;\theta) \cdot \nabla_\theta \log \pi(a_t|s_t;\theta)\)

Back in post 12, we showed how the chain rule can be used to define the gradient of the policy:

\(\nabla_\theta \pi(a_t|s_t;\theta) = \pi(a_t|s_t;\theta) \cdot \nabla_\theta \log \pi(a_t|s_t;\theta)\)

As a result, we can substitute this identity back into our expected value:

\(\mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \,\big|\, \text{prefix} \big] = \sum_{a_t} \nabla_\theta \pi(a_t|s_t;\theta)\)

As we’ve done before, we can distribute the gradient across the sum:

\(\sum_{a_t} \nabla_\theta \pi(a_t|s_t;\theta) = \nabla_\theta \left[ \sum_{a_t} \pi(a_t|s_t;\theta) \right]\)

Take a careful look at the summation now. It is asking us to sum across all probabilities of choosing action a given state s. That is a simple probability distribution. By definition, all the probabilities in a distribution must sum to 1.

\(\sum_{a_t} \pi(a_t|s_t;\theta) = 1\)

The derivative (and, by extension, the gradient) of a constant is 0.

\(\nabla_\theta [1] = 0\)

We’ve now shown that this expected value is 0 for any prefix. Combined with the r_j factoring from the earlier step, this means the entire conditional expectation is zero as well.

\(\mathbb{E}\big[ \nabla_\theta \log \pi(a_t|s_t;\theta) \,\big|\, \text{prefix} \big] = \sum_{a_t} \pi(a_t|s_t;\theta) \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) = \nabla_\theta \left[ \sum_{a_t} \pi(a_t|s_t;\theta) \right] = 0\)

We’re now ready to zoom all the way back out to show how we only need to worry about the return-to-go G_t when using the policy gradient.

Putting It All Together

At this point, we showed that for j < t, we have:

\(\mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \right] = 0\)

Working backward through our derivation, we can now substitute 0 in for this term wherever it shows up, which means that the inner expected value over the prefix (timesteps before t) must be 0.

\(\mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot \left(\sum_{j=1}^{t} \gamma^{j-1} r_j\right) \right] = \sum_{j=1}^{t} \gamma^{j-1} \cdot \mathbb{E}_{\tau \sim \pi_\theta}\left[ \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot r_j \right] = \sum_{j=1}^{t} \gamma^{j-1} \cdot 0 = 0\)

This sets the entire expected value prefix term of the policy gradient to 0.

\(\nabla J(\theta) = \sum_{t} \left( 0 + \mathbb{E}_{\tau \sim \pi_\theta}\left[ \gamma^t \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G_t \right] \right)\)

We can then pull the expected value out to get the simplified definition:

\(\nabla J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[ \sum_{t} \gamma^t \cdot \nabla_\theta \log \pi(a_t|s_t;\theta) \cdot G_t \right]\)

That was a lot of work to show how the past terms can effectively be ignored when looking at the policy gradient. This means that we only need to use the return-to-go G_t in our calculations, which helps to keep the variance low, and in turn, allows us to rely on few samples (i.e. sample efficiency) to make policy improvements.

REINFORCE Algorithm

Now that we have a complete understanding of the policy gradient and how to apply the causality trick, we can finally demonstrate an algorithm that uses it. Created by Ronald Williams in 1992 (see this paper), it relies on Monte Carlo sampling to estimate the policy gradient in order to update the parameters, θ, in a neural network. Remember that the neural network is the policy, as we showed in post 12.

We won’t spend much time analyzing REINFORCE. We show it as a simple instantiation of a policy gradient method as we continue our march toward PPO.

Initialize:
    θ: policy parameters, initialized arbitrarily (e.g. random NN weights)
    α: learning rate
    γ: discount factor

Loop forever (for each episode):
    Generate an episode S0, A0, R1, S1, A1, R2, ..., ST-1, AT-1, RT following π(·|·;θ)
    Loop for each step t = 0, 1, ..., T-1:
        G ← Σ{k=t+1..T} [γk-t-1 · Rk]
        θ ← θ + α · γ^t · G · ∇θ log π(At|St;θ)

Note that G given here is exactly G_t. It is the return-to-go from any given timestep t, computing the total (discounted) rewards from t+1 to the episode termination at T. We can leave out the rewards prior to t, thanks to the causality trick we derived.

You can also see that we update the parameters at each timestep in the episode. There are other versions of REINFORCE that accumulate the gradient across the whole episode before performing a single update at the end, but the classic text book (i.e. Sutton & Barto) example performs one update per timestep.

REINFORCE is on-policy, as the episode is generated using the current policy π(·|·;θ) prior to the update. You cannot use an older or different policy (unlike Q-learning’s off-policy approach), as the policy gradient update calculations require samples from the current policy.

The algorithm suffers from a number of limitations that prevented it from seeing widespread use:

• Even with the causality trick, REINFORCE sees a lot of variance. Each gradient update is built from a single sampled episode, which can vary wildly between episodes.
• A full Monte Carlo episode is required before updates can be made, which means REINFORCE (as presented) cannot be used for continuing (non-terminating) tasks.
• The on-policy nature means that you can’t use older experience from a replay buffer.
• Like all gradient-based methods, there is no guarantee of finding a globally optimal policy. You can only hope for a good enough local optimum.
• The base algorithm shown calculates an action’s update purely by the absolute value of G_t with no comparison to the expected value from a given state.

That last bullet point is important, as it provides the framing for our next post. We will establish how a baseline works as a reference point for comparing the actual discounted return G_t from a state to the expected value from that state. This becomes a crucial comparison in the advantage function, which is another building block in PPO.

Conclusion

We spent a good amount of time looking at one (seemingly) minor detail in policy gradient methods: the causality trick that allows us to use the return-to-go G_t instead of the full episodic return G(τ) when computing (or estimating) the policy gradient. The lesson here is that reducing variance can help policies converge more quickly with fewer sampled episodes. In the next post, we’ll examine how the advantage function builds on this idea, using a baseline to make each update more informative.

If you have any questions or comments, please let me know!