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
@@ -348,6 +348,7 @@ class UniversalGaussianDistribution(SB3_Distribution):

    def get_noise(self, latent_sde: th.Tensor) -> th.Tensor:
        latent_sde = latent_sde if self.learn_features else latent_sde.detach()
+        latent_sde = latent_sde[..., -self.latent_sde_dim:]
        latent_sde = th.nn.functional.normalize(latent_sde, dim=-1)
        # Default case: only one exploration matrix
        if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices):