fix errors of SAC and SAC-Discrete caused by torch>=1.4.0

Author: toshikwa

agents/actor_critic_agents/SAC.py

@@ -144,10 +144,12 @@ def learn(self):
         """Runs a learning iteration for the actor, both critics and (if specified) the temperature parameter"""
         state_batch, action_batch, reward_batch, next_state_batch, mask_batch = self.sample_experiences()
         qf1_loss, qf2_loss = self.calculate_critic_losses(state_batch, action_batch, reward_batch, next_state_batch, mask_batch)
+        self.update_critic_parameters(qf1_loss, qf2_loss)
+
         policy_loss, log_pi = self.calculate_actor_loss(state_batch)
         if self.automatic_entropy_tuning: alpha_loss = self.calculate_entropy_tuning_loss(log_pi)
         else: alpha_loss = None
-        self.update_all_parameters(qf1_loss, qf2_loss, policy_loss, alpha_loss)
+        self.update_actor_parameters(policy_loss, alpha_loss)
 
     def sample_experiences(self):
         return  self.memory.sample()
@@ -182,18 +184,21 @@ def calculate_entropy_tuning_loss(self, log_pi):
         alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
         return alpha_loss
 
-    def update_all_parameters(self, critic_loss_1, critic_loss_2, actor_loss, alpha_loss):
-        """Updates the parameters for the actor, both critics and (if specified) the temperature parameter"""
+    def update_critic_parameters(self, critic_loss_1, critic_loss_2):
+        """Updates the parameters for both critics"""
         self.take_optimisation_step(self.critic_optimizer, self.critic_local, critic_loss_1,
                                     self.hyperparameters["Critic"]["gradient_clipping_norm"])
         self.take_optimisation_step(self.critic_optimizer_2, self.critic_local_2, critic_loss_2,
                                     self.hyperparameters["Critic"]["gradient_clipping_norm"])
-        self.take_optimisation_step(self.actor_optimizer, self.actor_local, actor_loss,
-                                    self.hyperparameters["Actor"]["gradient_clipping_norm"])
         self.soft_update_of_target_network(self.critic_local, self.critic_target,
                                            self.hyperparameters["Critic"]["tau"])
         self.soft_update_of_target_network(self.critic_local_2, self.critic_target_2,
                                            self.hyperparameters["Critic"]["tau"])
+
+    def update_actor_parameters(self, actor_loss, alpha_loss):
+        """Updates the parameters for the actor and (if specified) the temperature parameter"""
+        self.take_optimisation_step(self.actor_optimizer, self.actor_local, actor_loss,
+                                    self.hyperparameters["Actor"]["gradient_clipping_norm"])
         if alpha_loss is not None:
             self.take_optimisation_step(self.alpha_optim, None, alpha_loss, None)
             self.alpha = self.log_alpha.exp()