From 440735cbd01512fe9d11fb793ae54d24b1230436 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Sat, 17 Sep 2022 11:10:03 +0200 Subject: [PATCH] Fix loading a model with different number of environments (#1058) * Fix loading with new `n_envs` * Update tests * Update changelog * Fix the fix * Remove `self._setup_model()` from `set_env()` * Raise `AssertionError` when setting env with a different `n_envs` * Update unitests Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 1 + stable_baselines3/common/base_class.py | 8 ++++++++ tests/test_save_load.py | 10 ++++++++++ 3 files changed, 19 insertions(+) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 0d745e5..085cfed 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -29,6 +29,7 @@ Bug Fixes: - Fixed missing verbose parameter passing in the ``EvalCallback`` constructor (@burakdmb) - Fixed the issue that when updating the target network in DQN, SAC, TD3, the ``running_mean`` and ``running_var`` properties of batch norm layers are not updated (@honglu2875) - Fixed incorrect type annotation of the replay_buffer_class argument in ``common.OffPolicyAlgorithm`` initializer, where an instance instead of a class was required (@Rocamonde) +- Fixed loading saved model with different number of envrionments - Removed ``forward()`` abstract method declaration from ``common.policies.BaseModel`` (already defined in ``torch.nn.Module``) to fix type errors in subclasses (@Rocamonde) Deprecations: diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 72428e0..dc04874 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -516,6 +516,11 @@ class BaseAlgorithm(ABC): # if it is not a VecEnv, make it a VecEnv # and do other transformations (dict obs, image transpose) if needed env = self._wrap_env(env, self.verbose) + assert env.num_envs == self.n_envs, ( + "The number of environments to be set is different from the number of environments in the model: " + f"({env.num_envs} != {self.n_envs}), whereas `set_env` requires them to be the same. To load a model with " + f"a different number of environments, you must use `{self.__class__.__name__}.load(path, env)` instead" + ) # Check that the observation spaces match check_for_correct_spaces(env, self.observation_space, self.action_space) # Update VecNormalize object @@ -730,6 +735,9 @@ class BaseAlgorithm(ABC): # See issue https://github.com/DLR-RM/stable-baselines3/issues/597 if force_reset and data is not None: data["_last_obs"] = None + # `n_envs` must be updated. See issue https://github.com/DLR-RM/stable-baselines3/issues/1018 + if data is not None: + data["n_envs"] = env.num_envs else: # Use stored env, if one exists. If not, continue as is (can be used for predict) if "env" in data: diff --git a/tests/test_save_load.py b/tests/test_save_load.py index d7a74c5..988d432 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -174,6 +174,7 @@ def test_set_env(tmp_path, model_class): env = DummyVecEnv([lambda: select_env(model_class)]) env2 = DummyVecEnv([lambda: select_env(model_class)]) env3 = select_env(model_class) + env4 = DummyVecEnv([lambda: select_env(model_class) for _ in range(2)]) kwargs = {} if model_class in {DQN, DDPG, SAC, TD3}: @@ -199,6 +200,10 @@ def test_set_env(tmp_path, model_class): # learn again model.learn(total_timesteps=64) + # num_env must be the same + with pytest.raises(AssertionError): + model.set_env(env4) + # Keep the same env, disable reset model.set_env(model.get_env(), force_reset=False) assert model._last_obs is not None @@ -223,6 +228,11 @@ def test_set_env(tmp_path, model_class): model.learn(total_timesteps=64, reset_num_timesteps=False) assert model.num_timesteps == 3 * 64 + del model + # Load the model with a different number of environments + model = model_class.load(tmp_path / "test_save.zip", env=env4) + model.learn(total_timesteps=64) + # Clear saved file os.remove(tmp_path / "test_save.zip")