Updated type hint and extended docstring in make_vec_env and make_atari_env (#1085)

* Updated type hint and extended docstring in make_vec_env

The function itself was already working with callables, but it wasn't considerent in the type hint of the function's signature.

Extended the description of the wrapper_class parameter with a link to a Github issue containing more details on the matter.

* Updated type hint in make_atari_env

The function itself was already working with callables, but it wasn't considerent in the type hint of the function's signature.

* Updated docstring in make_atari_env

When modifying the type hint of the parameter 'env_id' (in this commit: fda6872f73c11075901ba88f2520f6316f818d1d), I forgot to update its description in the docstrig.
Doing it now.

* Removed redundant type in env_id's type hint in make_vec_env and make_atari_env

Callable[..., gym.Env] already includes Type[gym.Env], as pointed out here: https://github.com/DLR-RM/stable-baselines3/pull/1085#issuecomment-1269685218

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>

docs/misc/changelog.rst

@@ -2,6 +2,33 @@
 
 Changelog
 ==========
+Release 1.6.2 WIP
+---------------------------
+
+**Bug fix release**
+
+Breaking Changes:
+^^^^^^^^^^^^^^^^^
+
+New Features:
+^^^^^^^^^^^^^
+
+SB3-Contrib
+^^^^^^^^^^^
+
+Bug Fixes:
+^^^^^^^^^^
+
+Deprecations:
+^^^^^^^^^^^^^
+
+Others:
+^^^^^^^
+- Fixed type hint of the ``env_id`` parameter in ``make_vec_env`` and ``make_atari_env`` (@AlexPasqua)
+
+Documentation:
+^^^^^^^^^^^^^^
+- Extended docstring of the ``wrapper_class`` parameter in ``make_vec_env`` (@AlexPasqua)
 
 Release 1.6.1 (2022-09-29)
 ---------------------------

stable_baselines3/common/env_util.py

@@ -36,7 +36,7 @@ def is_wrapped(env: Type[gym.Env], wrapper_class: Type[gym.Wrapper]) -> bool:
 
 
 def make_vec_env(
-    env_id: Union[str, Type[gym.Env]],
+    env_id: Union[str, Callable[..., gym.Env]],
     n_envs: int = 1,
     seed: Optional[int] = None,
     start_index: int = 0,
@@ -53,7 +53,7 @@ def make_vec_env(
     By default it uses a ``DummyVecEnv`` which is usually faster
     than a ``SubprocVecEnv``.
 
-    :param env_id: the environment ID or the environment class
+    :param env_id: either the env ID, the env class or a callable returning an env
     :param n_envs: the number of environments you wish to have in parallel
     :param seed: the initial seed for the random number generator
     :param start_index: start rank index
@@ -62,6 +62,9 @@ def make_vec_env(
         in a Monitor wrapper to provide additional information about training.
     :param wrapper_class: Additional wrapper to use on the environment.
         This can also be a function with single argument that wraps the environment in many things.
+        Note: the wrapper specified by this parameter will be applied after the ``Monitor`` wrapper.
+        if some cases (e.g. with TimeLimit wrapper) this can lead to undesired behavior.
+        See here for more details: https://github.com/DLR-RM/stable-baselines3/issues/894
     :param env_kwargs: Optional keyword argument to pass to the env constructor
     :param vec_env_cls: A custom ``VecEnv`` class constructor. Default: None.
     :param vec_env_kwargs: Keyword arguments to pass to the ``VecEnv`` class constructor.
@@ -106,7 +109,7 @@ def make_vec_env(
 
 
 def make_atari_env(
-    env_id: Union[str, Type[gym.Env]],
+    env_id: Union[str, Callable[..., gym.Env]],
     n_envs: int = 1,
     seed: Optional[int] = None,
     start_index: int = 0,
@@ -121,7 +124,7 @@ def make_atari_env(
     Create a wrapped, monitored VecEnv for Atari.
     It is a wrapper around ``make_vec_env`` that includes common preprocessing for Atari games.
 
-    :param env_id: the environment ID or the environment class
+    :param env_id: either the env ID, the env class or a callable returning an env
     :param n_envs: the number of environments you wish to have in parallel
     :param seed: the initial seed for the random number generator
     :param start_index: start rank index