やったこと
pytorchはv1.5.1を使用．
Linearの出力を(x, y, vxx, vyy, vxy)として，分散共分散行列を作りMultivariateNormalに渡す．
fc = nn.Linear(n, 5)
output = fc(x)
mean = output[:2]
vxx, vyy = nn.Softplus()(output[2:4])
vxy = output[-1]
covariance_matrix = torch.zeros(2, 2)
covariance_matrix[0, 0] += vxx
covariance_matrix[0, 1] += vxy
covariance_matrix[1, 0] += vxy
covariance_matrix[1, 1] += vyy
dist = MultivariateNormal(mean, covariance_matrix)
RuntimeError: cholesky_cuda: For batch 0: U(6,6) is zero, singular U.
MultivariateNormalでは分散共分散行列をコレスキー分解しているため，正定値行列を与える必要がある．このままではcovariance_matrixが正定値行列であることを保証していないので上記エラーが起こる．
解決法: コレスキー分解後の下三角行列を渡す．
fc = nn.Linear(n, 5)
output = fc(x)
mean = output[:2]
a, c = nn.Softplus()(output[2:4])
b = output[-1]
L = torch.zeros(2, 2)
L[0, 0] += a
L[1, 0] += b
L[1, 1] += c
dist = MultivariateNormal(mean, scale_tril=L)
scale_tril (Tensor) – lower-triangular factor of covariance, with positive-valued diagonal
とあるので対角成分の$a, c$はsoftplusで正の値にしている．
分散共分散行列にするときにはコレスキー分解
\Sigma = LL^{T}
に従い，
covariance_matrix = np.dot(L, L.T)
とすればよい．