Fix env loading (#203)

* Fix loading bug

* [ci skip] Fix docstring

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:
 ^^^^^^^^^^^^^

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:

stable_baselines3/version.txt

@@ -1 +1 @@
-0.10.0a1
+0.10.0a2

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")