make obs batch require grads as requirement for loss

Denys88 · Oct 29, 2024 · 633bd4c · 633bd4c
1 parent 32c47df
commit 633bd4c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/rl_games/algos_torch/a2c_continuous.py b/rl_games/algos_torch/a2c_continuous.py
@@ -112,6 +112,9 @@ def calc_gradients(self, input_dict):
         lr_mul = 1.0
         curr_e_clip = self.e_clip
 
+        # set requires_grad to True for gradient penalty loss
+        obs_batch = obs_batch.requires_grad_(True)
+
         batch_dict = {
             'is_train': True,
             'prev_actions': actions_batch, 

diff --git a/rl_games/algos_torch/a2c_discrete.py b/rl_games/algos_torch/a2c_discrete.py
@@ -137,6 +137,9 @@ def calc_gradients(self, input_dict):
         lr_mul = 1.0
         curr_e_clip = lr_mul * self.e_clip
 
+        # set requires_grad to True for gradient penalty loss
+        obs_batch = obs_batch.requires_grad_(True)
+
         batch_dict = {
             'is_train': True,
             'prev_actions': actions_batch,