Documentation fixes (#514)

* Update multiprocessing example

* Add VecEnvWrapper example

* Update docs/guide/vec_envs.rst

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

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

docs/guide/examples.rst

@@ -109,7 +109,7 @@ Multiprocessing: Unleashing the Power of Vectorized Environments
   import numpy as np
 
   from stable_baselines3 import PPO
-  from stable_baselines3.common.vec_env import SubprocVecEnv
+  from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
   from stable_baselines3.common.env_util import make_vec_env
   from stable_baselines3.common.utils import set_random_seed
 
@@ -136,8 +136,9 @@ Multiprocessing: Unleashing the Power of Vectorized Environments
       env = SubprocVecEnv([make_env(env_id, i) for i in range(num_cpu)])
 
       # Stable Baselines provides you with make_vec_env() helper
-      # which does exactly the previous steps for you:
-      # env = make_vec_env(env_id, n_envs=num_cpu, seed=0)
+      # which does exactly the previous steps for you.
+      # You can choose between `DummyVecEnv` (usually faster) and `SubprocVecEnv`
+      # env = make_vec_env(env_id, n_envs=num_cpu, seed=0, vec_env_cls=SubprocVecEnv)
 
       model = PPO('MlpPolicy', env, verbose=1)
       model.learn(total_timesteps=25000)

docs/guide/vec_envs.rst

@@ -44,6 +44,51 @@ SubprocVecEnv ✔️       ✔️           ✔️        ✔️         ✔️
     For more information, see Python's `multiprocessing guidelines <https://docs.python.org/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods>`_.
 
 
+Vectorized Environments Wrappers
+--------------------------------
+
+If you want to alter or augment a ``VecEnv`` without redefining it completely (e.g. stack multiple frames, monitor the ``VecEnv``, normalize the observation, ...), you can use ``VecEnvWrapper`` for that.
+They are the vectorized equivalents (i.e., they act on multiple environments at the same time) of ``gym.Wrapper``.
+
+You can find below an example for extracting one key from the observation:
+
+.. code-block:: python
+
+	import numpy as np
+
+	from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper
+
+
+	class VecExtractDictObs(VecEnvWrapper):
+	    """
+	    A vectorized wrapper for filtering a specific key from dictionary observations.
+	    Similar to Gym's FilterObservation wrapper:
+	        https://github.com/openai/gym/blob/master/gym/wrappers/filter_observation.py
+
+	    :param venv: The vectorized environment
+	    :param key: The key of the dictionary observation
+	    """
+
+	    def __init__(self, venv: VecEnv, key: str):
+	        self.key = key
+	        super().__init__(venv=venv, observation_space=venv.observation_space.spaces[self.key])
+
+	    def reset(self) -> np.ndarray:
+	        obs = self.venv.reset()
+	        return obs[self.key]
+
+	    def step_async(self, actions: np.ndarray) -> None:
+	        self.venv.step_async(actions)
+
+	    def step_wait(self) -> VecEnvStepReturn:
+	        obs, reward, done, info = self.venv.step_wait()
+	        return obs[self.key], reward, done, info
+
+	env = DummyVecEnv([lambda: gym.make("FetchReach-v1")])
+	# Wrap the VecEnv
+	env = VecExtractDictObs(env, key="observation")
+
+
 VecEnv
 ------
 

docs/misc/changelog.rst

@@ -26,6 +26,8 @@ Others:
 
 Documentation:
 ^^^^^^^^^^^^^^
+- Updated multiprocessing example
+- Added example of ``VecEnvWrapper``
 
 
 Release 1.1.0 (2021-07-01)