Fix SAC target network parameters (#871)

* Fix SAC target network to include any CNN parameters for target updates * Update changelog
Stable-Baselines-Team · May 28, 2020 · 257ab9c · 257ab9c
1 parent c4c31cb
commit 257ab9c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -28,6 +28,7 @@ Bug Fixes:
 - Partially fix tensorboard indexing for PPO2 (@enderdead)
 - Fixed potential bug in ``DummyVecEnv`` where ``copy()`` was used instead of ``deepcopy()``
 - Fixed a bug in ``GAIL`` where the dataloader was not available after saving, causing an error when using ``CheckpointCallback``
+- Fixed a bug in ``SAC`` where any convolutional layers were not included in the target network parameters.
 
 Deprecations:
 ^^^^^^^^^^^^^

diff --git a/stable_baselines/sac/sac.py b/stable_baselines/sac/sac.py
@@ -258,8 +258,8 @@ def setup_model(self):
                     value_optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_ph)
                     values_params = tf_util.get_trainable_vars('model/values_fn')
 
-                    source_params = tf_util.get_trainable_vars("model/values_fn/vf")
-                    target_params = tf_util.get_trainable_vars("target/values_fn/vf")
+                    source_params = tf_util.get_trainable_vars("model/values_fn")
+                    target_params = tf_util.get_trainable_vars("target/values_fn")
 
                     # Polyak averaging for target variables
                     self.target_update_op = [
@@ -304,7 +304,7 @@ def setup_model(self):
 
                 # Retrieve parameters that must be saved
                 self.params = tf_util.get_trainable_vars("model")
-                self.target_params = tf_util.get_trainable_vars("target/values_fn/vf")
+                self.target_params = tf_util.get_trainable_vars("target/values_fn")
 
                 # Initialize Variables and target network
                 with self.sess.as_default():