Fixed numerical issues with Wasserstein

2022-08-17 23:25:24 +02:00
parent 9fffe048af
commit a9e3f295b2
2 changed files with 5 additions and 0 deletions
@@ -232,8 +232,11 @@ class UniversalGaussianDistribution(SB3_Distribution):

        if nobatch:
            cov = th.mm(cov_sqrt.mT, cov_sqrt)
+            cov += th.eye(cov.shape[-1])*(self.epsilon)
        else:
            cov = th.bmm(cov_sqrt.mT, cov_sqrt)
+            cov += th.eye(cov.shape[-1]).expand(cov.shape)*(self.epsilon)
+
        chol = th.linalg.cholesky(cov)

        if vec: