From 919dfee45239ddd2c23291baf722f0ac870c5ad4 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Tue, 17 Dec 2019 11:14:44 +0100 Subject: [PATCH] Try to clip grad norm --- torchy_baselines/sac/sac.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index e170cd0..ee5ee4c 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -202,6 +202,9 @@ class SAC(BaseRLModel): if ent_coef_loss is not None: self.ent_coef_optimizer.zero_grad() ent_coef_loss.backward() + # Clip grad norm + if self.max_grad_norm is not None: + th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.ent_coef_optimizer.step() @@ -227,6 +230,9 @@ class SAC(BaseRLModel): # Optimize the critic self.critic.optimizer.zero_grad() critic_loss.backward() + # Clip grad norm + if self.max_grad_norm is not None: + th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.critic.optimizer.step() # Compute actor loss