[RLlib] Fix no gradient clipping happening in QMix. (#25656)

ray-project · Jun 10, 2022 · c364592 · c364592
1 parent 730df43
commit c364592
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rllib/algorithms/qmix/qmix_policy.py b/rllib/algorithms/qmix/qmix_policy.py
@@ -447,6 +447,7 @@ def to_batches(arr, dtype):
         # Optimise
         self.rmsprop_optimizer.zero_grad()
         loss_out.backward()
+        grad_norm_info = apply_grad_clipping(self, self.rmsprop_optimizer, loss_out)
         self.rmsprop_optimizer.step()
 
         mask_elems = mask.sum().item()
@@ -456,7 +457,8 @@ def to_batches(arr, dtype):
             "q_taken_mean": (chosen_action_qvals * mask).sum().item() / mask_elems,
             "target_mean": (targets * mask).sum().item() / mask_elems,
         }
-        stats.update(apply_grad_clipping(self, self.rmsprop_optimizer, loss_out))
+        stats.update(grad_norm_info)
+
         return {LEARNER_STATS_KEY: stats}
 
     @override(TorchPolicy)