Correct time to update checking in SAC/TD3 (#755)

* Update sac.py Corrected minor bug in sac where train_freq was compared to local variable `step` instead of `agent.num_timesteps`. * Update changelog.rst * Update changelog.rst * Update td3.py * Update changelog.rst * Update changelog.rst Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
Stable-Baselines-Team · Mar 21, 2020 · ae4f6c5 · ae4f6c5
1 parent 49b1ba6
commit ae4f6c5
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -19,6 +19,7 @@ Bug Fixes:
 ^^^^^^^^^^
 - Fixed DDPG sampling empty replay buffer when combined with HER  (@tirafesi)
 - Fixed a bug in ``HindsightExperienceReplayWrapper``, where the openai-gym signature for ``compute_reward`` was not matched correctly (@johannes-dornheim)
+- Fixed SAC/TD3 checking time to update on learn steps instead of total steps (@solliet)
 
 Deprecations:
 ^^^^^^^^^^^^^

diff --git a/stable_baselines/sac/sac.py b/stable_baselines/sac/sac.py
@@ -445,7 +445,7 @@ def learn(self, total_timesteps, callback=None,
                     tf_util.total_episode_reward_logger(self.episode_reward, ep_reward,
                                                         ep_done, writer, self.num_timesteps)
 
-                if step % self.train_freq == 0:
+                if self.num_timesteps % self.train_freq == 0:
                     callback.on_rollout_end()
 
                     mb_infos_vals = []

diff --git a/stable_baselines/td3/td3.py b/stable_baselines/td3/td3.py
@@ -364,7 +364,7 @@ def learn(self, total_timesteps, callback=None,
                     tf_util.total_episode_reward_logger(self.episode_reward, ep_reward,
                                                         ep_done, writer, self.num_timesteps)
 
-                if step % self.train_freq == 0:
+                if self.num_timesteps % self.train_freq == 0:
                     callback.on_rollout_end()
 
                     mb_infos_vals = []