Fix issues when providing diagonal matrix sqrt

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