Hello there,
I need to use - or implement - a means of calculating the probability density function of a diagonal, multivariate Gaussian distribution.
The context is that I am building an “Active Inference” agent based model to solve a version of the CartPole-v1 environment: Cart Pole - Gymnasium Documentation
The agent receives a noisy observation vector and must operate on the basis of this noisy/imperfect observation.
I will omit the details as to what my method actually is, I’ll only say that it necessitates the ability to calculate the probability density function (pdf) of a diagonal multivariate Gaussian distribution. I know that PyTorch offers this functionality out of the box with the MultivariateNormal class: Probability distributions - torch.distributions — PyTorch 2.0 documentation. Indeed, I have attempted to use this class, in addition to my own “from scratch” implementation, but I get the same problem in any case.
After importing:
import torch
import torch.distributions.multivariate_normal as mvn
import numpy as np
My implementation is as follows - note that the print statements are for debugging only:
def diag_mvg_pdf_torch(samples, means, diag_vars):
"""
Compute the PDF of a diagonal multivariate Gaussian distribution using PyTorch.
Parameters:
- samples (torch.Tensor): The samples tensor (batch_size x num_features).
- means (torch.Tensor): The means tensor (batch_size x num_features).
- diag_vars (torch.Tensor): The diagonal elements of the covariance matrix (batch_size x num_features).
Returns:
- pdf (torch.Tensor): The average PDF value over the batch.
"""
batch_size = len(samples)
final_pdf = 0
for i in range(batch_size):
mvn_dist_i = mvn.MultivariateNormal(
loc=means[i],
covariance_matrix=torch.diag(diag_vars[i])
)
pdf_i = torch.exp(mvn_dist_i.log_prob(samples[i]))
final_pdf += pdf_i
print(f"\npdf_i: {pdf_i}")
print(f"samples[i]: {samples[i]}")
print(f"means[i]: {means[i]}")
print(f"diag_vars[i]: {diag_vars[i]}\n")
# Compute the average PDF over the batch
pdf = final_pdf / batch_size
return pdf
This function is supposed to instantiate a multivariate Gaussian pdf, parameterised by the mean vector and variances in each row of “means” and “diag_vars” (respectively). The pdf value of each respective sample in “samples” is computed and the average value of these pdfs should be returned.
The batch size is 64, so this code should instantiate 64 diagonal multivariate Gaussians, compute the pdf of each sample with respect to exactly one, unique, instantiated Gaussian and then return the average of these pdf values.
My “samples”, “means” and “diag_vars” tensors look like this - I’ve concatenated their full extent to save space:
# reparameterized samples
samples = torch.tensor([
[ 0.0317, 0.0163, -0.0062, 0.0694],
[ 0.1244, -0.1245, 0.0240, 0.0954],
...
[-0.1707, 0.4822, 0.0498, 0.1098],
[ 0.1741, 0.0472, 0.1415, 0.1108]
])
# predicted mean vectors
means = torch.tensor([
[-0.0248, 0.1706, 0.0314, 0.1002],
[-0.0248, 0.1706, 0.0314, 0.1002],
...
[ 0.0237, 0.1656, 0.0014, 0.1111],
[ 0.0237, 0.1656, 0.0014, 0.1111]
])
# predicted variances
diag_vars = torch.tensor([
[8.1372e-03, 2.6815e-02, 3.7143e-02, 1.4694e-03],
[8.1372e-03, 2.6815e-02, 3.7143e-02, 1.4694e-03],
...
[1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07],
[1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07]
])
Now the output of this following code:
pdf = diag_mvg_pdf_torch(samples, means, diag_vars)
print(f"\npdf: {pdf}\n")
is this:
...
pdf_i: 1966.2965087890625
samples[i]: tensor([0.0185, 0.1974, 0.2239, 0.1114])
means[i]: tensor([0.0237, 0.1656, 0.0014, 0.1111])
diag_vars[i]: tensor([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])
pdf_i: 39.062076568603516
samples[i]: tensor([ 0.1559, 0.2052, -0.1978, 0.1238])
means[i]: tensor([0.0253, 0.1559, 0.0055, 0.1171])
diag_vars[i]: tensor([2.0956e-02, 6.7220e-02, 2.8870e-02, 9.5371e-06])
pdf_i: 250.84556579589844
samples[i]: tensor([-0.1707, 0.4822, 0.0498, 0.1098])
means[i]: tensor([0.0237, 0.1656, 0.0014, 0.1111])
diag_vars[i]: tensor([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])
pdf_i: 1444.880859375
samples[i]: tensor([0.1741, 0.0472, 0.1415, 0.1108])
means[i]: tensor([0.0237, 0.1656, 0.0014, 0.1111])
diag_vars[i]: tensor([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])
pdf: 485.08538818359375
where this is only the last few lines of the output.
Clearly this is not working correctly, since the final “averaged” pdf value is about 485. Indeed, each of the 64 “pdf_i” values are usually vastly larger than 1.0.
I have tried several things, like implementing the pdf calculation “from scratch”:
def diag_mvg_pdf_from_scratch(samples, means, diag_vars):
"""
Compute the PDF of a diagonal multivariate Gaussian distribution using PyTorch.
Parameters:
- samples (torch.Tensor): The samples tensor (batch_size x num_features).
- means (torch.Tensor): The means tensor (batch_size x num_features).
- diag_vars (torch.Tensor): The diagonal elements of the covariance matrix (batch_size x num_features).
Returns:
- pdf (torch.Tensor): The average PDF value over the batch.
"""
batch_size = len(samples)
final_pdf = 0
for i in range(batch_size):
exponent = (-(samples[i] - means[i]) ** 2)/(2 * diag_vars[i])
normalization_term = (2*np.pi*diag_vars[i]) ** (-1/2)
pdf_i_dims = normalization_term * torch.exp(exponent)
pdf_i = torch.prod(pdf_i_dims)
print(f"\nexponent: {exponent}")
print(f"normalization_term: {normalization_term}")
print(f"pdf_i_dims: {pdf_i_dims}")
print(f"pdf_i: {pdf_i}")
print(f"samples[i]: {samples[i]}")
print(f"means[i]: {means[i]}")
print(f"diag_vars[i]: {diag_vars[i]}\n")
# Compute the average PDF over the batch
pdf = final_pdf / batch_size
return pdf
but this has exactly the same result.
I suspect that the issue will have to do with the “samples”, “means” and “log_vars” themselves, since “means” and “log_vars” are actually predicted by a feed-forward neural network. “samples” are samples drawn from the this distribution, by means of the “reparameterization trick”.
I would appreciate any thoughts/assistance that anyone might have to offer. I’m very happy to elaborate on anything that you desire.
Many thanks for reading this far!