From e49d97bf98a0340f4173d5d8d51b42a6170dc453 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Wed, 18 Dec 2019 13:45:33 +0100 Subject: [PATCH] Fix infs in SAC by bounding the mean --- torchy_baselines/sac/policies.py | 2 ++ 1 file changed, 2 insertions(+) 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)