Allow reduced latent sde dim

2023-01-27 13:34:28 +01:00
parent f421dc2ab5
commit ffbf2b3fe5
3 changed files with 10 additions and 2 deletions
@@ -100,6 +100,7 @@ class ActorCriticPolicy(BasePolicy):
        optimizer_kwargs: Optional[Dict[str, Any]] = None,
        dist_kwargs: Optional[Dict[str, Any]] = None,
        sqrt_induced_gaussian: bool = False,
+        latent_dim_sde=None,
    ):

        if optimizer_kwargs is None:
@@ -155,6 +156,7 @@ class ActorCriticPolicy(BasePolicy):
        self.dist_kwargs = dist_kwargs

        self.sqrt_induced_gaussian = sqrt_induced_gaussian
+        self.latent_dim_sde = latent_dim_sde

        # Action distribution
        self.action_dist = make_proba_distribution(
@@ -244,7 +246,7 @@ class ActorCriticPolicy(BasePolicy):
                latent_dim=latent_dim_pi)
        elif isinstance(self.action_dist, UniversalGaussianDistribution):
            self.action_net, self.chol_net = self.action_dist.proba_distribution_net(
-                latent_dim=latent_dim_pi, latent_sde_dim=latent_dim_pi, std_init=math.exp(
+                latent_dim=latent_dim_pi, latent_sde_dim=self.latent_dim_sde or latent_dim_pi, std_init=math.exp(
                    self.log_std_init)
            )
        else: