Fix ent coef loading for SAC (#429)

* Fix ent coef loading for SAC

* Better fix and add comment

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.1.0a6 (WIP)
+Release 1.1.0a7 (WIP)
 ---------------------------
 
 **Dict observation support, timeout handling and refactored HER**

docs/modules/dqn.rst

@@ -134,3 +134,6 @@ DQN Policies
 
 .. autoclass:: CnnPolicy
   :members:
+
+.. autoclass:: MultiInputPolicy
+  :members:

stable_baselines3/common/base_class.py

@@ -665,6 +665,13 @@ class BaseAlgorithm(ABC):
         # put other pytorch variables back in place
         if pytorch_variables is not None:
             for name in pytorch_variables:
+                # Skip if PyTorch variable was not defined (to ensure backward compatibility).
+                # This happens when using SAC/TQC.
+                # SAC has an entropy coefficient which can be fixed or optimized.
+                # If it is optimized, an additional PyTorch variable `log_ent_coef` is defined,
+                # otherwise it is initialized to `None`.
+                if pytorch_variables[name] is None:
+                    continue
                 # Set the data attribute directly to avoid issue when using optimizers
                 # See https://github.com/DLR-RM/stable-baselines3/issues/391
                 recursive_setattr(model, name + ".data", pytorch_variables[name].data)

stable_baselines3/sac/sac.py

@@ -304,9 +304,9 @@ class SAC(OffPolicyAlgorithm):
 
     def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
         state_dicts = ["policy", "actor.optimizer", "critic.optimizer"]
-        saved_pytorch_variables = ["log_ent_coef"]
         if self.ent_coef_optimizer is not None:
+            saved_pytorch_variables = ["log_ent_coef"]
             state_dicts.append("ent_coef_optimizer")
         else:
-            saved_pytorch_variables.append("ent_coef_tensor")
+            saved_pytorch_variables = ["ent_coef_tensor"]
         return state_dicts, saved_pytorch_variables

stable_baselines3/version.txt

@@ -1 +1 @@
-1.1.0a6
+1.1.0a7

tests/test_save_load.py

@@ -239,16 +239,35 @@ def test_save_load_pytorch_var(tmp_path):
     save_path = str(tmp_path / "sac_pendulum")
     model.save(save_path)
     env = model.get_env()
-    ent_coef_before = model.log_ent_coef
+    log_ent_coef_before = model.log_ent_coef
 
     del model
 
     model = SAC.load(save_path, env=env)
-    assert th.allclose(ent_coef_before, model.log_ent_coef)
+    assert th.allclose(log_ent_coef_before, model.log_ent_coef)
     model.learn(200)
-    ent_coef_after = model.log_ent_coef
+    log_ent_coef_after = model.log_ent_coef
     # Check that the entropy coefficient is still optimized
-    assert not th.allclose(ent_coef_before, ent_coef_after)
+    assert not th.allclose(log_ent_coef_before, log_ent_coef_after)
+
+    # With a fixed entropy coef
+    model = SAC("MlpPolicy", "Pendulum-v0", seed=3, ent_coef=0.01, policy_kwargs=dict(net_arch=[64], n_critics=1))
+    model.learn(200)
+    save_path = str(tmp_path / "sac_pendulum")
+    model.save(save_path)
+    env = model.get_env()
+    assert model.log_ent_coef is None
+    ent_coef_before = model.ent_coef_tensor
+
+    del model
+
+    model = SAC.load(save_path, env=env)
+    assert th.allclose(ent_coef_before, model.ent_coef_tensor)
+    model.learn(200)
+    ent_coef_after = model.ent_coef_tensor
+    assert model.log_ent_coef is None
+    # Check that the entropy coefficient is still the same
+    assert th.allclose(ent_coef_before, ent_coef_after)
 
 
 @pytest.mark.parametrize("model_class", [A2C, TD3])