Now softmax sde latent is actually working...

This commit is contained in:
Dominik Moritz Roth 2023-01-28 19:48:43 +01:00
parent e985d892ca
commit b86498fa9a

View File

@ -351,7 +351,7 @@ class UniversalGaussianDistribution(SB3_Distribution):
latent_sde = latent_sde if self.learn_features else latent_sde.detach() latent_sde = latent_sde if self.learn_features else latent_sde.detach()
latent_sde = latent_sde[..., -self.latent_sde_dim:] latent_sde = latent_sde[..., -self.latent_sde_dim:]
if self.sde_latent_softmax: if self.sde_latent_softmax:
latent_sde = th.softmax(dim=-1) latent_sde = latent_sde.softmax(-1)
latent_sde = th.nn.functional.normalize(latent_sde, dim=-1) latent_sde = th.nn.functional.normalize(latent_sde, dim=-1)
# Default case: only one exploration matrix # Default case: only one exploration matrix
if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices): if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices):