Fixing sde's bugs

metastable_baselines/distributions/distributions.py

@@ -143,7 +143,7 @@ class UniversalGaussianDistribution(SB3_Distribution):
     :param action_dim:  Dimension of the action space.
     """
 
-    def __init__(self, action_dim: int, use_sde: bool = False, neural_strength: Strength = Strength.DIAG, cov_strength: Strength = Strength.DIAG, parameterization_type: ParametrizationType = ParametrizationType.NONE, enforce_positive_type: EnforcePositiveType = EnforcePositiveType.ABS, prob_squashing_type: ProbSquashingType = ProbSquashingType.NONE, epsilon=1e-6):
+    def __init__(self, action_dim: int, use_sde: bool = False, neural_strength: Strength = Strength.DIAG, cov_strength: Strength = Strength.DIAG, parameterization_type: ParametrizationType = ParametrizationType.NONE, enforce_positive_type: EnforcePositiveType = EnforcePositiveType.ABS, prob_squashing_type: ProbSquashingType = ProbSquashingType.NONE, epsilon=1e-6, sde_learn_features=False, full_sde=None):
         super(UniversalGaussianDistribution, self).__init__()
         self.action_dim = action_dim
         self.par_strength = cast_to_enum(neural_strength, Strength)
@@ -161,6 +161,10 @@ class UniversalGaussianDistribution(SB3_Distribution):
         self.gaussian_actions = None
 
         self.use_sde = use_sde
+        self.learn_features = sde_learn_features
+
+        if full_sde != None:
+            print('[!] Argument full_sde is only provided to remain compatible with vanilla SB3 PPO. It does not serve any function!')
 
         assert (self.par_type != ParametrizationType.NONE) == (
             self.cov_strength == Strength.FULL), 'You should set an ParameterizationType iff the cov-strength is full'