Moves all batch info into no_grad

pfnet · Apr 13, 2024 · 2842cb6 · 2842cb6
1 parent b29533b
commit 2842cb6
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 22 deletions.
diff --git a/pfrl/agents/ddpg.py b/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

diff --git a/pfrl/agents/soft_actor_critic.py b/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)

diff --git a/pfrl/agents/td3.py b/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()
             )