saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Sync with Stable-Baselines

Browse source
This commit is contained in:
Antonin RAFFIN 2020-05-05 16:28:38 +02:00
parent d542732c8d
commit 2c34a4d694
14 changed files with 832 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.coveragerc
View file

@ -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 =

23
Makefile
View file

@ -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/*

4
README.md
View file

@ -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}},
}
```

10
docs/index.rst
View file

@ -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

35
docs/misc/changelog.rst
View file

@ -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

2
setup.py
View file

@ -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",

121
stable_baselines3/common/bit_flipping_env.py Normal file
View file

@ -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

222
stable_baselines3/common/env_checker.py Normal file
View file

@ -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)

2
stable_baselines3/common/vec_env/__init__.py
View file

@ -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:

86
stable_baselines3/common/vec_env/vec_check_nan.py Normal file
View file

@ -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)

112
stable_baselines3/common/vec_env/vec_video_recorder.py Normal file
View file

@ -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()

2
stable_baselines3/version.txt
View file

@ -1 +1 @@
0.5.0
0.6.0a0

149
tests/test_envs.py Normal file
View file

@ -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, {}))

72
tests/test_vec_check_nan.py Normal file
View file

@ -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]]))