Bug fix in SAC with constant ent coeff + try batch sde matrices

torchy_baselines/common/distributions.py

@@ -290,7 +290,7 @@ class StateDependentNoiseDistribution(Distribution):
         # Reduce the number of parameters:
         return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
 
-    def sample_weights(self, log_std):
+    def sample_weights(self, log_std, batch_size=1):
         """
         Sample weights for the noise exploration matrix,
         using a centered gaussian distribution.
@@ -300,6 +300,7 @@ class StateDependentNoiseDistribution(Distribution):
         std = self.get_std(log_std)
         self.weights_dist = Normal(th.zeros_like(std), std)
         self.exploration_mat = self.weights_dist.rsample()
+        self.exploration_matrices = self.weights_dist.rsample((batch_size,))
 
     def proba_distribution_net(self, latent_dim, log_std_init=-2.0, latent_sde_dim=None):
         """
@@ -354,7 +355,13 @@ class StateDependentNoiseDistribution(Distribution):
 
     def get_noise(self, latent_sde):
         latent_sde = latent_sde if self.learn_features else latent_sde.detach()
-        return th.mm(latent_sde, self.exploration_mat)
+        if len(latent_sde) != len(self.exploration_matrices):
+            return th.mm(latent_sde, self.exploration_mat)
+        # (batch_size, n_features) -> (batch_size, 1, n_features)
+        latent_sde = latent_sde.unsqueeze(1)
+        # (batch_size, 1, n_actions)
+        noise = th.bmm(latent_sde, self.exploration_matrices)
+        return noise.squeeze(1)
 
     def sample(self, latent_sde):
         noise = self.get_noise(latent_sde)

torchy_baselines/sac/policies.py

@@ -53,11 +53,13 @@ class Actor(BaseNetwork):
         """
         return self.action_dist.get_std(self.log_std)
 
-    def reset_noise(self):
+    def reset_noise(self, batch_size=1):
         """
         Sample new weights for the exploration matrix, when using SDE.
+
+        :param batch_size: (int)
         """
-        self.action_dist.sample_weights(self.log_std)
+        self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
 
     def get_action_dist_params(self, obs):
         latent = self.actor_net(obs)

torchy_baselines/sac/sac.py

@@ -87,6 +87,7 @@ class SAC(BaseRLModel):
         self.n_episodes_rollout = n_episodes_rollout
         self.action_noise = action_noise
         self.gamma = gamma
+        self.ent_coef_optimizer = None
 
         if _init_setup_model:
             self._setup_model()
@@ -124,7 +125,7 @@ class SAC(BaseRLModel):
             # Force conversion to float
             # this will throw an error if a malformed string (different from 'auto')
             # is passed
-            self.ent_coef = float(self.ent_coef)
+            self.ent_coef = th.tensor(float(self.ent_coef)).to(self.device)
 
         self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
         self.policy = self.policy(self.observation_space, self.action_space, learning_rate=self.learning_rate,
@@ -174,15 +175,15 @@ class SAC(BaseRLModel):
             # or sample again the noise matrix
             # otherwise the intermediate step `std = th.exp(log_std)`
             # is lost and we cannot backpropagate through again
-            # if self.use_sde:
-            #     self.actor.reset_noise()
+            if self.use_sde:
+                self.actor.reset_noise(batch_size=batch_size)
 
             # Action by the current actor for the sampled state
             action_pi, log_prob = self.actor.action_log_prob(obs)
             log_prob = log_prob.reshape(-1, 1)
 
             ent_coef_loss = None
-            if not isinstance(self.ent_coef, float):
+            if self.ent_coef_optimizer is not None:
                 # Important: detach the variable from the graph
                 # so we don't change it with other losses
                 # see https://github.com/rail-berkeley/softlearning/issues/60
@@ -200,6 +201,8 @@ class SAC(BaseRLModel):
 
 
             with th.no_grad():
+                if self.use_sde:
+                    self.actor.reset_noise(batch_size=batch_size)
                 # Select action according to policy
                 next_action, next_log_prob = self.actor.action_log_prob(next_obs)
                 # Compute the target Q value
@@ -231,8 +234,8 @@ class SAC(BaseRLModel):
             self.actor.optimizer.zero_grad()
             # Cf comment above, otherwise pytorch raises an error
             # ("Trying to backward through the graph a second time")
-            retain_graph = True if self.use_sde and gradient_steps > 1 else False
-            actor_loss.backward(retain_graph=retain_graph)
+            # retain_graph = True if self.use_sde and gradient_steps > 1 else False
+            actor_loss.backward(retain_graph=False)
             self.actor.optimizer.step()
 
             # Update target networks