Allow reduced latent sde dim

metastable_baselines/sac/policies.py

@@ -69,6 +69,7 @@ class Actor(BasePolicy):
         clip_mean: float = 2.0,
         normalize_images: bool = True,
         dist_kwargs={},
+        latent_dim_sde=None,
     ):
         super().__init__(
             observation_space,
@@ -90,6 +91,8 @@ class Actor(BasePolicy):
         self.full_std = full_std
         self.clip_mean = clip_mean
 
+        self.latent_dim_sde = latent_dim_sde
+
         if sde_net_arch is not None:
             warnings.warn(
                 "sde_net_arch is deprecated and will be removed in SB3 v2.4.0.", DeprecationWarning)
@@ -111,7 +114,7 @@ class Actor(BasePolicy):
             self.action_dist = UniversalGaussianDistribution(
                 action_dim, **dist_kwargs)
             self.mu_net, self.chol_net = self.action_dist.proba_distribution_net(
-                latent_dim=last_layer_dim, latent_sde_dim=last_layer_dim, std_init=math.exp(
+                latent_dim=last_layer_dim, latent_sde_dim=self.latent_dim_sde or last_layer_dim, std_init=math.exp(
                     self.log_std_init)
             )
             # self.action_dist = StateDependentNoiseDistribution(
@@ -274,6 +277,7 @@ class SACPolicy(BasePolicy):
         n_critics: int = 2,
         share_features_extractor: bool = True,
         dist_kwargs={},
+        latent_dim_sde=None,
     ):
         super().__init__(
             observation_space,
@@ -316,6 +320,7 @@ class SACPolicy(BasePolicy):
         }
 
         self.actor_kwargs.update(sde_kwargs)
+        self.actor.kwargs.update({'latent_dim_sde': latent_dim_sde})
         self.critic_kwargs = self.net_args.copy()
         self.critic_kwargs.update(
             {