Added wrapper_kwargs argument to make_vec_env (#448)

* Added wrapper_kwargs to make_vec_env

* code black format

* Tmp fix for atari-py

* Update changelog

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

.github/workflows/ci.yml

@@ -32,6 +32,8 @@ jobs:
         pip install .[extra,tests,docs]
         # Use headless version
         pip install opencv-python-headless
+        # Tmp fix: ROM missing in the newest atari-py version
+        pip install atari-py==0.2.5
     - name: Build the doc
       run: |
         make doc

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.1.0a8 (WIP)
+Release 1.1.0a9 (WIP)
 ---------------------------
 
 **Dict observation support, timeout handling and refactored HER**
@@ -44,6 +44,7 @@ New Features:
 - Added support for image observation when using ``HER``
 - Added ``replay_buffer_class`` and ``replay_buffer_kwargs`` arguments to off-policy algorithms
 - Added ``kl_divergence`` helper for ``Distribution`` classes (@09tangriro)
+- Added ``wrapper_kwargs`` argument to ``make_vec_env`` (@amy12xx)
 
 Bug Fixes:
 ^^^^^^^^^^
@@ -689,4 +690,4 @@ And all the contributors:
 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio
 @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber @thisray
 @tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @JadenTravnik @decodyng @ardabbour @lorenz-h @mschweizer @lorepieri8 @vwxyzjn
-@ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro
+@ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro @amy12xx

stable_baselines3/common/env_util.py

@@ -46,6 +46,7 @@ def make_vec_env(
     vec_env_cls: Optional[Type[Union[DummyVecEnv, SubprocVecEnv]]] = None,
     vec_env_kwargs: Optional[Dict[str, Any]] = None,
     monitor_kwargs: Optional[Dict[str, Any]] = None,
+    wrapper_kwargs: Optional[Dict[str, Any]] = None,
 ) -> VecEnv:
     """
     Create a wrapped, monitored ``VecEnv``.
@@ -65,11 +66,13 @@ def make_vec_env(
     :param vec_env_cls: A custom ``VecEnv`` class constructor. Default: None.
     :param vec_env_kwargs: Keyword arguments to pass to the ``VecEnv`` class constructor.
     :param monitor_kwargs: Keyword arguments to pass to the ``Monitor`` class constructor.
+    :param wrapper_kwargs: Keyword arguments to pass to the ``Wrapper`` class constructor.
     :return: The wrapped environment
     """
     env_kwargs = {} if env_kwargs is None else env_kwargs
     vec_env_kwargs = {} if vec_env_kwargs is None else vec_env_kwargs
     monitor_kwargs = {} if monitor_kwargs is None else monitor_kwargs
+    wrapper_kwargs = {} if wrapper_kwargs is None else wrapper_kwargs
 
     def make_env(rank):
         def _init():
@@ -89,7 +92,7 @@ def make_vec_env(
             env = Monitor(env, filename=monitor_path, **monitor_kwargs)
             # Optionally, wrap the environment with the provided wrapper
             if wrapper_class is not None:
-                env = wrapper_class(env)
+                env = wrapper_class(env, **wrapper_kwargs)
             return env
 
         return _init

stable_baselines3/version.txt

@@ -1 +1 @@
-1.1.0a8
+1.1.0a9

tests/test_utils.py

@@ -7,7 +7,7 @@ import pytest
 import torch as th
 
 from stable_baselines3 import A2C, PPO
-from stable_baselines3.common.atari_wrappers import ClipRewardEnv
+from stable_baselines3.common.atari_wrappers import ClipRewardEnv, MaxAndSkipEnv
 from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_vec_env, unwrap_wrapper
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.monitor import Monitor
@@ -70,6 +70,11 @@ def test_vec_env_kwargs():
     assert env.get_attr("goal_velocity")[0] == 0.11
 
 
+def test_vec_env_wrapper_kwargs():
+    env = make_vec_env("MountainCarContinuous-v0", n_envs=1, seed=0, wrapper_class=MaxAndSkipEnv, wrapper_kwargs={"skip": 3})
+    assert env.get_attr("_skip")[0] == 3
+
+
 def test_vec_env_monitor_kwargs():
     env = make_vec_env("MountainCarContinuous-v0", n_envs=1, seed=0, monitor_kwargs={"allow_early_resets": False})
     assert env.get_attr("allow_early_resets")[0] is False