Fix issues when providing diagonal matrix sqrt

metastable_projections/misc/distTools.py

@@ -27,11 +27,14 @@ def get_mean_and_chol(p: AnyDistribution, expand=False):
 
 
 def get_mean_and_sqrt(p: UniversalGaussianDistribution, expand=False):
+    mean, chol = get_mean_and_chol(p, expand=False)
     if not hasattr(p, 'cov_sqrt'):
+        if len(p.distribution.scale.shape)==2: # In the factorized case, chol = matSqrt
+            sqrt_cov = p.distribution.scale
+        else:
             raise Exception(
                 'Distribution was not induced from sqrt. On-demand calculation is not supported.')
     else:
-        mean, chol = get_mean_and_chol(p, expand=False)
         sqrt_cov = p.cov_sqrt
     if mean.shape[0] != sqrt_cov.shape[0]:
         shape = list(sqrt_cov.shape)
@@ -123,12 +126,14 @@ def new_dist_like_from_sqrt(orig_p: AnyDistribution, mean: th.Tensor, cov_sqrt:
 
 
 def _sqrt_to_chol(cov_sqrt, only_diag=False):
+    if only_diag:
+        return cov_sqrt
     cov = th.bmm(cov_sqrt.mT, cov_sqrt)
     cov += th.eye(cov.shape[-1]).expand(cov.shape)*(1e-6)
 
     chol = th.linalg.cholesky(cov)
 
-    if only_diag:
-        chol = th.diagonal(chol, dim1=-2, dim2=-1)
+    #if only_diag:
+    #    chol = th.diagonal(chol, dim1=-2, dim2=-1)
 
     return chol

metastable_projections/projections/base_projection_layer.py

@@ -86,6 +86,7 @@ class BaseProjectionLayer(object):
         Returns:
             projected_dist, old_dist (from rollouts)
         """
+
         old_distribution = self.new_dist_like(dist, rollout_data.mean, rollout_data.cov_decomp)
         return self(dist, old_distribution, **kwargs), old_distribution
 

metastable_projections/projections/frob_projection_layer.py

@@ -4,7 +4,7 @@ from typing import Tuple
 from .base_projection_layer import BaseProjectionLayer, mean_projection
 
 from ..misc.norm import mahalanobis, frob_sq
-from ..misc.distTools import get_mean_and_chol, get_cov, new_dist_like
+from ..misc.distTools import get_mean_and_chol, get_cov, new_dist_like, has_diag_cov
 
 
 class FrobeniusProjectionLayer(BaseProjectionLayer):
@@ -57,6 +57,9 @@ class FrobeniusProjectionLayer(BaseProjectionLayer):
         else:
             proj_chol = chol
 
+        if has_diag_cov(p):
+            proj_chol = th.diagonal(proj_chol, dim1=-2, dim2=-1)
+
         proj_p = new_dist_like(p, proj_mean, proj_chol)
         return proj_p
 

metastable_projections/projections/w2_projection_layer.py

@@ -63,6 +63,9 @@ class WassersteinProjectionLayer(BaseProjectionLayer):
         else:
             proj_sqrt = sqrt
 
+        if has_diag_cov(p):
+            proj_sqrt = th.diagonal(proj_sqrt, dim1=-2, dim2=-1)
+
         proj_p = new_dist_like_from_sqrt(p, proj_mean, proj_sqrt)
         return proj_p
 
@@ -110,6 +113,7 @@ class WassersteinProjectionLayer(BaseProjectionLayer):
                 mean, scale_tril=cov_sqrt)
         else:
             raise Exception('Dist-Type not implemented (of sb3 dist)')
+        p_out.cov_sqrt = cov_sqrt
         return p_out