Fixed all the bugs in TRPL

metastable_baselines/distributions/distributions.py

@@ -221,13 +221,19 @@ class UniversalGaussianDistribution(SB3_Distribution):
 
     def _sqrt_to_chol(self, cov_sqrt):
         vec = False
-        if len(cov_sqrt.shape) == 2:
+        nobatch = False
+        if len(cov_sqrt.shape) <= 2:
             vec = True
+        if len(cov_sqrt.shape) == 1:
+            nobatch = True
 
         if vec:
             cov_sqrt = th.diag_embed(cov_sqrt)
 
-        cov = th.bmm(cov_sqrt.mT, cov_sqrt)
+        if nobatch:
+            cov = th.mm(cov_sqrt.mT, cov_sqrt)
+        else:
+            cov = th.bmm(cov_sqrt.mT, cov_sqrt)
         chol = th.linalg.cholesky(cov)
 
         if vec:

metastable_baselines/misc/distTools.py

@@ -26,7 +26,12 @@ def get_mean_and_sqrt(p: UniversalGaussianDistribution, expand=False):
     else:
         mean, chol = get_mean_and_chol(p, expand=False)
         sqrt_cov = p.cov_sqrt
-        if expand and len(sqrt_cov.shape) == 2:
+        if mean.shape[0] != sqrt_cov.shape[0]:
+            shape = list(sqrt_cov.shape)
+            shape[0] = mean.shape[0]
+            shape = tuple(shape)
+            sqrt_cov = sqrt_cov.expand(shape)
+        if expand and len(sqrt_cov.shape) <= 2:
             sqrt_cov = th.diag_embed(sqrt_cov)
         return mean, sqrt_cov
 

metastable_baselines/ppo/policies.py

@@ -315,9 +315,19 @@ class ActorCriticPolicy(BasePolicy):
         elif isinstance(self.action_dist, UniversalGaussianDistribution):
             if self.sqrt_induced_gaussian:
                 chol_sqrt_cov = self.chol_net(latent_pi)
-                if len(chol_sqrt_cov.shape) == 2:
+                unembed = False
+                squeeze = False
+                if len(chol_sqrt_cov.shape) <= 2:
+                    unembed = True
                     chol_sqrt_cov = th.diag_embed(chol_sqrt_cov)
+                if len(chol_sqrt_cov.shape) <= 2:
+                    squeeze = True
+                    chol_sqrt_cov = chol_sqrt_cov.unsqueeze(0)
                 cov_sqrt = th.bmm(chol_sqrt_cov.mT, chol_sqrt_cov)
+                if squeeze and False:
+                    cov_sqrt = cov_sqrt.squeeze()
+                if unembed:
+                    cov_sqrt = th.diagonal(cov_sqrt, dim1=-2, dim2=-1)
                 dist = self.action_dist.proba_distribution_from_sqrt(
                     mean_actions, cov_sqrt, latent_pi)
                 mean, chol = get_mean_and_chol(dist, expand=False)