From 503425932f5dc59880f854c4f0db3255a3aa8c1e Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Sun, 18 Jul 2021 20:51:41 +0200 Subject: [PATCH] Documentation fixes (#514) * Update multiprocessing example * Add VecEnvWrapper example * Update docs/guide/vec_envs.rst Co-authored-by: Anssi Co-authored-by: Anssi --- docs/guide/examples.rst | 7 ++++--- docs/guide/vec_envs.rst | 45 +++++++++++++++++++++++++++++++++++++++++ docs/misc/changelog.rst | 2 ++ 3 files changed, 51 insertions(+), 3 deletions(-) diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst index 7ed2748..b02b8c3 100644 --- a/docs/guide/examples.rst +++ b/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) diff --git a/docs/guide/vec_envs.rst b/docs/guide/vec_envs.rst index 7958fe0..b074dad 100644 --- a/docs/guide/vec_envs.rst +++ b/docs/guide/vec_envs.rst @@ -44,6 +44,51 @@ SubprocVecEnv ✔️ ✔️ ✔️ ✔️ ✔️ For more information, see Python's `multiprocessing guidelines `_. +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 ------ diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index e3a5daf..1e89257 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -26,6 +26,8 @@ Others: Documentation: ^^^^^^^^^^^^^^ +- Updated multiprocessing example +- Added example of ``VecEnvWrapper`` Release 1.1.0 (2021-07-01)