Moves all batch info into no_grad

Author: prabhatnagarajan

pfrl/agents/ddpg.py

@@ -148,16 +148,16 @@ def sync_target_network(self):
     def compute_critic_loss(self, batch):
         """Compute loss for critic."""
 
-        batch_next_state = batch["next_state"]
-        batch_rewards = batch["reward"]
-        batch_terminal = batch["is_state_terminal"]
-        batch_state = batch["state"]
-        batch_actions = batch["action"]
-        batchsize = len(batch_rewards)
-
         with torch.no_grad():
+            batch_state = batch["state"]
+            batch_actions = batch["action"]
+            batch_rewards = batch["reward"]
+            batchsize = len(batch_rewards)
             assert not self.recurrent
+            batch_next_state = batch["next_state"]
+            batch_terminal = batch["is_state_terminal"]
             next_actions = self.target_policy(batch_next_state).sample()
+            
             next_q = self.target_q_function((batch_next_state, next_actions))
             target_q = batch_rewards + self.gamma * (
                 1.0 - batch_terminal

pfrl/agents/soft_actor_critic.py

@@ -213,17 +213,16 @@ def sync_target_network(self):
 
     def update_q_func(self, batch):
         """Compute loss for a given Q-function."""
-
-        batch_next_state = batch["next_state"]
-        batch_rewards = batch["reward"]
-        batch_terminal = batch["is_state_terminal"]
-        batch_state = batch["state"]
-        batch_actions = batch["action"]
-        batch_discount = batch["discount"]
-
         with torch.no_grad(), pfrl.utils.evaluating(self.policy), pfrl.utils.evaluating(
             self.target_q_func1
         ), pfrl.utils.evaluating(self.target_q_func2):
+            batch_state = batch["state"]
+            batch_actions = batch["action"]
+            batch_next_state = batch["next_state"]
+            batch_rewards = batch["reward"]
+            batch_terminal = batch["is_state_terminal"]
+            batch_discount = batch["discount"]
+
             next_action_distrib = self.policy(batch_next_state)
             next_actions = next_action_distrib.sample()
             next_log_prob = next_action_distrib.log_prob(next_actions)

pfrl/agents/td3.py

@@ -181,18 +181,18 @@ def sync_target_network(self):
     def update_q_func(self, batch):
         """Compute loss for a given Q-function."""
 
-        batch_next_state = batch["next_state"]
-        batch_rewards = batch["reward"]
-        batch_terminal = batch["is_state_terminal"]
-        batch_state = batch["state"]
-        batch_actions = batch["action"]
-        batch_discount = batch["discount"]
-
         with torch.no_grad(), pfrl.utils.evaluating(
             self.target_policy
         ), pfrl.utils.evaluating(self.target_q_func1), pfrl.utils.evaluating(
             self.target_q_func2
         ):
+            batch_state = batch["state"]
+            batch_actions = batch["action"]
+            batch_next_state = batch["next_state"]
+            batch_rewards = batch["reward"]
+            batch_terminal = batch["is_state_terminal"]
+            batch_discount = batch["discount"]
+
             next_actions = self.target_policy_smoothing_func(
                 self.target_policy(batch_next_state).sample()
             )