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

Merge pull request #75 from Antonin-Raffin/feat/open-source

Browse source 
Prepare open source
This commit is contained in:
Raffin, Antonin 2020-05-05 17:56:16 +02:00 committed by GitHub Enterprise
commit 7c22225391
81 changed files with 1095 additions and 223 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.coveragerc
View file

@ -4,7 +4,9 @@ omit =
tests/* tests/*
setup.py setup.py
# Require graphical interface # Require graphical interface
torchy_baselines/common/results_plotter.py stable_baselines3/common/results_plotter.py
# Require ffmpeg
stable_baselines3/common/vec_env/vec_video_recorder.py
[report] [report]
exclude_lines = exclude_lines =

4
.github/ISSUE_TEMPLATE/issue-template.md vendored
View file

@ -13,7 +13,7 @@ If you are submitting a bug report, please fill in the following details.
If your issue is related to a custom gym environment, please check it first using: If your issue is related to a custom gym environment, please check it first using:
```python ```python
from torchy_baselines.common.env_checker import check_env from stable_baselines3.common.env_checker import check_env
env = CustomEnv(arg1, ...) env = CustomEnv(arg1, ...)
# It will check your custom environment and output additional warnings if needed # It will check your custom environment and output additional warnings if needed
@ -30,7 +30,7 @@ Please use the [markdown code blocks](https://help.github.com/en/articles/creati
for both code and stack traces. for both code and stack traces.
```python ```python
from torchy_baselines import ... from stable_baselines3 import ...
``` ```

23
Makefile
View file

@ -16,3 +16,26 @@ clean:
cd docs && make clean cd docs && make clean
.PHONY: clean spelling doc .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/*

2
NOTICE
View file

@ -1,4 +1,4 @@
Large portion of the code of Torchy-Baselines (in `common/`) were ported from Stable-Baselines, a fork of OpenAI Baselines, Large portion of the code of Stable-Baselines3 (in `common/`) were ported from Stable-Baselines, a fork of OpenAI Baselines,
both licensed under the MIT License: both licensed under the MIT License:
before the fork (June 2018): before the fork (June 2018):

8
README.md
View file

@ -2,7 +2,7 @@
[![Build Status](https://travis-ci.com/hill-a/stable-baselines.svg?branch=master)](https://travis-ci.com/hill-a/stable-baselines) [![Documentation Status](https://readthedocs.org/projects/stable-baselines/badge/?version=master)](https://stable-baselines.readthedocs.io/en/master/?badge=master) [![Build Status](https://travis-ci.com/hill-a/stable-baselines.svg?branch=master)](https://travis-ci.com/hill-a/stable-baselines) [![Documentation Status](https://readthedocs.org/projects/stable-baselines/badge/?version=master)](https://stable-baselines.readthedocs.io/en/master/?badge=master)
# Torchy Baselines # Stable Baselines3
PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines), a set of improved implementations of reinforcement learning algorithms. PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines), a set of improved implementations of reinforcement learning algorithms.
@ -56,12 +56,12 @@ make spelling
To cite this repository in publications: 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}, author = {Raffin, Antonin and Hill, Ashley and Ernestus, Maximilian and Gleave, Adam and Kanervisto, Anssi and Dormann, Noah},
title = {Torchy Baselines}, title = {Stable Baselines3},
year = {2019}, year = {2019},
publisher = {GitHub}, publisher = {GitHub},
journal = {GitHub repository}, journal = {GitHub repository},
howpublished = {\url{https://github.com/araffin/torchy-baselines}}, howpublished = {\url{https://github.com/DLR-RM/stable-baselines3}},
} }
``` ```

28
docs/conf.py
View file

@ -44,19 +44,19 @@ MOCK_MODULES = []
sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES) sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)
import torchy_baselines import stable_baselines3
# -- Project information ----------------------------------------------------- # -- Project information -----------------------------------------------------
project = 'Torchy Baselines' project = 'Stable Baselines3'
copyright = '2020, Torchy Baselines' copyright = '2020, Stable Baselines3'
author = 'Torchy Baselines Contributors' author = 'Stable Baselines3 Contributors'
# The short X.Y version # The short X.Y version
version = 'master (' + torchy_baselines.__version__ + ' )' version = 'master (' + stable_baselines3.__version__ + ' )'
# The full version, including alpha/beta/rc tags # The full version, including alpha/beta/rc tags
release = torchy_baselines.__version__ release = stable_baselines3.__version__
# -- General configuration --------------------------------------------------- # -- General configuration ---------------------------------------------------
@ -152,7 +152,7 @@ html_static_path = ['_static']
# -- Options for HTMLHelp output --------------------------------------------- # -- Options for HTMLHelp output ---------------------------------------------
# Output file base name for HTML help builder. # Output file base name for HTML help builder.
htmlhelp_basename = 'TorchyBaselinesdoc' htmlhelp_basename = 'StableBaselines3doc'
# -- Options for LaTeX output ------------------------------------------------ # -- Options for LaTeX output ------------------------------------------------
@ -179,8 +179,8 @@ latex_elements = {
# (source start file, target name, title, # (source start file, target name, title,
# author, documentclass [howto, manual, or own class]). # author, documentclass [howto, manual, or own class]).
latex_documents = [ latex_documents = [
(master_doc, 'TorchyBaselines.tex', 'Torchy Baselines Documentation', (master_doc, 'StableBaselines3.tex', 'Stable Baselines3 Documentation',
'Torchy Baselines Contributors', 'manual'), 'Stable Baselines3 Contributors', 'manual'),
] ]
@ -189,7 +189,7 @@ latex_documents = [
# One entry per manual page. List of tuples # One entry per manual page. List of tuples
# (source start file, name, description, authors, manual section). # (source start file, name, description, authors, manual section).
man_pages = [ man_pages = [
(master_doc, 'torchybaselines', 'Torchy Baselines Documentation', (master_doc, 'stablebaselines3', 'Stable Baselines3 Documentation',
[author], 1) [author], 1)
] ]
@ -200,8 +200,8 @@ man_pages = [
# (source start file, target name, title, author, # (source start file, target name, title, author,
# dir menu entry, description, category) # dir menu entry, description, category)
texinfo_documents = [ texinfo_documents = [
(master_doc, 'TorchyBaselines', 'Torchy Baselines Documentation', (master_doc, 'StableBaselines3', 'Stable Baselines3 Documentation',
author, 'TorchyBaselines', 'One line description of project.', author, 'StableBaselines3', 'One line description of project.',
'Miscellaneous'), 'Miscellaneous'),
] ]
@ -214,3 +214,7 @@ texinfo_documents = [
# 'numpy': ('http://docs.scipy.org/doc/numpy/', None), # 'numpy': ('http://docs.scipy.org/doc/numpy/', None),
# 'torch': ('http://pytorch.org/docs/master/', None), # 'torch': ('http://pytorch.org/docs/master/', None),
# } # }
# kornia's hack to get rtd builder to install latest pytorch
# if 'READTHEDOCS' in os.environ:
# os.system('pip install torch==1.5.0+cpu torchvision==0.6.0+cpu -f https://download.pytorch.org/whl/torch_stable.html')

8
docs/guide/quickstart.rst
View file

@ -12,9 +12,9 @@ Here is a quick example of how to train and run SAC on a Pendulum environment:
import gym import gym
from torchy_baselines.sac.policies import MlpPolicy from stable_baselines3.sac.policies import MlpPolicy
from torchy_baselines.common.vec_env import DummyVecEnv from stable_baselines3.common.vec_env import DummyVecEnv
from torchy_baselines import SAC from stable_baselines3 import SAC
env = gym.make('Pendulum-v0') env = gym.make('Pendulum-v0')
@ -34,6 +34,6 @@ the policy is registered:
.. code-block:: python .. code-block:: python
from torchy_baselines import SAC from stable_baselines3 import SAC
model = SAC('MlpPolicy', 'Pendulum-v0').learn(10000) model = SAC('MlpPolicy', 'Pendulum-v0').learn(10000)

2
docs/guide/vec_envs.rst
View file

@ -1,6 +1,6 @@
.. _vec_env: .. _vec_env:
.. automodule:: torchy_baselines.common.vec_env .. automodule:: stable_baselines3.common.vec_env
Vectorized Environments Vectorized Environments
======================= =======================

20
docs/index.rst
View file

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

39
docs/misc/changelog.rst
View file

@ -3,6 +3,34 @@
Changelog 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)
- Added determinism tests
Bug Fixes:
^^^^^^^^^^
- Fixed a bug that prevented model trained on cpu to be loaded on gpu
Deprecations:
^^^^^^^^^^^^^
Others:
^^^^^^^
- Renamed to Stable-Baseline3
Documentation:
^^^^^^^^^^^^^^
Pre-Release 0.5.0 (2020-05-05) Pre-Release 0.5.0 (2020-05-05)
------------------------------ ------------------------------
@ -111,7 +139,7 @@ Pre-Release 0.2.0 (2020-02-14)
Breaking Changes: Breaking Changes:
^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^
- Python 2 support was dropped, Torchy Baselines now requires Python 3.6 or above - Python 2 support was dropped, Stable Baselines3 now requires Python 3.6 or above
- Return type of ``evaluation.evaluate_policy()`` has been changed - Return type of ``evaluation.evaluate_policy()`` has been changed
- Refactored the replay buffer to avoid transformation between PyTorch and NumPy - Refactored the replay buffer to avoid transformation between PyTorch and NumPy
- Created `OffPolicyRLModel` base class - Created `OffPolicyRLModel` base class
@ -160,10 +188,17 @@ New Features:
Maintainers Maintainers
----------- -----------
Torchy-Baselines 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/ .. _Antonin Raffin: https://araffin.github.io/
.. _Maximilian Ernestus: https://github.com/erniejunior
.. _Adam Gleave: https://gleave.me/
.. _@araffin: https://github.com/araffin .. _@araffin: https://github.com/araffin
.. _@AdamGleave: https://github.com/adamgleave
.. _Anssi Kanervisto: https://github.com/Miffyli
.. _@Miffyli: https://github.com/Miffyli

8
docs/modules/a2c.rst
View file

@ -1,6 +1,6 @@
.. _a2c: .. _a2c:
.. automodule:: torchy_baselines.a2c .. automodule:: stable_baselines3.a2c
A2C A2C
@ -44,9 +44,9 @@ Train a A2C agent on `CartPole-v1` using 4 processes.
import gym import gym
from torchy_baselines.common.policies import MlpPolicy from stable_baselines3.common.policies import MlpPolicy
from torchy_baselines.common import make_vec_env from stable_baselines3.common import make_vec_env
from torchy_baselines import A2C from stable_baselines3 import A2C
# Parallel environments # Parallel environments
env = make_vec_env('CartPole-v1', n_envs=4) env = make_vec_env('CartPole-v1', n_envs=4)

2
docs/modules/base.rst
View file

@ -1,6 +1,6 @@
.. _base_algo: .. _base_algo:
.. automodule:: torchy_baselines.common.base_class .. automodule:: stable_baselines3.common.base_class
Base RL Class Base RL Class

8
docs/modules/ppo.rst
View file

@ -1,6 +1,6 @@
.. _ppo2: .. _ppo2:
.. automodule:: torchy_baselines.ppo .. automodule:: stable_baselines3.ppo
PPO PPO
=== ===
@ -53,9 +53,9 @@ Train a PPO agent on `Pendulum-v0` using 4 processes.
import gym import gym
from torchy_baselines.ppo.policies import MlpPolicy from stable_baselines3.ppo.policies import MlpPolicy
from torchy_baselines.common.vec_env import SubprocVecEnv from stable_baselines3.common.vec_env import SubprocVecEnv
from torchy_baselines import PPO from stable_baselines3 import PPO
# multiprocess environment # multiprocess environment
n_cpu = 4 n_cpu = 4

10
docs/modules/sac.rst
View file

@ -1,6 +1,6 @@
.. _sac: .. _sac:
.. automodule:: torchy_baselines.sac .. automodule:: stable_baselines3.sac
SAC SAC
@ -14,7 +14,7 @@ A key feature of SAC, and a major difference with common RL algorithms, is that
.. warning:: .. warning::
The SAC model does not support ``torchy_baselines.common.policies`` because it uses double q-values The SAC model does not support ``stable_baselines3.common.policies`` because it uses double q-values
and value estimation, as a result it must use its own policy models (see :ref:`sac_policies`). and value estimation, as a result it must use its own policy models (see :ref:`sac_policies`).
@ -72,9 +72,9 @@ Example
import gym import gym
import numpy as np import numpy as np
from torchy_baselines.sac.policies import MlpPolicy from stable_baselines3.sac.policies import MlpPolicy
from torchy_baselines.common.vec_env import DummyVecEnv from stable_baselines3.common.vec_env import DummyVecEnv
from torchy_baselines import SAC from stable_baselines3 import SAC
env = gym.make('Pendulum-v0') env = gym.make('Pendulum-v0')
env = DummyVecEnv([lambda: env]) env = DummyVecEnv([lambda: env])

10
docs/modules/td3.rst
View file

@ -1,6 +1,6 @@
.. _td3: .. _td3:
.. automodule:: torchy_baselines.td3 .. automodule:: stable_baselines3.td3
TD3 TD3
@ -14,7 +14,7 @@ We recommend reading `OpenAI Spinning guide on TD3 <https://spinningup.openai.co
.. warning:: .. warning::
The TD3 model does not support ``torchy_baselines.common.policies`` because it uses double q-values The TD3 model does not support ``stable_baselines3.common.policies`` because it uses double q-values
estimation, as a result it must use its own policy models (see :ref:`td3_policies`). estimation, as a result it must use its own policy models (see :ref:`td3_policies`).
@ -64,9 +64,9 @@ Example
import numpy as np import numpy as np
from torchy_baselines import TD3 from stable_baselines3 import TD3
from torchy_baselines.td3.policies import MlpPolicy from stable_baselines3.td3.policies import MlpPolicy
from torchy_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
# The noise objects for TD3 # The noise objects for TD3
n_actions = env.action_space.shape[-1] n_actions = env.action_space.shape[-1]

2
setup.cfg
View file

@ -18,4 +18,4 @@ filterwarnings =
ignore::UserWarning:gym ignore::UserWarning:gym
[pytype] [pytype]
inputs = torchy_baselines inputs = stable_baselines3

8
setup.py
View file

@ -3,13 +3,13 @@ import sys
import subprocess import subprocess
from setuptools import setup, find_packages from setuptools import setup, find_packages
with open(os.path.join('torchy_baselines', 'version.txt'), 'r') as file_handler: with open(os.path.join('stable_baselines3', 'version.txt'), 'r') as file_handler:
__version__ = file_handler.read() __version__ = file_handler.read()
setup(name='torchy_baselines', setup(name='stable_baselines3',
packages=[package for package in find_packages() packages=[package for package in find_packages()
if package.startswith('torchy_baselines')], if package.startswith('stable_baselines3')],
install_requires=[ install_requires=[
'gym[classic_control]>=0.11', 'gym[classic_control]>=0.11',
'numpy', 'numpy',
@ -45,7 +45,7 @@ setup(name='torchy_baselines',
}, },
description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.', description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',
author='Antonin Raffin', author='Antonin Raffin',
url='', url='https://github.com/DLR-RM/stable-baselines3',
author_email='antonin.raffin@dlr.de', author_email='antonin.raffin@dlr.de',
keywords="reinforcement-learning-algorithms reinforcement-learning machine-learning " keywords="reinforcement-learning-algorithms reinforcement-learning machine-learning "
"gym openai stable baselines toolbox python data-science", "gym openai stable baselines toolbox python data-science",

8
torchy_baselines/__init__.py → stable_baselines3/__init__.py
View file

@ -1,9 +1,9 @@
import os import os
from torchy_baselines.a2c import A2C from stable_baselines3.a2c import A2C
from torchy_baselines.ppo import PPO from stable_baselines3.ppo import PPO
from torchy_baselines.sac import SAC from stable_baselines3.sac import SAC
from torchy_baselines.td3 import TD3 from stable_baselines3.td3 import TD3
# Read version from file # Read version from file
version_file = os.path.join(os.path.dirname(__file__), 'version.txt') version_file = os.path.join(os.path.dirname(__file__), 'version.txt')

2
stable_baselines3/a2c/__init__.py Normal file
View file

@ -0,0 +1,2 @@
from stable_baselines3.a2c.a2c import A2C
from stable_baselines3.ppo.policies import MlpPolicy

10
torchy_baselines/a2c/a2c.py → stable_baselines3/a2c/a2c.py
View file

@ -3,11 +3,11 @@ import torch.nn.functional as F
from gym import spaces from gym import spaces
from typing import Type, Union, Callable, Optional, Dict, Any from typing import Type, Union, Callable, Optional, Dict, Any
from torchy_baselines.common import logger from stable_baselines3.common import logger
from torchy_baselines.common.type_aliases import GymEnv, MaybeCallback from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from torchy_baselines.common.utils import explained_variance from stable_baselines3.common.utils import explained_variance
from torchy_baselines.ppo.policies import PPOPolicy from stable_baselines3.ppo.policies import PPOPolicy
from torchy_baselines.ppo.ppo import PPO from stable_baselines3.ppo.ppo import PPO
class A2C(PPO): class A2C(PPO):

0
torchy_baselines/common/__init__.py → stable_baselines3/common/__init__.py
View file

51
torchy_baselines/common/base_class.py → stable_baselines3/common/base_class.py
View file

@ -11,17 +11,17 @@ import gym
import torch as th import torch as th
import numpy as np import numpy as np
from torchy_baselines.common import logger from stable_baselines3.common import logger
from torchy_baselines.common.policies import BasePolicy, get_policy_from_name from stable_baselines3.common.policies import BasePolicy, get_policy_from_name
from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate, get_device from stable_baselines3.common.utils import set_random_seed, get_schedule_fn, update_learning_rate, get_device
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize, VecTransposeImage from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize, VecTransposeImage
from torchy_baselines.common.preprocessing import is_image_space from stable_baselines3.common.preprocessing import is_image_space
from torchy_baselines.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr from stable_baselines3.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr
from torchy_baselines.common.type_aliases import GymEnv, TensorDict, RolloutReturn, MaybeCallback from stable_baselines3.common.type_aliases import GymEnv, TensorDict, RolloutReturn, MaybeCallback
from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback from stable_baselines3.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback
from torchy_baselines.common.monitor import Monitor from stable_baselines3.common.monitor import Monitor
from torchy_baselines.common.noise import ActionNoise from stable_baselines3.common.noise import ActionNoise
from torchy_baselines.common.buffers import ReplayBuffer from stable_baselines3.common.buffers import ReplayBuffer
class BaseRLModel(ABC): class BaseRLModel(ABC):
@ -219,9 +219,9 @@ class BaseRLModel(ABC):
def get_vec_normalize_env(self) -> Optional[VecNormalize]: def get_vec_normalize_env(self) -> Optional[VecNormalize]:
""" """
Return the `VecNormalize` wrapper of the training env Return the ``VecNormalize`` wrapper of the training env
if it exists. if it exists.
:return: Optional[VecNormalize] The `VecNormalize` env. :return: Optional[VecNormalize] The ``VecNormalize`` env.
""" """
return self._vec_normalize_env return self._vec_normalize_env
@ -267,7 +267,7 @@ class BaseRLModel(ABC):
def get_torch_variables(self) -> Tuple[List[str], List[str]]: def get_torch_variables(self) -> Tuple[List[str], List[str]]:
""" """
Get the name of the torch variable that will be saved. Get the name of the torch variable that will be saved.
`th.save` and `th.load` will be used with the right device ``th.save`` and ``th.load`` will be used with the right device
instead of the default pickling strategy. instead of the default pickling strategy.
:return: (Tuple[List[str], List[str]]) :return: (Tuple[List[str], List[str]])
@ -297,7 +297,7 @@ class BaseRLModel(ABC):
:param tb_log_name: (str) the name of the run for tensorboard log :param tb_log_name: (str) the name of the run for tensorboard log
:param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging) :param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging)
:param eval_env: (gym.Env) Environment that will be used to evaluate the agent :param eval_env: (gym.Env) Environment that will be used to evaluate the agent
:param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little) :param eval_freq: (int) Evaluate the agent every ``eval_freq`` timesteps (this may vary a little)
:param n_eval_episodes: (int) Number of episode to evaluate the agent :param n_eval_episodes: (int) Number of episode to evaluate the agent
:param eval_log_path: (Optional[str]) Path to a folder where the evaluations will be saved :param eval_log_path: (Optional[str]) Path to a folder where the evaluations will be saved
:param reset_num_timesteps: (bool) :param reset_num_timesteps: (bool)
@ -333,6 +333,11 @@ class BaseRLModel(ABC):
""" """
data, params, tensors = cls._load_from_file(load_path) data, params, tensors = cls._load_from_file(load_path)
if 'policy_kwargs' in data:
for arg_to_remove in ['device']:
if arg_to_remove in data['policy_kwargs']:
del data['policy_kwargs'][arg_to_remove]
if 'policy_kwargs' in kwargs and kwargs['policy_kwargs'] != data['policy_kwargs']: if 'policy_kwargs' in kwargs and kwargs['policy_kwargs'] != data['policy_kwargs']:
raise ValueError(f"The specified policy kwargs do not equal the stored policy kwargs." raise ValueError(f"The specified policy kwargs do not equal the stored policy kwargs."
f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}") f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}")
@ -354,7 +359,7 @@ class BaseRLModel(ABC):
model.__dict__.update(data) model.__dict__.update(data)
model.__dict__.update(kwargs) model.__dict__.update(kwargs)
if not hasattr(model, "_setup_model") and len(params) > 0: if not hasattr(model, "_setup_model") and len(params) > 0:
raise NotImplementedError(f"{cls} has no `_setup_model()` method") raise NotImplementedError(f"{cls} has no ``_setup_model()`` method")
model._setup_model() model._setup_model()
# put state_dicts back in place # put state_dicts back in place
@ -417,7 +422,7 @@ class BaseRLModel(ABC):
file_content.write(tensor_file.read()) file_content.write(tensor_file.read())
# go to start of file # go to start of file
file_content.seek(0) file_content.seek(0)
# load the parameters with the right `map_location` # load the parameters with the right ``map_location``
tensors = th.load(file_content, map_location=device) tensors = th.load(file_content, map_location=device)
# check for all other .pth files # check for all other .pth files
@ -434,7 +439,7 @@ class BaseRLModel(ABC):
file_content.write(opt_param_file.read()) file_content.write(opt_param_file.read())
# go to start of file # go to start of file
file_content.seek(0) file_content.seek(0)
# load the parameters with the right `map_location` # load the parameters with the right ``map_location``
params[os.path.splitext(file_path)[0]] = th.load(file_content, map_location=device) params[os.path.splitext(file_path)[0]] = th.load(file_content, map_location=device)
except zipfile.BadZipFile: except zipfile.BadZipFile:
@ -502,7 +507,7 @@ class BaseRLModel(ABC):
:param eval_freq: (int) :param eval_freq: (int)
:param n_eval_episodes: (int) :param n_eval_episodes: (int)
:param log_path (Optional[str]): Path to a log folder :param log_path (Optional[str]): Path to a log folder
:param reset_num_timesteps: (bool) Whether to reset or not the `num_timesteps` attribute :param reset_num_timesteps: (bool) Whether to reset or not the ``num_timesteps`` attribute
:return: (BaseCallback) :return: (BaseCallback)
""" """
self.start_time = time.time() self.start_time = time.time()
@ -516,7 +521,7 @@ class BaseRLModel(ABC):
self.num_timesteps = 0 self.num_timesteps = 0
self._episode_num = 0 self._episode_num = 0
# Avoid resetting the environment when calling `.learn()` consecutive times # Avoid resetting the environment when calling ``.learn()`` consecutive times
if reset_num_timesteps or self._last_obs is None: if reset_num_timesteps or self._last_obs is None:
self._last_obs = self.env.reset() self._last_obs = self.env.reset()
# Retrieve unnormalized observation for saving into the buffer # Retrieve unnormalized observation for saving into the buffer
@ -762,9 +767,9 @@ class OffPolicyRLModel(BaseRLModel):
:param env: (VecEnv) The training environment :param env: (VecEnv) The training environment
:param n_episodes: (int) Number of episodes to use to collect rollout data :param n_episodes: (int) Number of episodes to use to collect rollout data
You can also specify a `n_steps` instead You can also specify a ``n_steps`` instead
:param n_steps: (int) Number of steps to use to collect rollout data :param n_steps: (int) Number of steps to use to collect rollout data
You can also specify a `n_episodes` instead. You can also specify a ``n_episodes`` instead.
:param action_noise: (Optional[ActionNoise]) Action noise that will be used for exploration :param action_noise: (Optional[ActionNoise]) Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC. in addition to the stochastic policy for SAC.
@ -772,7 +777,7 @@ class OffPolicyRLModel(BaseRLModel):
(and at the beginning and end of the rollout) (and at the beginning and end of the rollout)
:param learning_starts: (int) Number of steps before learning for the warm-up phase. :param learning_starts: (int) Number of steps before learning for the warm-up phase.
:param replay_buffer: (ReplayBuffer) :param replay_buffer: (ReplayBuffer)
:param log_interval: (int) Log data every `log_interval` episodes :param log_interval: (int) Log data every ``log_interval`` episodes
:return: (RolloutReturn) :return: (RolloutReturn)
""" """
episode_rewards, total_timesteps = [], [] episode_rewards, total_timesteps = [], []

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

6
torchy_baselines/common/buffers.py → stable_baselines3/common/buffers.py
View file

@ -4,9 +4,9 @@ import numpy as np
import torch as th import torch as th
from gym import spaces from gym import spaces
from torchy_baselines.common.vec_env import VecNormalize from stable_baselines3.common.vec_env import VecNormalize
from torchy_baselines.common.type_aliases import RolloutBufferSamples, ReplayBufferSamples from stable_baselines3.common.type_aliases import RolloutBufferSamples, ReplayBufferSamples
from torchy_baselines.common.preprocessing import get_action_dim, get_obs_shape from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
class BaseBuffer(object): class BaseBuffer(object):

8
torchy_baselines/common/callbacks.py → stable_baselines3/common/callbacks.py
View file

@ -7,12 +7,12 @@ from typing import Union, List, Dict, Any, Optional
import gym import gym
import numpy as np import numpy as np
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, sync_envs_normalization from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, sync_envs_normalization
from torchy_baselines.common.evaluation import evaluate_policy from stable_baselines3.common.evaluation import evaluate_policy
from torchy_baselines.common.logger import Logger from stable_baselines3.common.logger import Logger
if typing.TYPE_CHECKING: if typing.TYPE_CHECKING:
from torchy_baselines.common.base_class import BaseRLModel # pytype: disable=pyi-error from stable_baselines3.common.base_class import BaseRLModel # pytype: disable=pyi-error
class BaseCallback(ABC): class BaseCallback(ABC):

2
torchy_baselines/common/distributions.py → stable_baselines3/common/distributions.py
View file

@ -6,7 +6,7 @@ import torch.nn as nn
from torch.distributions import Normal, Categorical from torch.distributions import Normal, Categorical
from gym import spaces from gym import spaces
from torchy_baselines.common.preprocessing import get_action_dim from stable_baselines3.common.preprocessing import get_action_dim
class Distribution(object): class Distribution(object):

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_baselines3.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
torchy_baselines/common/evaluation.py → stable_baselines3/common/evaluation.py
View file

@ -1,7 +1,7 @@
# Copied from stable_baselines # Copied from stable_baselines
import numpy as np import numpy as np
from torchy_baselines.common.vec_env import VecEnv from stable_baselines3.common.vec_env import VecEnv
def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,

68
torchy_baselines/common/identity_env.py → stable_baselines3/common/identity_env.py
View file

@ -1,30 +1,44 @@
from typing import List, Union from typing import List, Union, Optional
import numpy as np import numpy as np
from gym import Env from gym import Env, Space
from gym.spaces import Discrete, MultiDiscrete, MultiBinary, Box from gym.spaces import Discrete, MultiDiscrete, MultiBinary, Box
from torchy_baselines.common.type_aliases import GymStepReturn, GymObs from stable_baselines3.common.type_aliases import GymStepReturn, GymObs
class IdentityEnv(Env): class IdentityEnv(Env):
def __init__(self, dim, ep_length=100): def __init__(self,
dim: Optional[int] = None,
space: Optional[Space] = None,
ep_length: int = 100):
""" """
Identity environment for testing purposes Identity environment for testing purposes
:param dim: (int) the size of the dimensions you want to learn :param dim: the size of the action and observation dimension you want
:param ep_length: (int) the length of each episodes in timesteps to learn. Provide at most one of ``dim`` and ``space``. If both are
None, then initialization proceeds with ``dim=1`` and ``space=None``.
:param space: the action and observation space. Provide at most one of
``dim`` and ``space``.
:param ep_length: the length of each episode in timesteps
""" """
self.action_space = Discrete(dim) if space is None:
self.observation_space = self.action_space if dim is None:
dim = 1
space = Discrete(dim)
else:
assert dim is None, "arguments for both 'dim' and 'space' provided: at most one allowed"
self.action_space = self.observation_space = space
self.ep_length = ep_length self.ep_length = ep_length
self.current_step = 0 self.current_step = 0
self.dim = dim self.num_resets = -1 # Becomes 0 after __init__ exits.
self.reset() self.reset()
def reset(self) -> GymObs: def reset(self) -> GymObs:
self.current_step = 0 self.current_step = 0
self.num_resets += 1
self._choose_next_state() self._choose_next_state()
return self.state return self.state
@ -55,18 +69,11 @@ class IdentityEnvBox(IdentityEnv):
:param low: (float) the lower bound of the box dim :param low: (float) the lower bound of the box dim
:param high: (float) the upper bound of the box dim :param high: (float) the upper bound of the box dim
:param eps: (float) the epsilon bound for correct value :param eps: (float) the epsilon bound for correct value
:param ep_length: (int) the length of each episodes in timesteps :param ep_length: (int) the length of each episode in timesteps
""" """
super(IdentityEnvBox, self).__init__(1, ep_length) space = Box(low=low, high=high, shape=(1,), dtype=np.float32)
self.action_space = Box(low=low, high=high, shape=(1,), dtype=np.float32) super().__init__(ep_length=ep_length, space=space)
self.observation_space = self.action_space
self.eps = eps self.eps = eps
self.reset()
def reset(self) -> np.ndarray:
self.current_step = 0
self._choose_next_state()
return self.state
def step(self, action: np.ndarray) -> GymStepReturn: def step(self, action: np.ndarray) -> GymStepReturn:
reward = self._get_reward(action) reward = self._get_reward(action)
@ -75,39 +82,32 @@ class IdentityEnvBox(IdentityEnv):
done = self.current_step >= self.ep_length done = self.current_step >= self.ep_length
return self.state, reward, done, {} return self.state, reward, done, {}
def _choose_next_state(self) -> None:
self.state = self.observation_space.sample()
def _get_reward(self, action: np.ndarray) -> float: def _get_reward(self, action: np.ndarray) -> float:
return 1.0 if (self.state - self.eps) <= action <= (self.state + self.eps) else 0.0 return 1.0 if (self.state - self.eps) <= action <= (self.state + self.eps) else 0.0
class IdentityEnvMultiDiscrete(IdentityEnv): class IdentityEnvMultiDiscrete(IdentityEnv):
def __init__(self, dim: int, ep_length: int = 100): def __init__(self, dim: int = 1, ep_length: int = 100):
""" """
Identity environment for testing purposes Identity environment for testing purposes
:param dim: (int) the size of the dimensions you want to learn :param dim: (int) the size of the dimensions you want to learn
:param ep_length: (int) the length of each episodes in timesteps :param ep_length: (int) the length of each episode in timesteps
""" """
super(IdentityEnvMultiDiscrete, self).__init__(dim, ep_length) space = MultiDiscrete([dim, dim])
self.action_space = MultiDiscrete([dim, dim]) super().__init__(ep_length=ep_length, space=space)
self.observation_space = self.action_space
self.reset()
class IdentityEnvMultiBinary(IdentityEnv): class IdentityEnvMultiBinary(IdentityEnv):
def __init__(self, dim: int, ep_length: int = 100): def __init__(self, dim: int = 1, ep_length: int = 100):
""" """
Identity environment for testing purposes Identity environment for testing purposes
:param dim: (int) the size of the dimensions you want to learn :param dim: (int) the size of the dimensions you want to learn
:param ep_length: (int) the length of each episodes in timesteps :param ep_length: (int) the length of each episode in timesteps
""" """
super(IdentityEnvMultiBinary, self).__init__(dim, ep_length) space = MultiBinary(dim)
self.action_space = MultiBinary(dim) super().__init__(ep_length=ep_length, space=space)
self.observation_space = self.action_space
self.reset()
class FakeImageEnv(Env): class FakeImageEnv(Env):

0
torchy_baselines/common/logger.py → stable_baselines3/common/logger.py
View file

0
torchy_baselines/common/monitor.py → stable_baselines3/common/monitor.py
View file

0
torchy_baselines/common/noise.py → stable_baselines3/common/noise.py
View file

6
torchy_baselines/common/policies.py → stable_baselines3/common/policies.py
View file

@ -7,9 +7,9 @@ import torch as th
import torch.nn as nn import torch.nn as nn
import numpy as np import numpy as np
from torchy_baselines.common.preprocessing import preprocess_obs, get_flattened_obs_dim, is_image_space from stable_baselines3.common.preprocessing import preprocess_obs, get_flattened_obs_dim, is_image_space
from torchy_baselines.common.utils import get_device from stable_baselines3.common.utils import get_device
from torchy_baselines.common.vec_env import VecTransposeImage from stable_baselines3.common.vec_env import VecTransposeImage
class BaseFeaturesExtractor(nn.Module): class BaseFeaturesExtractor(nn.Module):

0
torchy_baselines/common/preprocessing.py → stable_baselines3/common/preprocessing.py
View file

2
torchy_baselines/common/results_plotter.py → stable_baselines3/common/results_plotter.py
View file

@ -6,7 +6,7 @@ import pandas as pd
# matplotlib.use('TkAgg') # Can change to 'Agg' for non-interactive mode # matplotlib.use('TkAgg') # Can change to 'Agg' for non-interactive mode
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
from torchy_baselines.common.monitor import load_results from stable_baselines3.common.monitor import load_results
X_TIMESTEPS = 'timesteps' X_TIMESTEPS = 'timesteps'

0
torchy_baselines/common/running_mean_std.py → stable_baselines3/common/running_mean_std.py
View file

0
torchy_baselines/common/save_util.py → stable_baselines3/common/save_util.py
View file

4
torchy_baselines/common/type_aliases.py → stable_baselines3/common/type_aliases.py
View file

@ -7,8 +7,8 @@ import numpy as np
import torch as th import torch as th
import gym import gym
from torchy_baselines.common.vec_env import VecEnv from stable_baselines3.common.vec_env import VecEnv
from torchy_baselines.common.callbacks import BaseCallback from stable_baselines3.common.callbacks import BaseCallback
GymEnv = Union[gym.Env, VecEnv] GymEnv = Union[gym.Env, VecEnv]

7
torchy_baselines/common/utils.py → stable_baselines3/common/utils.py
View file

@ -11,14 +11,17 @@ def set_random_seed(seed: int, using_cuda: bool = False) -> None:
:param seed: (int) :param seed: (int)
:param using_cuda: (bool) :param using_cuda: (bool)
""" """
# Seed python RNG
random.seed(seed) random.seed(seed)
# Seed numpy RNG
np.random.seed(seed) np.random.seed(seed)
# seed the RNG for all devices (both CPU and CUDA)
th.manual_seed(seed) th.manual_seed(seed)
if using_cuda: if using_cuda:
# Make CuDNN Determinist # Deterministic operations for CuDNN, it may impact performances
th.backends.cudnn.deterministic = True th.backends.cudnn.deterministic = True
th.cuda.manual_seed(seed) th.backends.cudnn.benchmark = False
# From stable baselines # From stable baselines

16
torchy_baselines/common/vec_env/__init__.py → stable_baselines3/common/vec_env/__init__.py
View file

@ -3,17 +3,19 @@ import typing
from typing import Optional, Union from typing import Optional, Union
from copy import deepcopy from copy import deepcopy
from torchy_baselines.common.vec_env.base_vec_env import (AlreadySteppingError, NotSteppingError, from stable_baselines3.common.vec_env.base_vec_env import (AlreadySteppingError, NotSteppingError,
VecEnv, VecEnvWrapper, CloudpickleWrapper) VecEnv, VecEnvWrapper, CloudpickleWrapper)
from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv
from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack from stable_baselines3.common.vec_env.vec_frame_stack import VecFrameStack
from torchy_baselines.common.vec_env.vec_normalize import VecNormalize from stable_baselines3.common.vec_env.vec_normalize import VecNormalize
from torchy_baselines.common.vec_env.vec_transpose import VecTransposeImage 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 # Avoid circular import
if typing.TYPE_CHECKING: if typing.TYPE_CHECKING:
from torchy_baselines.common.type_aliases import GymEnv from stable_baselines3.common.type_aliases import GymEnv
def unwrap_vec_normalize(env: Union['GymEnv', VecEnv]) -> Optional[VecNormalize]: def unwrap_vec_normalize(env: Union['GymEnv', VecEnv]) -> Optional[VecNormalize]:

0
torchy_baselines/common/vec_env/base_vec_env.py → stable_baselines3/common/vec_env/base_vec_env.py
View file

4
torchy_baselines/common/vec_env/dummy_vec_env.py → stable_baselines3/common/vec_env/dummy_vec_env.py
View file

@ -3,8 +3,8 @@ from copy import deepcopy
import numpy as np import numpy as np
from torchy_baselines.common.vec_env.base_vec_env import VecEnv from stable_baselines3.common.vec_env.base_vec_env import VecEnv
from torchy_baselines.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info from stable_baselines3.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info
class DummyVecEnv(VecEnv): class DummyVecEnv(VecEnv):

2
torchy_baselines/common/vec_env/subproc_vec_env.py → stable_baselines3/common/vec_env/subproc_vec_env.py
View file

@ -4,7 +4,7 @@ from collections import OrderedDict
import gym import gym
import numpy as np import numpy as np
from torchy_baselines.common.vec_env.base_vec_env import VecEnv, CloudpickleWrapper from stable_baselines3.common.vec_env.base_vec_env import VecEnv, CloudpickleWrapper
def _worker(remote, parent_remote, env_fn_wrapper): def _worker(remote, parent_remote, env_fn_wrapper):

0
torchy_baselines/common/vec_env/util.py → stable_baselines3/common/vec_env/util.py
View file

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)

2
torchy_baselines/common/vec_env/vec_frame_stack.py → stable_baselines3/common/vec_env/vec_frame_stack.py
View file

@ -3,7 +3,7 @@ import warnings
import numpy as np import numpy as np
from gym import spaces from gym import spaces
from torchy_baselines.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper
class VecFrameStack(VecEnvWrapper): class VecFrameStack(VecEnvWrapper):

4
torchy_baselines/common/vec_env/vec_normalize.py → stable_baselines3/common/vec_env/vec_normalize.py
View file

@ -2,8 +2,8 @@ import pickle
import numpy as np import numpy as np
from torchy_baselines.common.vec_env.base_vec_env import VecEnvWrapper from stable_baselines3.common.vec_env.base_vec_env import VecEnvWrapper
from torchy_baselines.common.running_mean_std import RunningMeanStd from stable_baselines3.common.running_mean_std import RunningMeanStd
class VecNormalize(VecEnvWrapper): class VecNormalize(VecEnvWrapper):

6
torchy_baselines/common/vec_env/vec_transpose.py → stable_baselines3/common/vec_env/vec_transpose.py
View file

@ -2,11 +2,11 @@ import typing
import numpy as np import numpy as np
from gym import spaces from gym import spaces
from torchy_baselines.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper
from torchy_baselines.common.preprocessing import is_image_space from stable_baselines3.common.preprocessing import is_image_space
if typing.TYPE_CHECKING: if typing.TYPE_CHECKING:
from torchy_baselines.common.type_aliases import GymStepReturn from stable_baselines3.common.type_aliases import GymStepReturn
class VecTransposeImage(VecEnvWrapper): class VecTransposeImage(VecEnvWrapper):

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/ppo/__init__.py Normal file
View file

@ -0,0 +1,2 @@
from stable_baselines3.ppo.ppo import PPO
from stable_baselines3.ppo.policies import MlpPolicy

4
torchy_baselines/ppo/policies.py → stable_baselines3/ppo/policies.py
View file

@ -6,10 +6,10 @@ import torch as th
import torch.nn as nn import torch.nn as nn
import numpy as np import numpy as np
from torchy_baselines.common.policies import (BasePolicy, register_policy, MlpExtractor, from stable_baselines3.common.policies import (BasePolicy, register_policy, MlpExtractor,
create_sde_features_extractor, NatureCNN, create_sde_features_extractor, NatureCNN,
BaseFeaturesExtractor, FlattenExtractor) BaseFeaturesExtractor, FlattenExtractor)
from torchy_baselines.common.distributions import (make_proba_distribution, Distribution, from stable_baselines3.common.distributions import (make_proba_distribution, Distribution,
DiagGaussianDistribution, CategoricalDistribution, DiagGaussianDistribution, CategoricalDistribution,
StateDependentNoiseDistribution) StateDependentNoiseDistribution)

16
torchy_baselines/ppo/ppo.py → stable_baselines3/ppo/ppo.py
View file

@ -14,14 +14,14 @@ import torch.nn.functional as F
# SummaryWriter = None # SummaryWriter = None
import numpy as np import numpy as np
from torchy_baselines.common import logger from stable_baselines3.common import logger
from torchy_baselines.common.base_class import BaseRLModel from stable_baselines3.common.base_class import BaseRLModel
from torchy_baselines.common.type_aliases import GymEnv, MaybeCallback from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from torchy_baselines.common.buffers import RolloutBuffer from stable_baselines3.common.buffers import RolloutBuffer
from torchy_baselines.common.utils import explained_variance, get_schedule_fn from stable_baselines3.common.utils import explained_variance, get_schedule_fn
from torchy_baselines.common.vec_env import VecEnv from stable_baselines3.common.vec_env import VecEnv
from torchy_baselines.common.callbacks import BaseCallback from stable_baselines3.common.callbacks import BaseCallback
from torchy_baselines.ppo.policies import PPOPolicy from stable_baselines3.ppo.policies import PPOPolicy
class PPO(BaseRLModel): class PPO(BaseRLModel):

0
torchy_baselines/py.typed → stable_baselines3/py.typed
View file

2
stable_baselines3/sac/__init__.py Normal file
View file

@ -0,0 +1,2 @@
from stable_baselines3.sac.sac import SAC
from stable_baselines3.sac.policies import MlpPolicy

6
torchy_baselines/sac/policies.py → stable_baselines3/sac/policies.py
View file

@ -4,11 +4,11 @@ import gym
import torch as th import torch as th
import torch.nn as nn import torch.nn as nn
from torchy_baselines.common.preprocessing import get_action_dim from stable_baselines3.common.preprocessing import get_action_dim
from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp, from stable_baselines3.common.policies import (BasePolicy, register_policy, create_mlp,
create_sde_features_extractor, NatureCNN, create_sde_features_extractor, NatureCNN,
BaseFeaturesExtractor, FlattenExtractor) BaseFeaturesExtractor, FlattenExtractor)
from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution from stable_baselines3.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
# CAP the standard deviation of the actor # CAP the standard deviation of the actor
LOG_STD_MAX = 2 LOG_STD_MAX = 2

10
torchy_baselines/sac/sac.py → stable_baselines3/sac/sac.py
View file

@ -4,11 +4,11 @@ import torch as th
import torch.nn.functional as F import torch.nn.functional as F
import numpy as np import numpy as np
from torchy_baselines.common import logger from stable_baselines3.common import logger
from torchy_baselines.common.base_class import OffPolicyRLModel from stable_baselines3.common.base_class import OffPolicyRLModel
from torchy_baselines.common.type_aliases import GymEnv, MaybeCallback from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from torchy_baselines.common.noise import ActionNoise from stable_baselines3.common.noise import ActionNoise
from torchy_baselines.sac.policies import SACPolicy from stable_baselines3.sac.policies import SACPolicy
class SAC(OffPolicyRLModel): class SAC(OffPolicyRLModel):

2
stable_baselines3/td3/__init__.py Normal file
View file

@ -0,0 +1,2 @@
from stable_baselines3.td3.td3 import TD3
from stable_baselines3.td3.policies import MlpPolicy

6
torchy_baselines/td3/policies.py → stable_baselines3/td3/policies.py
View file

@ -4,11 +4,11 @@ import gym
import torch as th import torch as th
import torch.nn as nn import torch.nn as nn
from torchy_baselines.common.preprocessing import get_action_dim from stable_baselines3.common.preprocessing import get_action_dim
from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp, from stable_baselines3.common.policies import (BasePolicy, register_policy, create_mlp,
create_sde_features_extractor, NatureCNN, create_sde_features_extractor, NatureCNN,
BaseFeaturesExtractor, FlattenExtractor) BaseFeaturesExtractor, FlattenExtractor)
from torchy_baselines.common.distributions import StateDependentNoiseDistribution from stable_baselines3.common.distributions import StateDependentNoiseDistribution
class Actor(BasePolicy): class Actor(BasePolicy):

10
torchy_baselines/td3/td3.py → stable_baselines3/td3/td3.py
View file

@ -2,11 +2,11 @@ import torch as th
import torch.nn.functional as F import torch.nn.functional as F
from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
from torchy_baselines.common import logger from stable_baselines3.common import logger
from torchy_baselines.common.base_class import OffPolicyRLModel from stable_baselines3.common.base_class import OffPolicyRLModel
from torchy_baselines.common.noise import ActionNoise from stable_baselines3.common.noise import ActionNoise
from torchy_baselines.common.type_aliases import GymEnv, MaybeCallback from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from torchy_baselines.td3.policies import TD3Policy from stable_baselines3.td3.policies import TD3Policy
class TD3(OffPolicyRLModel): class TD3(OffPolicyRLModel):

1
stable_baselines3/version.txt Normal file
View file

@ -0,0 +1 @@
0.6.0a0

4
tests/test_callbacks.py
View file

@ -4,8 +4,8 @@ import shutil
import pytest import pytest
import gym import gym
from torchy_baselines import A2C, PPO, SAC, TD3 from stable_baselines3 import A2C, PPO, SAC, TD3
from torchy_baselines.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback, from stable_baselines3.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback,
EveryNTimesteps, StopTrainingOnRewardThreshold) EveryNTimesteps, StopTrainingOnRewardThreshold)

4
tests/test_cnn.py
View file

@ -3,8 +3,8 @@ import os
import numpy as np import numpy as np
import pytest import pytest
from torchy_baselines import A2C, PPO, SAC, TD3 from stable_baselines3 import A2C, PPO, SAC, TD3
from torchy_baselines.common.identity_env import FakeImageEnv from stable_baselines3.common.identity_env import FakeImageEnv
SAVE_PATH = './cnn_model.zip' SAVE_PATH = './cnn_model.zip'

2
tests/test_custom_policy.py
View file

@ -1,7 +1,7 @@
import pytest import pytest
import torch as th import torch as th
from torchy_baselines import A2C, PPO, SAC, TD3 from stable_baselines3 import A2C, PPO, SAC, TD3
@pytest.mark.parametrize('net_arch', [ @pytest.mark.parametrize('net_arch', [

36
tests/test_deterministic.py Normal file
View file

@ -0,0 +1,36 @@
import pytest
from stable_baselines3 import A2C, PPO, SAC, TD3
from stable_baselines3.common.noise import NormalActionNoise
N_STEPS_TRAINING = 3000
SEED = 0
@pytest.mark.parametrize("algo", [A2C, PPO, SAC, TD3])
def test_deterministic_training_common(algo):
results = [[], []]
rewards = [[], []]
# Smaller network
kwargs = {'policy_kwargs': dict(net_arch=[64])}
if algo in [TD3, SAC]:
env_id = 'Pendulum-v0'
kwargs.update({'action_noise': NormalActionNoise(0.0, 0.1),
'learning_starts': 100})
else:
env_id = 'CartPole-v1'
# if algo == DQN:
# kwargs.update({'learning_starts': 100})
for i in range(2):
model = algo('MlpPolicy', env_id, seed=SEED, **kwargs)
model.learn(N_STEPS_TRAINING)
env = model.get_env()
obs = env.reset()
for _ in range(100):
action, _ = model.predict(obs, deterministic=False)
obs, reward, _, _ = env.step(action)
results[i].append(action)
rewards[i].append(reward)
assert sum(results[0]) == sum(results[1]), results
assert sum(rewards[0]) == sum(rewards[1]), rewards

6
tests/test_distributions.py
View file

@ -1,11 +1,11 @@
import pytest import pytest
import torch as th import torch as th
from torchy_baselines import A2C, PPO from stable_baselines3 import A2C, PPO
from torchy_baselines.common.distributions import (DiagGaussianDistribution, TanhBijector, from stable_baselines3.common.distributions import (DiagGaussianDistribution, TanhBijector,
StateDependentNoiseDistribution, StateDependentNoiseDistribution,
CategoricalDistribution, SquashedDiagGaussianDistribution) CategoricalDistribution, SquashedDiagGaussianDistribution)
from torchy_baselines.common.utils import set_random_seed from stable_baselines3.common.utils import set_random_seed
N_ACTIONS = 2 N_ACTIONS = 2

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, FakeImageEnv,
IdentityEnvMultiBinary, IdentityEnvMultiDiscrete,)
ENV_CLASSES = [BitFlippingEnv, IdentityEnv, IdentityEnvBox, IdentityEnvMultiBinary,
IdentityEnvMultiDiscrete, FakeImageEnv]
@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 = FakeImageEnv()
# 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 = FakeImageEnv()
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, {}))

8
tests/test_identity.py
View file

@ -1,10 +1,10 @@
import numpy as np import numpy as np
import pytest import pytest
from torchy_baselines import A2C, PPO, SAC, TD3 from stable_baselines3 import A2C, PPO, SAC, TD3
from torchy_baselines.common.identity_env import IdentityEnvBox, IdentityEnv from stable_baselines3.common.identity_env import IdentityEnvBox, IdentityEnv
from torchy_baselines.common.evaluation import evaluate_policy from stable_baselines3.common.evaluation import evaluate_policy
from torchy_baselines.common.noise import NormalActionNoise from stable_baselines3.common.noise import NormalActionNoise
@pytest.mark.parametrize("model_class", [A2C, PPO]) @pytest.mark.parametrize("model_class", [A2C, PPO])

4
tests/test_logger.py
View file

@ -4,7 +4,7 @@ import shutil
import pytest import pytest
import numpy as np import numpy as np
from torchy_baselines.common.logger import (make_output_format, read_csv, read_json, DEBUG, ScopedConfigure, from stable_baselines3.common.logger import (make_output_format, read_csv, read_json, DEBUG, ScopedConfigure,
info, debug, set_level, configure, logkv, logkvs, info, debug, set_level, configure, logkv, logkvs,
dumpkvs, logkv_mean, warn, error, reset) dumpkvs, logkv_mean, warn, error, reset)
@ -18,7 +18,7 @@ KEY_VALUES = {
"g": np.array([[[1]]]), "g": np.array([[[1]]]),
} }
LOG_DIR = '/tmp/torchy_baselines/' LOG_DIR = '/tmp/stable_baselines3/'
def test_main(): def test_main():

2
tests/test_monitor.py
View file

@ -5,7 +5,7 @@ import os
import pandas import pandas
import gym import gym
from torchy_baselines.common.monitor import Monitor, get_monitor_files, load_results from stable_baselines3.common.monitor import Monitor, get_monitor_files, load_results
def test_monitor(tmp_path): def test_monitor(tmp_path):

4
tests/test_predict.py
View file

@ -1,8 +1,8 @@
import gym import gym
import pytest import pytest
from torchy_baselines import A2C, PPO, SAC, TD3 from stable_baselines3 import A2C, PPO, SAC, TD3
from torchy_baselines.common.vec_env import DummyVecEnv from stable_baselines3.common.vec_env import DummyVecEnv
MODEL_LIST = [ MODEL_LIST = [
PPO, PPO,

4
tests/test_run.py
View file

@ -1,8 +1,8 @@
import numpy as np import numpy as np
import pytest import pytest
from torchy_baselines import A2C, PPO, SAC, TD3 from stable_baselines3 import A2C, PPO, SAC, TD3
from torchy_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)) normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1))

8
tests/test_save_load.py
View file

@ -5,10 +5,10 @@ import pytest
import numpy as np import numpy as np
import torch as th import torch as th
from torchy_baselines import A2C, PPO, SAC, TD3 from stable_baselines3 import A2C, PPO, SAC, TD3
from torchy_baselines.common.identity_env import IdentityEnvBox from stable_baselines3.common.identity_env import IdentityEnvBox
from torchy_baselines.common.vec_env import DummyVecEnv from stable_baselines3.common.vec_env import DummyVecEnv
from torchy_baselines.common.identity_env import FakeImageEnv from stable_baselines3.common.identity_env import FakeImageEnv
MODEL_LIST = [ MODEL_LIST = [

2
tests/test_sde.py
View file

@ -2,7 +2,7 @@ import pytest
import torch as th import torch as th
from torch.distributions import Normal from torch.distributions import Normal
from torchy_baselines import A2C, TD3, SAC, PPO from stable_baselines3 import A2C, TD3, SAC, PPO
def test_state_dependent_exploration_grad(): def test_state_dependent_exploration_grad():

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_baselines3.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]]))

2
tests/test_vec_envs.py
View file

@ -7,7 +7,7 @@ import pytest
import gym import gym
import numpy as np import numpy as np
from torchy_baselines.common.vec_env import DummyVecEnv, SubprocVecEnv, VecNormalize, VecFrameStack from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecNormalize, VecFrameStack
N_ENVS = 3 N_ENVS = 3
VEC_ENV_CLASSES = [DummyVecEnv, SubprocVecEnv] VEC_ENV_CLASSES = [DummyVecEnv, SubprocVecEnv]

6
tests/test_vec_normalize.py
View file

@ -2,9 +2,9 @@ import gym
import pytest import pytest
import numpy as np import numpy as np
from torchy_baselines.common.running_mean_std import RunningMeanStd from stable_baselines3.common.running_mean_std import RunningMeanStd
from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize, VecFrameStack, sync_envs_normalization, unwrap_vec_normalize from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize, VecFrameStack, sync_envs_normalization, unwrap_vec_normalize
from torchy_baselines import SAC, TD3 from stable_baselines3 import SAC, TD3
ENV_ID = 'Pendulum-v0' ENV_ID = 'Pendulum-v0'

2
torchy_baselines/a2c/__init__.py
View file

@ -1,2 +0,0 @@
from torchy_baselines.a2c.a2c import A2C
from torchy_baselines.ppo.policies import MlpPolicy

2
torchy_baselines/ppo/__init__.py
View file

@ -1,2 +0,0 @@
from torchy_baselines.ppo.ppo import PPO
from torchy_baselines.ppo.policies import MlpPolicy

2
torchy_baselines/sac/__init__.py
View file

@ -1,2 +0,0 @@
from torchy_baselines.sac.sac import SAC
from torchy_baselines.sac.policies import MlpPolicy

2
torchy_baselines/td3/__init__.py
View file

@ -1,2 +0,0 @@
from torchy_baselines.td3.td3 import TD3
from torchy_baselines.td3.policies import MlpPolicy

1
torchy_baselines/version.txt
View file

@ -1 +0,0 @@
0.5.0