dodox/metastable-baselines
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fixed bug with SDE

Browse Source
This commit is contained in:
Dominik Moritz Roth 2022-08-22 13:36:17 +02:00
parent a9e3f295b2
commit 197de7997c
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
metastable_baselines/distributions/distributions.py
View File

@ -365,7 +365,12 @@ class UniversalGaussianDistribution(SB3_Distribution):
if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices):
chol = th.diag_embed(self.distribution.stddev)
return (th.mm(latent_sde, self.exploration_mat) @ chol)[0]
chol = self.distribution.scale_tril
p = self.distribution
if isinstance(p, th.distributions.Normal) or isinstance(p, th.distributions.Independent):
chol = p.stddev
elif isinstance(p, th.distributions.MultivariateNormal):
chol = p.scale_tril
# Use batch matrix multiplication for efficient computation
# (batch_size, n_features) -> (batch_size, 1, n_features)
latent_sde = latent_sde.unsqueeze(dim=1)