diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index e4cab0f..945e513 100644 --- a/docs/misc/changelog.rst +++ b/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** diff --git a/docs/modules/dqn.rst b/docs/modules/dqn.rst index 2ada427..426de6b 100644 --- a/docs/modules/dqn.rst +++ b/docs/modules/dqn.rst @@ -134,3 +134,6 @@ DQN Policies .. autoclass:: CnnPolicy :members: + +.. autoclass:: MultiInputPolicy + :members: diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 7a97980..a0af69d 100644 --- a/stable_baselines3/common/base_class.py +++ b/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) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index a220955..bcbd165 100644 --- a/stable_baselines3/sac/sac.py +++ b/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 diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 1406d2f..5e7d254 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.1.0a6 +1.1.0a7 diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 5e4e9e7..7b1fef5 100644 --- a/tests/test_save_load.py +++ b/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])