saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-23 19:32:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Try to clip grad norm

Browse source
This commit is contained in:
Antonin Raffin 2019-12-17 11:14:44 +01:00
parent d63cef7693
commit 919dfee452
1 changed files with 6 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
torchy_baselines/sac/sac.py
View file

@ -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