Automatically wrap with a Monitor when possible (#237)

* Automatically wrap with a Monitor when possible

* Update stable_baselines3/common/base_class.py

Co-authored-by: Anssi <kaneran21@hotmail.com>

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

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 = maybe_make_env(env, self.verbose)
-            env = self._wrap_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])

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.