Policy gradients (part 2): using 'probability-until-now'

This post continues my policy gradient series. Hopefully today I will get to how to actually estimate the gradient, and what the variance is.

Politics gradients

Policy gradients is just doing more good actions and less bad actions

First let’s appreciate what the insight we obtain from the policy gradient theorem:

To incrementally improve your policy, increase the probability of actions that resulted in good outcomes, and decrease the probability of actions with bad outcomes.

Take a look at the equation again:

\(\text{Gradient of expected return} = \mathbb{E}_{x_{1:T} \sim \pi_\theta} \left[ G(x_{1:T}) \nabla_\theta \log \pi_\theta\left( x_{1:T} \right) \right]\)

The right-hand side is the average over gradients of log-probabilities. Gradients of log-probabilities are the directions of maximum increase in log-probability; if we move the parameters slightly in whatever direction they’re pointing then the probabilities of the taken actions will go up. If we move parameters in the opposite direction, the probabilities will go down.

In the equation, we weigh those gradients by the return of each trajectory G(x_{1:T}). This means we assign directions that got higher returns more weight in the average, and directions that got lower returns less (or even negative) weights in the average. The net effect is that the averaged direction is the one that maximally increases return.

A different fork in the derivation

Let’s re-do the policy gradients derivation, but without aggregating rewards into the return. We start with the gradient of expected return, renaming the total number of steps T into N to avoid confusion:

\(\nabla_\theta \left(\mathbb{E}_{x_{1: N} \sim \pi_\theta} \left[ \sum_{t=1}^N r(x_{1:t}) \right] \right)\)

By linearity of the expectation and derivatives, we can take the sum outside:

\( \sum_{t=1}^N \nabla_\theta \left(\mathbb{E}_{x_{1: N} \sim \pi_\theta} \left[ r(x_{1:t}) \right] \right)\)

Notice that the value on the inside of the expectation over x_{1:N} only depends on the states until the intermediate time-step t. Thus we can drop the states between t and N:

\( \sum_{t=1}^N \nabla_\theta \left(\mathbb{E}_{x_{1:t} \sim \pi_\theta} \left[ r(x_{1:t}) \right] \right)\)

(The change was in the expectation subscript, x_{1:N} became x_{1:t}.)

Finally, we apply the old policy gradients theorem to the inner part, to get:

\( \sum_{t=1}^N \mathbb{E}_{x_{1:t} \sim \pi_\theta} \left[ r(x_{1:t}) \nabla_\theta \log \pi_\theta(x_t | x_{1:t-1}) \right] \)

At first glance, this expression seems incompatible with the other. We cannot put the sum back inside the expectation, because each r(x_{1:t}) is multiplied by a different factor ∇_θ log π_θ .

However, it turns out they’re the same! In practice, when the expectation is not exact, this version has lower variance.

Why are these two different expressions the same? The baseline trick.

The answer is related to the fact that we can add any constant to the policy gradients return, and the gradient’s expectation remains the same. Going back to the old formulation:

\(\mathbb{E}_{x_{1:T} \sim \pi_\theta} \left[ G(x_{1:T}) \nabla_\theta \log \pi_\theta\left( x_{1:T} \right) \right] = \mathbb{E}_{x_{1:T} \sim \pi_\theta} \left[ \left(G(x_{1:T}) + c \right) \nabla_\theta \log \pi_\theta\left( x_{1:T} \right) \right]\)

The reason for this is that we can convert the constant part back to the gradient of an expectation:

\(\mathbb{E}_{x_{1:T} \sim \pi_\theta} \left[ c \cdot \nabla_\theta \log \pi_\theta\left( x_{1:T} \right) \right] = c \int \pi_\theta(x_{1:T}) \nabla_\theta \log \pi_\theta(x_{1:T}) \textsf{d}x_{1:T}\)

take the policy gradient identity, in reverse:

\(= c \int \nabla_\theta \pi_\theta(x_{1:T}) \textsf{d}x_{1:T}\)

then we interchange the integral and derivate and note that probability densities integrate to 1:

\(= c \,\nabla_\theta \int \pi_\theta(x_{1:T}) \textsf{d}x_{1:T} = c \, \nabla_\theta 1 = 0.\)

So the resulting contribution of adding a constant to the expectation is zero (multiplied by the gradient and such).

How is truncating log p the same as the baseline trick?

Let’s zoom out to proving that the weird expression we derived is the same as the original one. We want to show that we arrive at this:

\( \sum_{t=1}^N \mathbb{E}_{x_{1:t} \sim \pi_\theta} \left[ r(x_{1:t}) \nabla_\theta \log \pi_\theta(x_t | x_{1:t-1}) \right] \)

Starting from the original formulation:

\( \mathbb{E}_{x_{1:N} \sim \pi_\theta} \left[ \left( \sum_{t=1}^N r(x_{1:t})\right) \nabla_\theta \log \pi_\theta(x_{1:N}) \right] \)

we take the sum out of the expectation. Then, using the linearity of gradient and the fact that log-probabilities add, we split off the part of the trajectory after t:

\( \sum_{t=1}^N \mathbb{E}_{x_{1:N} \sim \pi_\theta} \left[ r(x_{1:t}) \left( \nabla_\theta\log \pi_\theta(x_t | x_{1:t-1} )+ \nabla_\theta\log \pi_\theta(x_{t:N} | x_{1:t-1}) \right) \right] \)

We can then use the law of total expectation to write

\( \sum_{t=1}^N \mathbb{E}_{x_{1:t} \sim \pi_\theta} \left[ r(x_{1:t}) \left( \nabla_\theta\log \pi_\theta(x_t | x_{1:t-1} )+ \mathbb{E}_{x_{t+1:N} \sim \pi_\theta(x_{1:t})} \left[ \nabla_\theta\log \pi_\theta(x_{t+1:N} | x_{1:t}) \right] \right) \right] \)

Then, using the derivation from the previous section, we can see that the expectation of the gradient of a log probability is zero! So it cancels out, and we derive the original expression.

In practice, their associated estimators have different variances.