diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index a47bff9..014630e 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -24,6 +24,7 @@ New Features: - Added ``env_is_wrapped()`` method for ``VecEnv`` to check if its environments are wrapped with given Gym wrappers. - Added ``monitor_kwargs`` parameter to ``make_vec_env`` and ``make_atari_env`` +- Wrap the environments automatically with a ``Monitor`` wrapper when possible. Bug Fixes: ^^^^^^^^^^ diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 908b999..32c1ce5 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -13,6 +13,7 @@ import torch as th from stable_baselines3.common import logger, utils from stable_baselines3.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback +from stable_baselines3.common.env_util import is_wrapped from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.policies import BasePolicy, get_policy_from_name @@ -37,11 +38,10 @@ from stable_baselines3.common.vec_env import ( from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper -def maybe_make_env(env: Union[GymEnv, str, None], monitor_wrapper: bool, verbose: int) -> Optional[GymEnv]: +def maybe_make_env(env: Union[GymEnv, str, None], verbose: int) -> Optional[GymEnv]: """If env is a string, make the environment; otherwise, return env. :param env: The environment to learn from. - :param monitor_wrapper: Whether to wrap env in a Monitor when creating env. :param verbose: logging verbosity :return A Gym (vector) environment. """ @@ -49,9 +49,6 @@ def maybe_make_env(env: Union[GymEnv, str, None], monitor_wrapper: bool, verbose if verbose >= 1: print(f"Creating environment from the given name '{env}'") env = gym.make(env) - if monitor_wrapper: - env = Monitor(env, filename=None) - return env @@ -151,10 +148,10 @@ class BaseAlgorithm(ABC): if env is not None: if isinstance(env, str): if create_eval_env: - self.eval_env = maybe_make_env(env, monitor_wrapper, self.verbose) + self.eval_env = maybe_make_env(env, self.verbose) - env = maybe_make_env(env, monitor_wrapper, self.verbose) - env = self._wrap_env(env, self.verbose) + env = maybe_make_env(env, self.verbose) + env = self._wrap_env(env, self.verbose, monitor_wrapper) self.observation_space = env.observation_space self.action_space = env.action_space @@ -170,8 +167,22 @@ class BaseAlgorithm(ABC): raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.") @staticmethod - def _wrap_env(env: GymEnv, verbose: int = 0) -> VecEnv: + def _wrap_env(env: GymEnv, verbose: int = 0, monitor_wrapper: bool = True) -> VecEnv: + """ " + Wrap environment with the appropriate wrappers if needed. + For instance, to have a vectorized environment + or to re-order the image channels. + + :param env: + :param verbose: + :param monitor_wrapper: Whether to wrap the env in a ``Monitor`` when possible. + :return: The wrapped environment. + """ if not isinstance(env, VecEnv): + if not is_wrapped(env, Monitor) and monitor_wrapper: + if verbose >= 1: + print("Wrapping the env with a `Monitor` wrapper") + env = Monitor(env) if verbose >= 1: print("Wrapping the env in a DummyVecEnv.") env = DummyVecEnv([lambda: env]) diff --git a/tests/test_utils.py b/tests/test_utils.py index d595fb6..3684900 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -84,6 +84,19 @@ def test_vec_env_monitor_kwargs(): assert env.get_attr("allow_early_resets")[0] is True +def test_env_auto_monitor_wrap(): + env = gym.make("Pendulum-v0") + model = A2C("MlpPolicy", env) + assert model.env.env_is_wrapped(Monitor)[0] is True + + env = Monitor(env) + model = A2C("MlpPolicy", env) + assert model.env.env_is_wrapped(Monitor)[0] is True + + model = A2C("MlpPolicy", "Pendulum-v0") + assert model.env.env_is_wrapped(Monitor)[0] is True + + def test_custom_vec_env(tmp_path): """ Stand alone test for a special case (passing a custom VecEnv class) to avoid doubling the number of tests.