Reduce variance

What's wrong with the policy gradient ?

Reduce variance

We knew that $c$, and there are many reasons that will cause high variance. Here is the most straightforward one.

Suppose we got a large negative reward and two small positive rewards, according to the update formula to $\theta$, the distribution of $\tau$ will move to the right a lot. However, if we add a constant to both positive and negative reward, here we have a small positive reward and two large positive rewards, the result is that the distribution of $\tau$ will move to the right a little bit.

From the above illustration, a sight change to rewards will severely influence the value of $\nabla_\theta J(\theta)$, which will cause high variance.

The worst case is that, if the two good samples have $r(\tau)=0$, it may take a long to converge or end in a sub-optimal solution.

Causality

$$\nabla_\theta J(\theta) =\frac{1}{N}\sum_{i=1}^N \left[ \left(\sum_{t=1}^T \nabla_\theta \log\pi_\theta (a_{i,t}|s_{i,t}) \right) \left(\sum_{t=1}^T r(s_{i,t},a_{i,t}) \right) \right]$$

The causality is that policy at time $t'$ cannot affect reward at time $t$ when $t<t'$. So the gradient should be

$$\nabla_\theta J(\theta) =\frac{1}{N}\sum_{i=1}^N \left[ \sum_{t=1}^T \nabla_\theta \log\pi_\theta (a_{i,t}|s_{i,t}) \left(\sum_{t'=t}^T r(s_{i,t'},a_{i,t'}) \right) \right] =\frac{1}{N}\sum_{i=1}^N \left[ \sum_{t=1}^T \nabla_\theta \log\pi_\theta (a_{i,t}|s_{i,t}) \hat{Q}_{i,t} \right]$$

The $\hat{Q}_{i,t}$ is reward to go from time $t$ for sample $i$. Since $t'=1$ becomes $t'=t$ , the number of rewards reduces, which leads to lower variance.

The causality always works well, so can be used every time.

Baseline

$$\nabla_\theta J(\theta)\approx\frac{1}{N}\sum_i^N \nabla_\theta \log \pi_\theta (\tau_i)r(\tau_i)$$

Our purpose is to make good trajectory have bigger probability and bad trajectory have smaller probability. The problem is that good trajectories don't always have big positive rewards and bad trajectories aren't always negative. Our methods is to subtract a baseline like that,

$$\nabla_\theta J(\theta)\approx\frac{1}{N}\sum_i^N \nabla_\theta \log \pi_\theta (\tau_i)(r(\tau_i)-b)$$

where $b=\frac{1}{N}\sum_{i=1}^N r(\tau_i)$

Actually, subtracting a baseline is unbiased in expectation

$$E[\nabla_\theta\log\pi_\theta(\tau)b]=\int \pi_\theta(\tau)\nabla_\theta\log\pi_\theta(\tau)bd\tau=\int \nabla_\theta \pi_\theta(\tau) bd\tau=b\nabla_\theta\int \pi_\theta(\tau)d\tau=b\nabla_\theta1=0$$

The last thing worth mentioning is that average reward is not the best baseline, but it's pretty good.

Analyzing variance

Subtracting a baseline is unbiased, but can we find the best baseline that has the lowest variance ?

$$\nabla_\theta J(\theta)=E_{\tau\sim \pi_\theta(\tau)}[\nabla_\theta\log\pi_\theta (\tau)(r(\tau)-b)]$$

And the variance

$$\begin{aligned} \text{Var} &= E[x^2]-E[x]^2 \\ &=E_{\tau\sim \pi_\theta(\tau)}[(\nabla_\theta\log\pi_\theta (\tau)(r(\tau)-b))^2]-E_{\tau\sim \pi_\theta(\tau)}[\nabla_\theta\log\pi_\theta (\tau)(r(\tau)-b)]^2 \\ &=E_{\tau\sim \pi_\theta(\tau)}[(\nabla_\theta\log\pi_\theta (\tau)(r(\tau)-b))^2]-E_{\tau\sim \pi_\theta(\tau)}[\nabla_\theta\log\pi_\theta (\tau)r(\tau)]^2 \\ \end{aligned}$$

The second equation is because baseline is unbiased in expectation.

Denote $g(\tau)=\nabla_\theta\log\pi_\theta (\tau)$

$$\begin{aligned} \frac{d\text{Var}}{db} &=\frac{d}{db}E[g(\tau)^2(r(\tau)-b)^2] \\ &=\frac{d}{db}\left(E[g(\tau)^2r(\tau)^2]-2E[g(\tau)^2r(\tau)b]+b^2E[g(\tau)] \right) \\ &=\frac{d}{db}\left(-2E[g(\tau)^2r(\tau)b]+b^2E[g(\tau)] \right) \\ &=-2E[g(\tau)^2r(\tau)]+2bE[g(\tau)] =0 \end{aligned}$$

So $b$ should be

$$b=\frac{E[g(\tau)^2r(\tau)]}{E[g(\tau)}$$

This is just expected reward, but weighted by gradient magnitudes.

In theory, the best baseline should be weighted expected reward, but in practice, we haven't find any difference between weighted expected reward and average reward. And since average have smaller computation, we just use average reward as baseline.