Sync with Stable-Baselines

.coveragerc

@@ -5,6 +5,8 @@ omit =
     setup.py
     # Require graphical interface
     stable_baselines3/common/results_plotter.py
+    # Require ffmpeg
+    stable_baselines3/common/vec_env/vec_video_recorder.py
 
 [report]
 exclude_lines =

Makefile

@@ -16,3 +16,26 @@ clean:
 	cd docs && make clean
 
 .PHONY: clean spelling doc
+
+# TODO: create Dockerfile
+# # Build docker images
+# # If you do export RELEASE=True, it will also push them
+# docker: docker-cpu docker-gpu
+#
+# docker-cpu:
+# 	./scripts/build_docker.sh
+#
+# docker-gpu:
+# 	USE_GPU=True ./scripts/build_docker.sh
+
+# PyPi package release
+release:
+	python setup.py sdist
+	python setup.py bdist_wheel
+	twine upload dist/*
+
+# Test PyPi package release
+test-release:
+	python setup.py sdist
+	python setup.py bdist_wheel
+	twine upload --repository-url https://test.pypi.org/legacy/ dist/*

README.md

@@ -56,12 +56,12 @@ make spelling
 To cite this repository in publications:
 
 ```
-@misc{torchy-baselines,
+@misc{stable-baselines3,
   author = {Raffin, Antonin and Hill, Ashley and Ernestus, Maximilian and Gleave, Adam and Kanervisto, Anssi and Dormann, Noah},
   title = {Stable Baselines3},
   year = {2019},
   publisher = {GitHub},
   journal = {GitHub repository},
-  howpublished = {\url{https://github.com/araffin/torchy-baselines}},
+  howpublished = {\url{https://github.com/DLR-RM/stable-baselines3}},
 }
 ```

docs/index.rst

@@ -3,8 +3,8 @@
    You can adapt this file completely to your liking, but it should at least
    contain the root `toctree` directive.
 
-Welcome to Stable Baselines3 docs! - Pytorch RL Baselines
+Welcome to Stable Baselines3 docs!
-========================================================
+==================================
 
 `Stable Baselines3 <https://github.com/hill-a/stable-baselines>`_ is the PyTorch version of `Stable Baselines <https://github.com/hill-a/stable-baselines>`_,
 a set of improved implementations of reinforcement learning algorithms.
@@ -42,18 +42,18 @@ RL Baselines zoo also offers a simple interface to train, evaluate agents and do
 
 
 Citing Stable Baselines3
------------------------
+------------------------
 To cite this project in publications:
 
 .. code-block:: bibtex
 
-    @misc{torchy-baselines,
+    @misc{stable-baselines3,
       author = {Raffin, Antonin and Hill, Ashley and Ernestus, Maximilian and Gleave, Adam and Kanervisto, Anssi and Dormann, Noah},
       title = {Stable Baselines3},
       year = {2019},
       publisher = {GitHub},
       journal = {GitHub repository},
-      howpublished = {\url{https://github.com/araffin/torchy-baselines}},
+      howpublished = {\url{https://github.com/DLR-RM/stable-baselines3}},
     }
 
 Indices and tables

docs/misc/changelog.rst

@@ -3,6 +3,32 @@
 Changelog
 ==========
 
+Pre-Release 0.6.0a (WIP)
+------------------------------
+
+
+Breaking Changes:
+^^^^^^^^^^^^^^^^^
+
+New Features:
+^^^^^^^^^^^^^
+- Added env checker (Sync with Stable Baselines)
+- Added ``VecCheckNan`` and ``VecVideoRecorder`` (Sync with Stable Baselines)
+
+Bug Fixes:
+^^^^^^^^^^
+
+Deprecations:
+^^^^^^^^^^^^^
+
+Others:
+^^^^^^^
+- Renamed to Stable-Baseline3
+
+Documentation:
+^^^^^^^^^^^^^^
+
+
 Pre-Release 0.5.0 (2020-05-05)
 ------------------------------
 
@@ -160,10 +186,17 @@ New Features:
 Maintainers
 -----------
 
-Stable-Baselines3 is currently maintained by `Antonin Raffin`_ (aka `@araffin`_).
+Stable-Baselines3 is currently maintained by `Antonin Raffin`_ (aka `@araffin`_), `Ashley Hill`_ (aka @hill-a),
+`Maximilian Ernestus`_ (aka @erniejunior), `Adam Gleave`_ (`@AdamGleave`_) and `Anssi Kanervisto`_ (aka `@Miffyli`_).
 
+.. _Ashley Hill: https://github.com/hill-a
 .. _Antonin Raffin: https://araffin.github.io/
+.. _Maximilian Ernestus: https://github.com/erniejunior
+.. _Adam Gleave: https://gleave.me/
 .. _@araffin: https://github.com/araffin
+.. _@AdamGleave: https://github.com/adamgleave
+.. _Anssi Kanervisto: https://github.com/Miffyli
+.. _@Miffyli: https://github.com/Miffyli
 
 
 

setup.py

@@ -45,7 +45,7 @@ setup(name='stable_baselines3',
       },
       description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',
       author='Antonin Raffin',
-      url='',
+      url='https://github.com/DLR-RM/stable-baselines3',
       author_email='antonin.raffin@dlr.de',
       keywords="reinforcement-learning-algorithms reinforcement-learning machine-learning "
                "gym openai stable baselines toolbox python data-science",

stable_baselines3/common/bit_flipping_env.py

@@ -0,0 +1,121 @@
+from collections import OrderedDict
+from typing import Optional, Union
+
+import numpy as np
+from gym import GoalEnv, spaces
+
+from stable_baselines3.common.type_aliases import GymStepReturn
+
+
+class BitFlippingEnv(GoalEnv):
+    """
+    Simple bit flipping env, useful to test HER.
+    The goal is to flip all the bits to get a vector of ones.
+    In the continuous variant, if the ith action component has a value > 0,
+    then the ith bit will be flipped.
+
+    :param n_bits: (int) Number of bits to flip
+    :param continuous: (bool) Whether to use the continuous actions version or not,
+        by default, it uses the discrete one
+    :param max_steps: (Optional[int]) Max number of steps, by default, equal to n_bits
+    :param discrete_obs_space: (bool) Whether to use the discrete observation
+        version or not, by default, it uses the MultiBinary one
+    """
+    def __init__(self, n_bits: int = 10,
+                 continuous: bool = False,
+                 max_steps: Optional[int] = None,
+                 discrete_obs_space: bool = False):
+        super(BitFlippingEnv, self).__init__()
+        # The achieved goal is determined by the current state
+        # here, it is a special where they are equal
+        if discrete_obs_space:
+            # In the discrete case, the agent act on the binary
+            # representation of the observation
+            self.observation_space = spaces.Dict({
+                'observation': spaces.Discrete(2 ** n_bits - 1),
+                'achieved_goal': spaces.Discrete(2 ** n_bits - 1),
+                'desired_goal': spaces.Discrete(2 ** n_bits - 1)
+            })
+        else:
+            self.observation_space = spaces.Dict({
+                'observation': spaces.MultiBinary(n_bits),
+                'achieved_goal': spaces.MultiBinary(n_bits),
+                'desired_goal': spaces.MultiBinary(n_bits)
+            })
+
+        self.obs_space = spaces.MultiBinary(n_bits)
+
+        if continuous:
+            self.action_space = spaces.Box(-1, 1, shape=(n_bits,), dtype=np.float32)
+        else:
+            self.action_space = spaces.Discrete(n_bits)
+        self.continuous = continuous
+        self.discrete_obs_space = discrete_obs_space
+        self.state = None
+        self.desired_goal = np.ones((n_bits,))
+        if max_steps is None:
+            max_steps = n_bits
+        self.max_steps = max_steps
+        self.current_step = 0
+        self.reset()
+
+    def convert_if_needed(self, state: np.ndarray) -> Union[int, np.ndarray]:
+        """
+        Convert to discrete space if needed.
+
+        :param state: (np.ndarray)
+        :return: (np.ndarray or int)
+        """
+        if self.discrete_obs_space:
+            # The internal state is the binary representation of the
+            # observed one
+            return int(sum([state[i] * 2**i for i in range(len(state))]))
+        return state
+
+    def _get_obs(self) -> OrderedDict:
+        """
+        Helper to create the observation.
+
+        :return: (OrderedDict<int or ndarray>)
+        """
+        return OrderedDict([
+            ('observation', self.convert_if_needed(self.state.copy())),
+            ('achieved_goal', self.convert_if_needed(self.state.copy())),
+            ('desired_goal', self.convert_if_needed(self.desired_goal.copy()))
+        ])
+
+    def reset(self) -> OrderedDict:
+        self.current_step = 0
+        self.state = self.obs_space.sample()
+        return self._get_obs()
+
+    def step(self, action: Union[np.ndarray, int]) -> GymStepReturn:
+        if self.continuous:
+            self.state[action > 0] = 1 - self.state[action > 0]
+        else:
+            self.state[action] = 1 - self.state[action]
+        obs = self._get_obs()
+        reward = self.compute_reward(obs['achieved_goal'], obs['desired_goal'], None)
+        done = reward == 0
+        self.current_step += 1
+        # Episode terminate when we reached the goal or the max number of steps
+        info = {'is_success': done}
+        done = done or self.current_step >= self.max_steps
+        return obs, reward, done, info
+
+    def compute_reward(self,
+                       achieved_goal: np.ndarray,
+                       desired_goal: np.ndarray,
+                       _info) -> float:
+        # Deceptive reward: it is positive only when the goal is achieved
+        if self.discrete_obs_space:
+            return 0.0 if achieved_goal == desired_goal else -1.0
+        return 0.0 if (achieved_goal == desired_goal).all() else -1.0
+
+    def render(self, mode: str = 'human') -> Optional[np.ndarray]:
+        if mode == 'rgb_array':
+            return self.state.copy()
+        print(self.state)
+
+    def close(self) -> None:
+        pass

stable_baselines3/common/env_checker.py

@@ -0,0 +1,222 @@
+import warnings
+from typing import Union
+
+import gym
+from gym import spaces
+import numpy as np
+
+from stable_baselines.common.vec_env import DummyVecEnv, VecCheckNan
+
+
+def _enforce_array_obs(observation_space: spaces.Space) -> bool:
+    """
+    Whether to check that the returned observation is a numpy array
+    it is not mandatory for `Dict` and `Tuple` spaces.
+    """
+    return not isinstance(observation_space, (spaces.Dict, spaces.Tuple))
+
+
+def _check_image_input(observation_space: spaces.Box) -> None:
+    """
+    Check that the input will be compatible with Stable-Baselines
+    when the observation is apparently an image.
+    """
+    if observation_space.dtype != np.uint8:
+        warnings.warn("It seems that your observation is an image but the `dtype` "
+                      "of your observation_space is not `np.uint8`. "
+                      "If your observation is not an image, we recommend you to flatten the observation "
+                      "to have only a 1D vector")
+
+    if np.any(observation_space.low != 0) or np.any(observation_space.high != 255):
+        warnings.warn("It seems that your observation space is an image but the "
+                      "upper and lower bounds are not in [0, 255]. "
+                      "Because the CNN policy normalize automatically the observation "
+                      "you may encounter issue if the values are not in that range."
+                      )
+
+    if observation_space.shape[0] < 36 or observation_space.shape[1] < 36:
+        warnings.warn("The minimal resolution for an image is 36x36 for the default CnnPolicy. "
+                      "You might need to use a custom `cnn_extractor` "
+                      "cf https://stable-baselines.readthedocs.io/en/master/guide/custom_policy.html")
+
+
+def _check_unsupported_obs_spaces(env: gym.Env, observation_space: spaces.Space) -> None:
+    """Emit warnings when the observation space used is not supported by Stable-Baselines."""
+
+    if isinstance(observation_space, spaces.Dict) and not isinstance(env, gym.GoalEnv):
+        warnings.warn("The observation space is a Dict but the environment is not a gym.GoalEnv "
+                      "(cf https://github.com/openai/gym/blob/master/gym/core.py), "
+                      "this is currently not supported by Stable Baselines "
+                      "(cf https://github.com/hill-a/stable-baselines/issues/133), "
+                      "you will need to use a custom policy. "
+                      )
+
+    if isinstance(observation_space, spaces.Tuple):
+        warnings.warn("The observation space is a Tuple,"
+                      "this is currently not supported by Stable Baselines "
+                      "(cf https://github.com/hill-a/stable-baselines/issues/133), "
+                      "you will need to flatten the observation and maybe use a custom policy. "
+                      )
+
+
+def _check_nan(env: gym.Env) -> None:
+    """Check for Inf and NaN using the VecWrapper."""
+    vec_env = VecCheckNan(DummyVecEnv([lambda: env]))
+    for _ in range(10):
+        action = [env.action_space.sample()]
+        _, _, _, _ = vec_env.step(action)
+
+
+def _check_obs(obs: Union[tuple, dict, np.ndarray, int],
+               observation_space: spaces.Space,
+               method_name: str) -> None:
+    """
+    Check that the observation returned by the environment
+    correspond to the declared one.
+    """
+    if not isinstance(observation_space, spaces.Tuple):
+        assert not isinstance(obs, tuple), ("The observation returned by the `{}()` "
+                                            "method should be a single value, not a tuple".format(method_name))
+
+    # The check for a GoalEnv is done by the base class
+    if isinstance(observation_space, spaces.Discrete):
+        assert isinstance(obs, int), "The observation returned by `{}()` method must be an int".format(method_name)
+    elif _enforce_array_obs(observation_space):
+        assert isinstance(obs, np.ndarray), ("The observation returned by `{}()` "
+                                             "method must be a numpy array".format(method_name))
+
+    assert observation_space.contains(obs), ("The observation returned by the `{}()` "
+                                             "method does not match the given observation space".format(method_name))
+
+
+def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action_space: spaces.Space) -> None:
+    """
+    Check the returned values by the env when calling `.reset()` or `.step()` methods.
+    """
+    # because env inherits from gym.Env, we assume that `reset()` and `step()` methods exists
+    obs = env.reset()
+
+    _check_obs(obs, observation_space, 'reset')
+
+    # Sample a random action
+    action = action_space.sample()
+    data = env.step(action)
+
+    assert len(data) == 4, "The `step()` method must return four values: obs, reward, done, info"
+
+    # Unpack
+    obs, reward, done, info = data
+
+    _check_obs(obs, observation_space, 'step')
+
+    # We also allow int because the reward will be cast to float
+    assert isinstance(reward, (float, int)), "The reward returned by `step()` must be a float"
+    assert isinstance(done, bool), "The `done` signal must be a boolean"
+    assert isinstance(info, dict), "The `info` returned by `step()` must be a python dictionary"
+
+    if isinstance(env, gym.GoalEnv):
+        # For a GoalEnv, the keys are checked at reset
+        assert reward == env.compute_reward(obs['achieved_goal'], obs['desired_goal'], info)
+
+
+def _check_spaces(env: gym.Env) -> None:
+    """
+    Check that the observation and action spaces are defined
+    and inherit from gym.spaces.Space.
+    """
+    # Helper to link to the code, because gym has no proper documentation
+    gym_spaces = " cf https://github.com/openai/gym/blob/master/gym/spaces/"
+
+    assert hasattr(env, 'observation_space'), "You must specify an observation space (cf gym.spaces)" + gym_spaces
+    assert hasattr(env, 'action_space'), "You must specify an action space (cf gym.spaces)" + gym_spaces
+
+    assert isinstance(env.observation_space,
+                      spaces.Space), "The observation space must inherit from gym.spaces" + gym_spaces
+    assert isinstance(env.action_space, spaces.Space), "The action space must inherit from gym.spaces" + gym_spaces
+
+
+def _check_render(env: gym.Env, warn: bool = True, headless: bool = False) -> None:
+    """
+    Check the declared render modes and the `render()`/`close()`
+    method of the environment.
+
+    :param env: (gym.Env) The environment to check
+    :param warn: (bool) Whether to output additional warnings
+    :param headless: (bool) Whether to disable render modes
+        that require a graphical interface. False by default.
+    """
+    render_modes = env.metadata.get('render.modes')
+    if render_modes is None:
+        if warn:
+            warnings.warn("No render modes was declared in the environment "
+                          " (env.metadata['render.modes'] is None or not defined), "
+                          "you may have trouble when calling `.render()`")
+
+    else:
+        # Don't check render mode that require a
+        # graphical interface (useful for CI)
+        if headless and 'human' in render_modes:
+            render_modes.remove('human')
+        # Check all declared render modes
+        for render_mode in render_modes:
+            env.render(mode=render_mode)
+        env.close()
+
+
+def check_env(env: gym.Env, warn: bool = True, skip_render_check: bool = True) -> None:
+    """
+    Check that an environment follows Gym API.
+    This is particularly useful when using a custom environment.
+    Please take a look at https://github.com/openai/gym/blob/master/gym/core.py
+    for more information about the API.
+
+    It also optionally check that the environment is compatible with Stable-Baselines.
+
+    :param env: (gym.Env) The Gym environment that will be checked
+    :param warn: (bool) Whether to output additional warnings
+        mainly related to the interaction with Stable Baselines
+    :param skip_render_check: (bool) Whether to skip the checks for the render method.
+        True by default (useful for the CI)
+    """
+    assert isinstance(env, gym.Env), ("You environment must inherit from gym.Env class "
+                                      " cf https://github.com/openai/gym/blob/master/gym/core.py")
+
+    # ============= Check the spaces (observation and action) ================
+    _check_spaces(env)
+
+    # Define aliases for convenience
+    observation_space = env.observation_space
+    action_space = env.action_space
+
+    # Warn the user if needed.
+    # A warning means that the environment may run but not work properly with Stable Baselines algorithms
+    if warn:
+        _check_unsupported_obs_spaces(env, observation_space)
+
+        # If image, check the low and high values, the type and the number of channels
+        # and the shape (minimal value)
+        if isinstance(observation_space, spaces.Box) and len(observation_space.shape) == 3:
+            _check_image_input(observation_space)
+
+        if isinstance(observation_space, spaces.Box) and len(observation_space.shape) not in [1, 3]:
+            warnings.warn("Your observation has an unconventional shape (neither an image, nor a 1D vector). "
+                          "We recommend you to flatten the observation "
+                          "to have only a 1D vector")
+
+        # Check for the action space, it may lead to hard-to-debug issues
+        if (isinstance(action_space, spaces.Box) and
+                (np.any(np.abs(action_space.low) != np.abs(action_space.high))
+                 or np.any(np.abs(action_space.low) > 1) or np.any(np.abs(action_space.high) > 1))):
+            warnings.warn("We recommend you to use a symmetric and normalized Box action space (range=[-1, 1]) "
+                          "cf https://stable-baselines.readthedocs.io/en/master/guide/rl_tips.html")
+
+    # ============ Check the returned values ===============
+    _check_returned_values(env, observation_space, action_space)
+
+    # ==== Check the render method and the declared render modes ====
+    if not skip_render_check:
+        _check_render(env, warn=warn)
+
+    # The check only works with numpy arrays
+    if _enforce_array_obs(observation_space):
+        _check_nan(env)

stable_baselines3/common/vec_env/__init__.py

@@ -10,6 +10,8 @@ from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv
 from stable_baselines3.common.vec_env.vec_frame_stack import VecFrameStack
 from stable_baselines3.common.vec_env.vec_normalize import VecNormalize
 from stable_baselines3.common.vec_env.vec_transpose import VecTransposeImage
+from stable_baselines3.common.vec_env.vec_video_recorder import VecVideoRecorder
+from stable_baselines3.common.vec_env.vec_check_nan import VecCheckNan
 
 # Avoid circular import
 if typing.TYPE_CHECKING:

stable_baselines3/common/vec_env/vec_check_nan.py

@@ -0,0 +1,86 @@
+import warnings
+
+import numpy as np
+
+from stable_baselines3.common.vec_env.base_vec_env import VecEnvWrapper
+
+
+class VecCheckNan(VecEnvWrapper):
+    """
+    NaN and inf checking wrapper for vectorized environment, will raise a warning by default,
+    allowing you to know from what the NaN of inf originated from.
+
+    :param venv: (VecEnv) the vectorized environment to wrap
+    :param raise_exception: (bool) Whether or not to raise a ValueError, instead of a UserWarning
+    :param warn_once: (bool) Whether or not to only warn once.
+    :param check_inf: (bool) Whether or not to check for +inf or -inf as well
+    """
+
+    def __init__(self, venv, raise_exception=False, warn_once=True, check_inf=True):
+        VecEnvWrapper.__init__(self, venv)
+        self.raise_exception = raise_exception
+        self.warn_once = warn_once
+        self.check_inf = check_inf
+        self._actions = None
+        self._observations = None
+        self._user_warned = False
+
+    def step_async(self, actions):
+        self._check_val(async_step=True, actions=actions)
+
+        self._actions = actions
+        self.venv.step_async(actions)
+
+    def step_wait(self):
+        observations, rewards, news, infos = self.venv.step_wait()
+
+        self._check_val(async_step=False, observations=observations, rewards=rewards, news=news)
+
+        self._observations = observations
+        return observations, rewards, news, infos
+
+    def reset(self):
+        observations = self.venv.reset()
+        self._actions = None
+
+        self._check_val(async_step=False, observations=observations)
+
+        self._observations = observations
+        return observations
+
+    def _check_val(self, *, async_step, **kwargs):
+        # if warn and warn once and have warned once: then stop checking
+        if not self.raise_exception and self.warn_once and self._user_warned:
+            return
+
+        found = []
+        for name, val in kwargs.items():
+            has_nan = np.any(np.isnan(val))
+            has_inf = self.check_inf and np.any(np.isinf(val))
+            if has_inf:
+                found.append((name, "inf"))
+            if has_nan:
+                found.append((name, "nan"))
+
+        if found:
+            self._user_warned = True
+            msg = ""
+            for i, (name, type_val) in enumerate(found):
+                msg += "found {} in {}".format(type_val, name)
+                if i != len(found) - 1:
+                    msg += ", "
+
+            msg += ".\r\nOriginated from the "
+
+            if not async_step:
+                if self._actions is None:
+                    msg += "environment observation (at reset)"
+                else:
+                    msg += "environment, Last given value was: \r\n\taction={}".format(self._actions)
+            else:
+                msg += "RL model, Last given value was: \r\n\tobservations={}".format(self._observations)
+
+            if self.raise_exception:
+                raise ValueError(msg)
+            else:
+                warnings.warn(msg, UserWarning)

stable_baselines3/common/vec_env/vec_video_recorder.py

@@ -0,0 +1,112 @@
+import os
+
+from gym.wrappers.monitoring import video_recorder
+
+from stable_baselines3.common import logger
+from stable_baselines3.common.vec_env.base_vec_env import VecEnvWrapper
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
+from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv
+from stable_baselines3.common.vec_env.vec_frame_stack import VecFrameStack
+from stable_baselines3.common.vec_env.vec_normalize import VecNormalize
+
+
+class VecVideoRecorder(VecEnvWrapper):
+    """
+    Wraps a VecEnv or VecEnvWrapper object to record rendered image as mp4 video.
+    It requires ffmpeg or avconv to be installed on the machine.
+
+    :param venv: (VecEnv or VecEnvWrapper)
+    :param video_folder: (str) Where to save videos
+    :param record_video_trigger: (func) Function that defines when to start recording.
+                                        The function takes the current number of step,
+                                        and returns whether we should start recording or not.
+    :param video_length: (int)  Length of recorded videos
+    :param name_prefix: (str) Prefix to the video name
+    """
+
+    def __init__(self, venv, video_folder, record_video_trigger,
+                 video_length=200, name_prefix='rl-video'):
+
+        VecEnvWrapper.__init__(self, venv)
+
+        self.env = venv
+        # Temp variable to retrieve metadata
+        temp_env = venv
+
+        # Unwrap to retrieve metadata dict
+        # that will be used by gym recorder
+        while isinstance(temp_env, VecNormalize) or isinstance(temp_env, VecFrameStack):
+            temp_env = temp_env.venv
+
+        if isinstance(temp_env, DummyVecEnv) or isinstance(temp_env, SubprocVecEnv):
+            metadata = temp_env.get_attr('metadata')[0]
+        else:
+            metadata = temp_env.metadata
+
+        self.env.metadata = metadata
+
+        self.record_video_trigger = record_video_trigger
+        self.video_recorder = None
+
+        self.video_folder = os.path.abspath(video_folder)
+        # Create output folder if needed
+        os.makedirs(self.video_folder, exist_ok=True)
+
+        self.name_prefix = name_prefix
+        self.step_id = 0
+        self.video_length = video_length
+
+        self.recording = False
+        self.recorded_frames = 0
+
+    def reset(self):
+        obs = self.venv.reset()
+        self.start_video_recorder()
+        return obs
+
+    def start_video_recorder(self):
+        self.close_video_recorder()
+
+        video_name = '{}-step-{}-to-step-{}'.format(self.name_prefix, self.step_id,
+                                                    self.step_id + self.video_length)
+        base_path = os.path.join(self.video_folder, video_name)
+        self.video_recorder = video_recorder.VideoRecorder(
+                env=self.env,
+                base_path=base_path,
+                metadata={'step_id': self.step_id}
+                )
+
+        self.video_recorder.capture_frame()
+        self.recorded_frames = 1
+        self.recording = True
+
+    def _video_enabled(self):
+        return self.record_video_trigger(self.step_id)
+
+    def step_wait(self):
+        obs, rews, dones, infos = self.venv.step_wait()
+
+        self.step_id += 1
+        if self.recording:
+            self.video_recorder.capture_frame()
+            self.recorded_frames += 1
+            if self.recorded_frames > self.video_length:
+                logger.info("Saving video to ", self.video_recorder.path)
+                self.close_video_recorder()
+        elif self._video_enabled():
+            self.start_video_recorder()
+
+        return obs, rews, dones, infos
+
+    def close_video_recorder(self):
+        if self.recording:
+            self.video_recorder.close()
+        self.recording = False
+        self.recorded_frames = 1
+
+    def close(self):
+        VecEnvWrapper.close(self)
+        self.close_video_recorder()
+
+    def __del__(self):
+        self.close()

stable_baselines3/version.txt

@@ -1 +1 @@
-0.5.0
+0.6.0a0

tests/test_envs.py

@@ -0,0 +1,149 @@
+import pytest
+import gym
+from gym import spaces
+import numpy as np
+
+from stable_baselines3.common.env_checker import check_env
+from stable_baselines3.common.bit_flipping_env import BitFlippingEnv
+from stable_baselines3.common.identity_env import (IdentityEnv, IdentityEnvBox,
+                                                  IdentityEnvMultiBinary, IdentityEnvMultiDiscrete,)
+
+ENV_CLASSES = [BitFlippingEnv, IdentityEnv, IdentityEnvBox, IdentityEnvMultiBinary,
+               IdentityEnvMultiDiscrete]
+
+
+@pytest.mark.parametrize("env_id", ['CartPole-v0', 'Pendulum-v0'])
+def test_env(env_id):
+    """
+    Check that environmnent integrated in Gym pass the test.
+
+    :param env_id: (str)
+    """
+    env = gym.make(env_id)
+    with pytest.warns(None) as record:
+        check_env(env)
+
+    # Pendulum-v0 will produce a warning because the action space is
+    # in [-2, 2] and not [-1, 1]
+    if env_id == 'Pendulum-v0':
+        assert len(record) == 1
+    else:
+        # The other environments must pass without warning
+        assert len(record) == 0
+
+
+@pytest.mark.parametrize("env_class", ENV_CLASSES)
+def test_custom_envs(env_class):
+    env = env_class()
+    check_env(env)
+
+
+def test_high_dimension_action_space():
+    """
+    Test for continuous action space
+    with more than one action.
+    """
+    env = gym.make('Pendulum-v0')
+    # Patch the action space
+    env.action_space = spaces.Box(low=-1, high=1, shape=(20,), dtype=np.float32)
+    # Patch to avoid error
+    def patched_step(_action):
+        return env.observation_space.sample(), 0.0, False, {}
+    env.step = patched_step
+    check_env(env)
+
+
+@pytest.mark.parametrize("new_obs_space", [
+    # Small image
+    spaces.Box(low=0, high=255, shape=(32, 32, 3), dtype=np.uint8),
+    # Range not in [0, 255]
+    spaces.Box(low=0, high=1, shape=(64, 64, 3), dtype=np.uint8),
+    # Wrong dtype
+    spaces.Box(low=0, high=255, shape=(64, 64, 3), dtype=np.float32),
+    # Not an image, it should be a 1D vector
+    spaces.Box(low=-1, high=1, shape=(64, 3), dtype=np.float32),
+    # Tuple space is not supported by SB
+    spaces.Tuple([spaces.Discrete(5), spaces.Discrete(10)]),
+    # Dict space is not supported by SB when env is not a GoalEnv
+    spaces.Dict({"position": spaces.Discrete(5)}),
+])
+def test_non_default_spaces(new_obs_space):
+    env = gym.make('BreakoutNoFrameskip-v4')
+    env.observation_space = new_obs_space
+    # Patch methods to avoid errors
+    env.reset = new_obs_space.sample
+
+    def patched_step(_action):
+        return new_obs_space.sample(), 0.0, False, {}
+
+    env.step = patched_step
+    with pytest.warns(UserWarning):
+        check_env(env)
+
+
+def check_reset_assert_error(env, new_reset_return):
+    """
+    Helper to check that the error is caught.
+    :param env: (gym.Env)
+    :param new_reset_return: (Any)
+    """
+
+    def wrong_reset():
+        return new_reset_return
+
+    # Patch the reset method with a wrong one
+    env.reset = wrong_reset
+    with pytest.raises(AssertionError):
+        check_env(env)
+
+
+def test_common_failures_reset():
+    """
+    Test that common failure cases of the `reset_method` are caught
+    """
+    env = IdentityEnvBox()
+    # Return an observation that does not match the observation_space
+    check_reset_assert_error(env, np.ones((3,)))
+    # The observation is not a numpy array
+    check_reset_assert_error(env, 1)
+
+    # Return not only the observation
+    check_reset_assert_error(env, (env.observation_space.sample(), False))
+
+
+def check_step_assert_error(env, new_step_return=()):
+    """
+    Helper to check that the error is caught.
+    :param env: (gym.Env)
+    :param new_step_return: (tuple)
+    """
+
+    def wrong_step(_action):
+        return new_step_return
+
+    # Patch the step method with a wrong one
+    env.step = wrong_step
+    with pytest.raises(AssertionError):
+        check_env(env)
+
+
+def test_common_failures_step():
+    """
+    Test that common failure cases of the `step` method are caught
+    """
+    env = IdentityEnvBox()
+
+    # Wrong shape for the observation
+    check_step_assert_error(env, (np.ones((4,)), 1.0, False, {}))
+    # Obs is not a numpy array
+    check_step_assert_error(env, (1, 1.0, False, {}))
+
+    # Return a wrong reward
+    check_step_assert_error(env, (env.observation_space.sample(), np.ones(1), False, {}))
+
+    # Info dict is not returned
+    check_step_assert_error(env, (env.observation_space.sample(), 0.0, False))
+
+    # Done is not a boolean
+    check_step_assert_error(env, (env.observation_space.sample(), 0.0, 3.0, {}))
+    check_step_assert_error(env, (env.observation_space.sample(), 0.0, 1, {}))

tests/test_vec_check_nan.py

@@ -0,0 +1,72 @@
+import gym
+from gym import spaces
+import numpy as np
+
+from stable_baselines.common.vec_env import DummyVecEnv, VecCheckNan
+
+
+class NanAndInfEnv(gym.Env):
+    """Custom Environment that raised NaNs and Infs"""
+    metadata = {'render.modes': ['human']}
+
+    def __init__(self):
+        super(NanAndInfEnv, self).__init__()
+        self.action_space = spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float64)
+        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float64)
+
+    @staticmethod
+    def step(action):
+        if np.all(np.array(action) > 0):
+            obs = float('NaN')
+        elif np.all(np.array(action) < 0):
+            obs = float('inf')
+        else:
+            obs = 0
+        return [obs], 0.0, False, {}
+
+    @staticmethod
+    def reset():
+        return [0.0]
+
+    def render(self, mode='human', close=False):
+        pass
+
+
+def test_check_nan():
+    """Test VecCheckNan Object"""
+
+    env = DummyVecEnv([NanAndInfEnv])
+    env = VecCheckNan(env, raise_exception=True)
+
+    env.step([[0]])
+
+    try:
+        env.step([[float('NaN')]])
+    except ValueError:
+        pass
+    else:
+        assert False
+
+    try:
+        env.step([[float('inf')]])
+    except ValueError:
+        pass
+    else:
+        assert False
+
+    try:
+        env.step([[-1]])
+    except ValueError:
+        pass
+    else:
+        assert False
+
+    try:
+        env.step([[1]])
+    except ValueError:
+        pass
+    else:
+        assert False
+
+
+    env.step(np.array([[0, 1], [0, 1]]))