From 2c34a4d6946d0f64a9d12026cd359b1e52a316e0 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 5 May 2020 16:28:38 +0200 Subject: [PATCH] Sync with Stable-Baselines --- .coveragerc | 2 + Makefile | 23 ++ README.md | 4 +- docs/index.rst | 10 +- docs/misc/changelog.rst | 35 ++- setup.py | 2 +- stable_baselines3/common/bit_flipping_env.py | 121 ++++++++++ stable_baselines3/common/env_checker.py | 222 ++++++++++++++++++ stable_baselines3/common/vec_env/__init__.py | 2 + .../common/vec_env/vec_check_nan.py | 86 +++++++ .../common/vec_env/vec_video_recorder.py | 112 +++++++++ stable_baselines3/version.txt | 2 +- tests/test_envs.py | 149 ++++++++++++ tests/test_vec_check_nan.py | 72 ++++++ 14 files changed, 832 insertions(+), 10 deletions(-) create mode 100644 stable_baselines3/common/bit_flipping_env.py create mode 100644 stable_baselines3/common/env_checker.py create mode 100644 stable_baselines3/common/vec_env/vec_check_nan.py create mode 100644 stable_baselines3/common/vec_env/vec_video_recorder.py create mode 100644 tests/test_envs.py create mode 100644 tests/test_vec_check_nan.py diff --git a/.coveragerc b/.coveragerc index 9020116..8ee0fb9 100644 --- a/.coveragerc +++ b/.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 = diff --git a/Makefile b/Makefile index 5a31e2d..4d95b7d 100644 --- a/Makefile +++ b/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/* diff --git a/README.md b/README.md index 5f35c82..bc0efa7 100644 --- a/README.md +++ b/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}}, } ``` diff --git a/docs/index.rst b/docs/index.rst index c97722b..9e74bd7 100644 --- a/docs/index.rst +++ b/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 `_ is the PyTorch version of `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 diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 3767b7c..86b7246 100644 --- a/docs/misc/changelog.rst +++ b/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 diff --git a/setup.py b/setup.py index 780f504..40aec1f 100644 --- a/setup.py +++ b/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", diff --git a/stable_baselines3/common/bit_flipping_env.py b/stable_baselines3/common/bit_flipping_env.py new file mode 100644 index 0000000..1dd8815 --- /dev/null +++ b/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) + """ + 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 diff --git a/stable_baselines3/common/env_checker.py b/stable_baselines3/common/env_checker.py new file mode 100644 index 0000000..751b204 --- /dev/null +++ b/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) diff --git a/stable_baselines3/common/vec_env/__init__.py b/stable_baselines3/common/vec_env/__init__.py index 2c89181..8f7719a 100644 --- a/stable_baselines3/common/vec_env/__init__.py +++ b/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: diff --git a/stable_baselines3/common/vec_env/vec_check_nan.py b/stable_baselines3/common/vec_env/vec_check_nan.py new file mode 100644 index 0000000..613be0d --- /dev/null +++ b/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) diff --git a/stable_baselines3/common/vec_env/vec_video_recorder.py b/stable_baselines3/common/vec_env/vec_video_recorder.py new file mode 100644 index 0000000..35409b0 --- /dev/null +++ b/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() diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 8f0916f..b7c431c 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.5.0 +0.6.0a0 diff --git a/tests/test_envs.py b/tests/test_envs.py new file mode 100644 index 0000000..9c11d71 --- /dev/null +++ b/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, {})) diff --git a/tests/test_vec_check_nan.py b/tests/test_vec_check_nan.py new file mode 100644 index 0000000..85dd884 --- /dev/null +++ b/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]]))