diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index cd1b35b..610b59e 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -48,6 +48,8 @@ class Actor(BaseNetwork): self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1], latent_sde_dim=latent_sde_dim, log_std_init=log_std_init) + # Avoid saturation by limiting the mean of the gaussian to be in [-1, 1] + self.mu = nn.Sequential(self.mu, nn.Tanh()) else: self.action_dist = SquashedDiagGaussianDistribution(action_dim) self.mu = nn.Linear(net_arch[-1], action_dim)