Calculating gradient for log policy when variance is not constant

Question

I've noticed that when modelling a continuous action space, the default thing to do is to estimate a mean and a variance where each is parameterized by a neural network or some other model.

I also often see that it is one network $\theta$ models both. The REINFORCE objective can be written as

$$\nabla \mathcal{J}(\theta) = \mathbb{E}_{\pi} [\nabla_\theta \log \pi(a_t|s_t) * R_t] $$

For discrete action space this makes sense since the output of the network is determined by a softmax. However, if we explicitly model the output of the network as a Gaussian, then the gradient of the log likelihood is of a different form,

$$\pi_\theta(a_t|s_t) = Normal(\mu_\theta(s_t), \Sigma_\theta(s_t))$$

and the log is:

$$\log \pi_\theta(a_t | s_t) = -\frac{1}{2} (a_t-\mu_\theta)^\top \Sigma^{-1}_\theta(a_t-\mu_\theta) + \log 2 \pi \det({\Sigma_\theta})$$

In the slides provided here (slide 18): http://www0.cs.ucl.ac.uk/staff/d.silver/web/Teaching_files/pg.pdf

IF the variance is held constant, then we can solve this analytically:

$$\nabla_\theta \log \pi_\theta(a_t|s_t) = (a_t - \mu_\theta) \Sigma^{-1} \phi(s)$$

But, are things always modelled assuming a constant variance? If it's not constant then we have to account for the inverse of the covariance matrix as well as the determinant?

I've taken a look at code online and from what I've seen, most of them assume the variance is constant.

---

@NielSlater

Using the reparameterization trick we would use a normal distribution with fixed parameters 0 and 1. $$ a_t \sim \mu_\theta(s_t) + \Sigma_\theta(s_t) * Normal(0, 1) $$

Which is the same as if we had actually sampled directly from a distribution, $ \pi_\theta(a_t | s_t) = Normal(\mu_\theta(s_t), \Sigma_\theta(s_t))$ and let's us calculate the corresponding $\log \pi_\theta(a_t|s_t)$ and $\nabla_\theta \log \pi_\theta(a_t | s_t)$ without having to differentiate through the actual density.

Answer 1

It's true that computing the log prob of a sample from a gaussian requires inverting a matrix and dealing with the determinant in the general case of a full covariance matrix. If you wanted to backprop through the log prob, you'd need to backprop through these operations, which by the way are available as differentiable operations in pytorch. See the gaussian distribution class which deals with these operations: https://github.com/pytorch/pytorch/blob/master/torch/distributions/multivariate_normal.py#L177 You can see there's a half_log_det and although you won't find an explicit matrix inversion, you'll see that it uses a differentiable triangular linear solve (torch.trtrs) at some point, which is essentially doing a matrix inversion in a more efficient way.