diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 965c9da..554965c 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,7 +4,7 @@ Changelog ========== -Pre-Release 0.10.0a1 (WIP) +Pre-Release 0.10.0a2 (WIP) ------------------------------ Breaking Changes: @@ -30,6 +30,8 @@ Bug Fixes: - Update the check for spaces unsupported by Stable Baselines 3 to include checks on the action space (@wmmc88) - Fixed feature extractor bug for target network where the same net was shared instead of being separate. This bug affects ``SAC``, ``DDPG`` and ``TD3`` when using ``CnnPolicy`` (or custom feature extractor) +- Fixed a bug when passing an environment when loading a saved model with a ``CnnPolicy``, the passed env was not wrapped properly + (the bug was introduced when implementing ``HER`` so it should not be present in previous versions) Deprecations: ^^^^^^^^^^^^^ diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index f270c60..b4195f7 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -592,7 +592,7 @@ class BaseAlgorithm(ABC): if env is not None: # Wrap first if needed - cls._wrap_env(env, data["verbose"]) + env = cls._wrap_env(env, data["verbose"]) # Check if given env is valid check_for_correct_spaces(env, data["observation_space"], data["action_space"]) else: diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 8dabd1f..bb625ca 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.10.0a1 +0.10.0a2 diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 3787ebc..f7c5521 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -227,6 +227,27 @@ def test_exclude_include_saved_params(tmp_path, model_class): os.remove(tmp_path / "test_save.zip") +@pytest.mark.parametrize("model_class", [A2C, TD3]) +def test_save_load_env_cnn(tmp_path, model_class): + """ + Test loading with an env that requires a ``CnnPolicy``. + This is to test wrapping and observation space check. + We test one on-policy and one off-policy + algorithm as the rest share the loading part. + """ + env = FakeImageEnv(screen_height=40, screen_width=40, n_channels=2, discrete=False) + kwargs = dict(policy_kwargs=dict(net_arch=[32])) + if model_class == TD3: + kwargs.update(dict(buffer_size=100, learning_starts=50)) + + model = model_class("CnnPolicy", env, **kwargs).learn(100) + model.save(tmp_path / "test_save") + # Test loading with env and continuing training + model = model_class.load(str(tmp_path / "test_save.zip"), env=env).learn(100) + # clear file from os + os.remove(tmp_path / "test_save.zip") + + @pytest.mark.parametrize("model_class", [SAC, TD3, DQN]) def test_save_load_replay_buffer(tmp_path, model_class): path = pathlib.Path(tmp_path / "logs/replay_buffer.pkl")