How to enforce covariance-matrix output as part of the last layer of a Policy Network?

Question

I have a continuous state space, and a continuous action space. The way I understand it, I can build a policy network which takes as input a continuous state vector and outputs both mean vector and covariance matrix of the action-distribution. To get a valid action I then sample from that distribution.

However, when trying to implement such a network, I get the error message that the parts of my output layer which I want to be the covariance matrix are singular/not positive-semi-definite. How can I fix this? I tried different activation-functions and initializations for the last layer, but once in a while I run into the same problem again.

How can I enforce that my network outputs a valid covariance matrix?

Answer 1

Usually it is assumed that there is no correlation between different actions, so the covariance matrix will be zero everywhere except on the main diagonal. Diagonal will represent variances of actions. Diagonal covariance matrix will be positive semidefinite if all values on diagonal are $\geq$ 0 so you need to insure that output of final layer is $\geq$ 0, which can be done with ReLU activation for example.

Answer 2

@Brale_ 's answer is correct, it is common practice in a multitude of models to learn a representation of an independent multivariate normal, but don't let that stop you from pushing the envelope for your needs.

You can actually learn a dependent form as well. Normally the independent form is done by learning the means and standard deviations, because sampling a standard normal can achieve your draw by $z \sim N(\mu, Diag(\sigma^2))$ by $z = \mu + \sigma \epsilon$ where $\epsilon$ is drawn from a unit normal.

But you can actually achieve a similar trick for a generalized multivariate normal distribution: lets assume your trying to learn $N(\mu, \Sigma)$ where $\Sigma$ is the covariance matrix. So what you would do is learn the Cholesky decomposition because where $\Sigma = AA^T$ you can now draw it through the parametrization trick: $z = \mu + A\epsilon$.

Answer 3

@mshlis, if $\sigma$ is covariance matrix, then exist $A$ s.t. $\sigma=AA^T$.

But if we generate $A$ by model, and calculate $\sigma = AA^T$, then Cholesky decomposition will not successful