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

More doc + sync VecEnvs + atari

Browse source
This commit is contained in:
Antonin RAFFIN 2020-05-07 16:08:23 +02:00
parent 73afaf157c
commit 8046a24719
20 changed files with 955 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
README.md
View file

@ -98,7 +98,7 @@ Install the Stable Baselines3 package:
pip install stable-baselines3[extra] pip install stable-baselines3[extra]
``` ```
This includes an optional dependency OpenCV to display the environments when using `SubprocVecEnv`. If you do not need it, you can install without OpenCV: This includes an optional dependencies like OpenCV or `atari-py` to train on atari games. If you do not need those, you can use:
``` ```
pip install stable-baselines3 pip install stable-baselines3
``` ```
@ -144,14 +144,14 @@ model = PPO('MlpPolicy', 'CartPole-v1').learn(10000)
Please read the [documentation](https://stable-baselines.readthedocs.io/) for more examples. Please read the [documentation](https://stable-baselines.readthedocs.io/) for more examples.
<!-- ## Try it online with Colab Notebooks ! ## Try it online with Colab Notebooks !
All the following examples can be executed online using Google colab notebooks: All the following examples can be executed online using Google colab notebooks:
- [Full Tutorial](https://github.com/araffin/rl-tutorial-jnrr19) <!-- - [Full Tutorial](https://github.com/araffin/rl-tutorial-jnrr19) -->
- [All Notebooks](https://github.com/Stable-Baselines-Team/rl-colab-notebooks) - [All Notebooks](https://github.com/Stable-Baselines-Team/rl-colab-notebooks/tree/sb3)
- [Getting Started](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/stable_baselines_getting_started.ipynb) - [Getting Started](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/stable_baselines_getting_started.ipynb)
- [Training, Saving, Loading](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/saving_loading_dqn.ipynb) <!-- - [Training, Saving, Loading](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/saving_loading_dqn.ipynb)
- [Multiprocessing](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/multiprocessing_rl.ipynb) - [Multiprocessing](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/multiprocessing_rl.ipynb)
- [Monitor Training and Plotting](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/monitor_training.ipynb) - [Monitor Training and Plotting](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/monitor_training.ipynb)
- [Atari Games](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/atari_games.ipynb) - [Atari Games](https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/atari_games.ipynb)

BIN
docs/_static/img/breakout.gif vendored Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 195 KiB

7
docs/_static/img/colab.svg vendored Normal file
View file

@ -0,0 +1,7 @@
<svg width="24px" height="15px" viewBox="0 0 24 15" version="1.1" xmlns="http://www.w3.org/2000/svg">
<path d="M1.977,11.77 C-0.69,9.493 -0.628,4.691 1.977,2.413 C2.919,3.057 3.522,4.075 4.49,4.691 C3.338,6.291 3.344,7.892 4.486,9.494 C3.522,10.111 2.918,11.126 1.977,11.77 Z" fill="#FDBA18"/>
<path d="M12.257,12.114 C10.49,10.481 9.772,8.456 10.139,6.094 C10.59,3.184 12.278,1.201 15.085,0.416 C17.65,-0.302 20.049,0.199 21.963,2.235 C21.079,2.978 20.256,3.782 19.529,4.681 C18.488,3.827 17.319,3.435 16,3.856 C13.596,4.623 12.954,7.097 14.506,9.5 C14.265,9.775 14.013,10.041 13.785,10.326 C13.295,10.939 12.511,11.3 12.257,12.114 Z" fill="#FCD93D"/>
<path d="M19.529,4.682 C20.256,3.783 21.079,2.979 21.963,2.236 C24.666,5.019 24.664,9.267 21.958,12 C19.31,14.674 15.022,14.725 12.257,12.115 C12.511,11.301 13.295,10.94 13.785,10.327 C14.013,10.042 14.265,9.775 14.506,9.501 C15.559,10.417 16.76,10.769 18.106,10.331 C20.502,9.551 21.151,6.927 19.529,4.682 Z" fill="#FDBA18"/>
<path d="M4.49,4.691 C3.522,4.075 2.919,3.057 1.977,2.413 C4.186,0.015 7.698,-0.529 10.453,1.058 C11.008,1.378 11.172,1.664 10.738,2.186 C10.581,2.374 10.48,2.608 10.347,2.817 C10.048,3.287 9.838,3.884 9.418,4.188 C8.933,4.539 8.523,3.847 8.021,3.746 C6.673,3.475 5.509,3.787 4.49,4.691 Z" fill="#FCD93D"/>
<path d="M1.977,11.77 C2.918,11.126 3.522,10.111 4.486,9.493 C5.859,10.645 7.336,10.926 8.936,9.992 C9.268,9.798 9.439,9.904 9.609,10.182 C9.995,10.817 10.362,11.467 10.79,12.072 C11.13,12.552 11.012,12.787 10.537,13.078 C7.84,14.73 4.205,14.188 1.977,11.77 Z" fill="#FCD93D"/>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 1.5 KiB

BIN
docs/_static/img/try_it.png vendored Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 12 KiB

48
docs/guide/algos.rst Normal file
View file

@ -0,0 +1,48 @@
RL Algorithms
=============
This table displays the rl algorithms that are implemented in the Stable Baselines3 project,
along with some useful characteristics: support for discrete/continuous actions, multiprocessing.
============ =========== ============ ================
Name ``Box`` ``Discrete`` Multi Processing
============ =========== ============ ================
A2C ✔️ ✔️ ✔️
PPO ✔️ ✔️ ✔️
SAC ✔️ ❌ ❌
TD3 ✔️ ❌ ❌
============ =========== ============ ================
.. note::
Non-array spaces such as ``Dict`` or ``Tuple`` are not currently supported by any algorithm.
Actions ``gym.spaces``:
- ``Box``: A N-dimensional box that contains every point in the action
space.
- ``Discrete``: A list of possible actions, where each timestep only
one of the actions can be used.
- ``MultiDiscrete``: A list of possible actions, where each timestep only one action of each discrete set can be used.
- ``MultiBinary``: A list of possible actions, where each timestep any of the actions can be used in any combination.
.. note::
Some logging values (like ``ep_rew_mean``, ``ep_len_mean``) are only available when using a ``Monitor`` wrapper
See `Issue #339 <https://github.com/hill-a/stable-baselines/issues/339>`_ for more info.
Reproducibility
---------------
Completely reproducible results are not guaranteed across Tensorflow releases or different platforms.
Furthermore, results need not be reproducible between CPU and GPU executions, even when using identical seeds.
In order to make computations deterministics, on your specific problem on one specific platform,
you need to pass a ``seed`` argument at the creation of a model.
If you pass an environment to the model using ``set_env()``, then you also need to seed the environment first.
Credit: part of the *Reproducibility* section comes from `PyTorch Documentation <https://pytorch.org/docs/stable/notes/randomness.html>`_

105
docs/guide/custom_policy.rst Normal file
View file

@ -0,0 +1,105 @@
.. _custom_policy:
Custom Policy Network
---------------------
Stable Baselines3 provides policy networks for images (CnnPolicies)
and other type of input features (MlpPolicies).
One way of customising the policy network architecture is to pass arguments when creating the model,
using ``policy_kwargs`` parameter:
.. code-block:: python
import gym
import torch as th
from stable_baselines3 import PPO
# Custom MLP policy of two layers of size 32 each with tanh activation function
policy_kwargs = dict(activation_fn=th.nn.ReLU, net_arch=[32, 32])
# Create the agent
model = PPO("MlpPolicy", "CartPole-v1", policy_kwargs=policy_kwargs, verbose=1)
# Retrieve the environment
env = model.get_env()
# Train the agent
model.learn(total_timesteps=100000)
# Save the agent
model.save("ppo-cartpole")
del model
# the policy_kwargs are automatically loaded
model = PPO.load("ppo-cartpole")
You can also easily define a custom architecture for the policy (or value) network:
.. note::
Defining a custom policy class is equivalent to passing ``policy_kwargs``.
However, it lets you name the policy and so makes usually the code clearer.
``policy_kwargs`` should be rather used when doing hyperparameter search.
The ``net_arch`` parameter of ``A2C`` and ``PPO`` policies allows to specify the amount and size of the hidden layers and how many
of them are shared between the policy network and the value network. It is assumed to be a list with the following
structure:
1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer.
If the number of ints is zero, there will be no shared layers.
2. An optional dict, to specify the following non-shared layers for the value network and the policy network.
It is formatted like ``dict(vf=[<value layer sizes>], pi=[<policy layer sizes>])``.
If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed.
In short: ``[<shared layers>, dict(vf=[<non-shared value network layers>], pi=[<non-shared policy network layers>])]``.
Examples
~~~~~~~~
Two shared layers of size 128: ``net_arch=[128, 128]``
.. code-block:: none
obs
|
<128>
|
<128>
/ \
action value
Value network deeper than policy network, first layer shared: ``net_arch=[128, dict(vf=[256, 256])]``
.. code-block:: none
obs
|
<128>
/ \
action <256>
|
<256>
|
value
Initially shared then diverging: ``[128, dict(vf=[256], pi=[16])]``
.. code-block:: none
obs
|
<128>
/ \
<16> <256>
| |
action value
If your task requires even more granular control over the policy architecture, you can redefine the policy directly.
**TODO**

396
docs/guide/examples.rst Normal file
View file

@ -0,0 +1,396 @@
.. _examples:
Examples
========
Try it online with Colab Notebooks!
-----------------------------------
All the following examples can be executed online using Google colab |colab|
notebooks:
- `Full Tutorial <https://github.com/araffin/rl-tutorial-jnrr19>`_
- `All Notebooks <https://github.com/Stable-Baselines-Team/rl-colab-notebooks/tree/sb3>`_
- `Getting Started`_
.. - `Training, Saving, Loading`_
.. - `Multiprocessing`_
.. - `Monitor Training and Plotting`_
.. - `Atari Games`_
.. - `Breakout`_ (trained agent included)
.. - `Hindsight Experience Replay`_
.. - `RL Baselines zoo`_
.. _Getting Started: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/stable_baselines_getting_started.ipynb
.. _Training, Saving, Loading: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/saving_loading_dqn.ipynb
.. _Multiprocessing: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/multiprocessing_rl.ipynb
.. _Monitor Training and Plotting: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/monitor_training.ipynb
.. _Atari Games: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/atari_games.ipynb
.. _Breakout: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/breakout.ipynb
.. _Hindsight Experience Replay: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/stable_baselines_her.ipynb
.. _RL Baselines zoo: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/rl-baselines-zoo.ipynb
.. |colab| image:: ../_static/img/colab.svg
Basic Usage: Training, Saving, Loading
--------------------------------------
In the following example, we will train, save and load a DQN model on the Lunar Lander environment.
..
.. .. image:: ../_static/img/try_it.png
.. :scale: 30 %
.. :target: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/saving_loading_dqn.ipynb
.. figure:: https://cdn-images-1.medium.com/max/960/1*f4VZPKOI0PYNWiwt0la0Rg.gif
Lunar Lander Environment
.. note::
LunarLander requires the python package ``box2d``.
You can install it using ``apt install swig`` and then ``pip install box2d box2d-kengz``
.. .. note::
.. ``load`` function re-creates model from scratch on each call, which can be slow.
.. If you need to e.g. evaluate same model with multiple different sets of parameters, consider
.. using ``load_parameters`` instead.
.. code-block:: python
import gym
from stable_baselines3 import A2C
from stable_baselines3.common.evaluation import evaluate_policy
# Create environment
env = gym.make('LunarLander-v2')
# Instantiate the agent
model = A2C('MlpPolicy', env, verbose=1)
# Train the agent
model.learn(total_timesteps=int(2e5))
# Save the agent
model.save("a2c_lunar")
del model # delete trained model to demonstrate loading
# Load the trained agent
model = A2C.load("a2c_lunar")
# Evaluate the agent
mean_reward, std_reward = evaluate_policy(model, model.get_env(), n_eval_episodes=10)
# Enjoy trained agent
obs = env.reset()
for i in range(1000):
action, _states = model.predict(obs)
obs, rewards, dones, info = env.step(action)
env.render()
Multiprocessing: Unleashing the Power of Vectorized Environments
----------------------------------------------------------------
..
.. .. image:: ../_static/img/try_it.png
.. :scale: 30 %
.. :target: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/multiprocessing_rl.ipynb
.. figure:: https://cdn-images-1.medium.com/max/960/1*h4WTQNVIsvMXJTCpXm_TAw.gif
CartPole Environment
.. code-block:: python
import gym
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.ppo import MlpPolicy
from stable_baselines3.common.vec_env import SubprocVecEnv
from stable_baselines3.common.cmd_util import make_vec_env
from stable_baselines3.common.utils import set_random_seed
def make_env(env_id, rank, seed=0):
"""
Utility function for multiprocessed env.
:param env_id: (str) the environment ID
:param num_env: (int) the number of environments you wish to have in subprocesses
:param seed: (int) the inital seed for RNG
:param rank: (int) index of the subprocess
"""
def _init():
env = gym.make(env_id)
env.seed(seed + rank)
return env
set_random_seed(seed)
return _init
if __name__ == '__main__':
env_id = "CartPole-v1"
num_cpu = 4 # Number of processes to use
# Create the vectorized environment
env = SubprocVecEnv([make_env(env_id, i) for i in range(num_cpu)])
# Stable Baselines provides you with make_vec_env() helper
# which does exactly the previous steps for you:
# env = make_vec_env(env_id, n_envs=num_cpu, seed=0)
model = PPO('MlpPolicy', env, verbose=1)
model.learn(total_timesteps=25000)
obs = env.reset()
for _ in range(1000):
action, _states = model.predict(obs)
obs, rewards, dones, info = env.step(action)
env.render()
Using Callback: Monitoring Training
-----------------------------------
.. note::
We recommend reading the `Callback section <callbacks.html>`_
You can define a custom callback function that will be called inside the agent.
This could be useful when you want to monitor training, for instance display live
learning curves in Tensorboard (or in Visdom) or save the best agent.
If your callback returns False, training is aborted early.
.. .. image:: ../_static/img/try_it.png
.. :scale: 30 %
.. :target: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/monitor_training.ipynb
..
.. .. figure:: ../_static/img/learning_curve.png
..
.. Learning curve of TD3 on LunarLanderContinuous environment
.. code-block:: python
import os
import gym
import numpy as np
import matplotlib.pyplot as plt
from stable_baselines3 import TD3
from stable_baselines3.td3 import MlpPolicy
from stable_baselines3.common import results_plotter
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback
class SaveOnBestTrainingRewardCallback(BaseCallback):
"""
Callback for saving a model (the check is done every ``check_freq`` steps)
based on the training reward (in practice, we recommend using ``EvalCallback``).
:param check_freq: (int)
:param log_dir: (str) Path to the folder where the model will be saved.
It must contains the file created by the ``Monitor`` wrapper.
:param verbose: (int)
"""
def __init__(self, check_freq: int, log_dir: str, verbose=1):
super(SaveOnBestTrainingRewardCallback, self).__init__(verbose)
self.check_freq = check_freq
self.log_dir = log_dir
self.save_path = os.path.join(log_dir, 'best_model')
self.best_mean_reward = -np.inf
def _init_callback(self) -> None:
# Create folder if needed
if self.save_path is not None:
os.makedirs(self.save_path, exist_ok=True)
def _on_step(self) -> bool:
if self.n_calls % self.check_freq == 0:
# Retrieve training reward
x, y = ts2xy(load_results(self.log_dir), 'timesteps')
if len(x) > 0:
# Mean training reward over the last 100 episodes
mean_reward = np.mean(y[-100:])
if self.verbose > 0:
print("Num timesteps: {}".format(self.num_timesteps))
print("Best mean reward: {:.2f} - Last mean reward per episode: {:.2f}".format(self.best_mean_reward, mean_reward))
# New best model, you could save the agent here
if mean_reward > self.best_mean_reward:
self.best_mean_reward = mean_reward
# Example for saving best model
if self.verbose > 0:
print("Saving new best model to {}".format(self.save_path))
self.model.save(self.save_path)
return True
# Create log dir
log_dir = "tmp/"
os.makedirs(log_dir, exist_ok=True)
# Create and wrap the environment
env = gym.make('LunarLanderContinuous-v2')
env = Monitor(env, log_dir)
# Add some action noise for exploration
n_actions = env.action_space.shape[-1]
action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))
# Because we use parameter noise, we should use a MlpPolicy with layer normalization
model = TD3(MlpPolicy, env, action_noise=action_noise, verbose=0)
# Create the callback: check every 1000 steps
callback = SaveOnBestTrainingRewardCallback(check_freq=1000, log_dir=log_dir)
# Train the agent
timesteps = 1e5
model.learn(total_timesteps=int(timesteps), callback=callback)
plot_results([log_dir], timesteps, results_plotter.X_TIMESTEPS, "TD3 LunarLander")
plt.show()
Atari Games
-----------
.. figure:: ../_static/img/breakout.gif
Trained A2C agent on Breakout
.. figure:: https://cdn-images-1.medium.com/max/960/1*UHYJE7lF8IDZS_U5SsAFUQ.gif
Pong Environment
Training a RL agent on Atari games is straightforward thanks to ``make_atari_env`` helper function.
It will do `all the preprocessing <https://danieltakeshi.github.io/2016/11/25/frame-skipping-and-preprocessing-for-deep-q-networks-on-atari-2600-games/>`_
and multiprocessing for you.
.. .. image:: ../_static/img/try_it.png
.. :scale: 30 %
.. :target: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/atari_games.ipynb
..
.. code-block:: python
from stable_baselines3.common.cmd_util import make_atari_env
from stable_baselines3.common.vec_env import VecFrameStack
from stable_baselines3 import A2C
# There already exists an environment generator
# that will make and wrap atari environments correctly.
# Here we are also multi-worker training (n_envs=4 => 4 environments)
env = make_atari_env('PongNoFrameskip-v4', n_envs=4, seed=0)
# Frame-stacking with 4 frames
env = VecFrameStack(env, n_stack=4)
model = A2C('CnnPolicy', env, verbose=1)
model.learn(total_timesteps=25000)
obs = env.reset()
while True:
action, _states = model.predict(obs)
obs, rewards, dones, info = env.step(action)
env.render()
PyBullet: Normalizing input features
------------------------------------
Normalizing input features may be essential to successful training of an RL agent
(by default, images are scaled but not other types of input),
for instance when training on `PyBullet <https://github.com/bulletphysics/bullet3/>`_. For that, a wrapper exists and
will compute a running average and standard deviation of input features (it can do the same for rewards).
.. note::
you need to install pybullet with ``pip install pybullet``
.. code-block:: python
import gym
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
from stable_baselines3 import PPO
env = DummyVecEnv([lambda: gym.make("HalfCheetahBulletEnv-v0")])
# Automatically normalize the input features
env = VecNormalize(env, norm_obs=True, norm_reward=True,
clip_obs=10.)
model = PPO2('MlpPolicy', env)
model.learn(total_timesteps=2000)
# Don't forget to save the VecNormalize statistics when saving the agent
log_dir = "/tmp/"
model.save(log_dir + "ppo_reacher")
env.save(os.path.join(log_dir, "vec_normalize.pkl"))
Record a Video
--------------
Record a mp4 video (here using a random agent).
.. note::
It requires ``ffmpeg`` or ``avconv`` to be installed on the machine.
.. code-block:: python
import gym
from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv
env_id = 'CartPole-v1'
video_folder = 'logs/videos/'
video_length = 100
env = DummyVecEnv([lambda: gym.make(env_id)])
obs = env.reset()
# Record the video starting at the first step
env = VecVideoRecorder(env, video_folder,
record_video_trigger=lambda x: x == 0, video_length=video_length,
name_prefix="random-agent-{}".format(env_id))
env.reset()
for _ in range(video_length + 1):
action = [env.action_space.sample()]
obs, _, _, _ = env.step(action)
# Save the video
env.close()
Bonus: Make a GIF of a Trained Agent
------------------------------------
.. note::
For Atari games, you need to use a screen recorder such as `Kazam <https://launchpad.net/kazam>`_.
And then convert the video using `ffmpeg <https://superuser.com/questions/556029/how-do-i-convert-a-video-to-gif-using-ffmpeg-with-reasonable-quality>`_
.. code-block:: python
import imageio
import numpy as np
from stable_baselines3 import A2C
model = A2C("MlpPolicy", "LunarLander-v2").learn(100000)
images = []
obs = model.env.reset()
img = model.env.render(mode='rgb_array')
for i in range(350):
images.append(img)
action, _ = model.predict(obs)
obs, _, _ ,_ = model.env.step(action)
img = model.env.render(mode='rgb_array')
imageio.mimsave('lander_a2c.gif', [np.array(img) for i, img in enumerate(images) if i%2 == 0], fps=29)

3
docs/guide/install.rst
View file

@ -29,8 +29,7 @@ To install Stable Baselines3 with pip, execute:
pip install stable-baselines3[extra] pip install stable-baselines3[extra]
This includes an optional dependency OpenCV to display the environments when using ``SubprocVecEnv``. If you do not need it, you can install without OpenCV: This includes an optional dependencies like OpenCV or ```atari-py``` to train on atari games. If you do not need those, you can use:
.. code-block:: bash .. code-block:: bash

3
docs/index.rst
View file

@ -36,8 +36,11 @@ Main Features
guide/quickstart guide/quickstart
guide/rl_tips guide/rl_tips
guide/rl guide/rl
guide/algos
guide/examples
guide/vec_envs guide/vec_envs
guide/custom_env guide/custom_env
guide/custom_policy
guide/callbacks guide/callbacks
guide/migration guide/migration
guide/checking_nan guide/checking_nan

6
docs/misc/changelog.rst
View file

@ -3,7 +3,7 @@
Changelog Changelog
========== ==========
Pre-Release 0.6.0a1 (WIP) Pre-Release 0.6.0a3 (WIP)
------------------------------ ------------------------------
@ -15,10 +15,12 @@ New Features:
- Added env checker (Sync with Stable Baselines) - Added env checker (Sync with Stable Baselines)
- Added ``VecCheckNan`` and ``VecVideoRecorder`` (Sync with Stable Baselines) - Added ``VecCheckNan`` and ``VecVideoRecorder`` (Sync with Stable Baselines)
- Added determinism tests - Added determinism tests
- Added ``cmd_utils`` and ``atari_wrappers``
Bug Fixes: Bug Fixes:
^^^^^^^^^^ ^^^^^^^^^^
- Fixed a bug that prevented model trained on cpu to be loaded on gpu - Fixed a bug that prevented model trained on cpu to be loaded on gpu
- Fixed version number that had a new line included
Deprecations: Deprecations:
^^^^^^^^^^^^^ ^^^^^^^^^^^^^
@ -26,6 +28,8 @@ Deprecations:
Others: Others:
^^^^^^^ ^^^^^^^
- Renamed to Stable-Baseline3 - Renamed to Stable-Baseline3
- Added Dockerfile
- Sync ``VecEnvs`` with Stable-Baselines
Documentation: Documentation:
^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^

2
setup.py
View file

@ -104,6 +104,8 @@ setup(name='stable_baselines3',
'extra': [ 'extra': [
# For render # For render
'opencv-python', 'opencv-python',
# For atari games,
'gym[classic_control,atari]>=0.11'
] ]
}, },
description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.', description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',

2
stable_baselines3/__init__.py
View file

@ -8,4 +8,4 @@ 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')
with open(version_file, 'r') as file_handler: with open(version_file, 'r') as file_handler:
__version__ = file_handler.read() __version__ = file_handler.read().strip()

2
stable_baselines3/common/__init__.py
View file

@ -0,0 +1,2 @@
from stable_baselines3.common.cmd_util import make_vec_env, make_atari_env
from stable_baselines3.common.utils import set_random_seed

76
stable_baselines3/common/atari_wrappers.py Normal file
View file

@ -0,0 +1,76 @@
import gym
from gym.wrappers import AtariPreprocessing
import numpy as np
from stable_baselines3.common.type_aliases import GymStepReturn
class AtariWrapper(gym.Wrapper):
"""
Atari 2600 preprocessings
It is a wrapper around the one found in gym.
It reshapes the observation to have an additional dimension and clip the reward.
See https://github.com/openai/gym/blob/master/gym/wrappers/atari_preprocessing.py
.
This class follows the guidelines in
Machado et al. (2018), "Revisiting the Arcade Learning Environment:
Evaluation Protocols and Open Problems for General Agents".
Specifically:
* NoopReset: obtain initial state by taking random number of no-ops on reset.
* Frame skipping: 4 by default
* Max-pooling: most recent two observations
* Termination signal when a life is lost: turned off by default. Not recommended by Machado et al. (2018).
* Resize to a square image: 84x84 by default
* Grayscale observation: by default
* Scale observation: optional
:param env: (gym.Env) gym environment
env (Env): environment
:param noop_max: (int): max number of no-ops
:param frame_skip: (int): the frequency at which the agent experiences the game.
:param screen_size: (int): resize Atari frame
:param terminal_on_life_loss: (bool): if True, then step() returns done=True whenever a
life is lost.
:param grayscale_obs: (bool): if True, then gray scale observation is returned, otherwise, RGB observation
is returned.
:param scale_obs: (bool): if True, then observation normalized in range [0,1] is returned. It also limits memory
optimization benefits of FrameStack Wrapper.
"""
def __init__(self, env: gym.Env,
noop_max: int = 30,
frame_skip: int = 4,
screen_size: int = 84,
terminal_on_life_loss: bool = False,
grayscale_obs: bool = True,
scale_obs: bool = False,
clip_reward: bool = True):
env = AtariPreprocessing(env, noop_max=noop_max, frame_skip=frame_skip, screen_size=screen_size,
terminal_on_life_loss=terminal_on_life_loss, grayscale_obs=grayscale_obs,
scale_obs=scale_obs)
# Add channel dimension
if grayscale_obs:
obs_space = env.observation_space
_low, _high, _obs_dtype = (0, 255, np.uint8) if not scale_obs else (0, 1, np.float32)
env.observation_space = gym.spaces.Box(low=_low, high=_high, shape=obs_space.shape + (1,),
dtype=_obs_dtype)
super(AtariWrapper, self).__init__(env)
self.clip_reward = clip_reward
def _add_axis(self, obs: np.ndarray) -> np.ndarray:
if self.env.grayscale_obs:
return obs[..., np.newaxis]
return obs
def reset(self) -> np.ndarray:
return self._add_axis(self.env.reset())
def step(self, action: int) -> GymStepReturn:
obs, reward, done, info = self.env.step(action)
# Bin reward to {+1, 0, -1} by its sign.
if self.clip_reward:
reward = np.sign(reward)
return self._add_axis(obs), reward, done, info

112
stable_baselines3/common/cmd_util.py Normal file
View file

@ -0,0 +1,112 @@
import os
import warnings
from typing import Dict, Any, Optional, Callable, Type, Union
import gym
from stable_baselines3.common import logger
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.utils import set_random_seed
from stable_baselines3.common.atari_wrappers import AtariWrapper
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
def make_vec_env(env_id: Union[str, Type[gym.Env]],
n_envs: int = 1,
seed: Optional[int] = None,
start_index: int = 0,
monitor_dir: Optional[str] = None,
wrapper_class: Optional[Callable] = None,
env_kwargs: Optional[Dict[str, Any]] = None,
vec_env_cls: Optional[Union[DummyVecEnv, SubprocVecEnv]] = None,
vec_env_kwargs: Optional[Dict[str, Any]] = None):
"""
Create a wrapped, monitored ``VecEnv``.
By default it uses a ``DummyVecEnv`` which is usually faster
than a ``SubprocVecEnv``.
:param env_id: (str or Type[gym.Env]) the environment ID or the environment class
:param n_envs: (int) the number of environments you wish to have in parallel
:param seed: (int) the initial seed for the random number generator
:param start_index: (int) start rank index
:param monitor_dir: (str) Path to a folder where the monitor files will be saved.
If None, no file will be written, however, the env will still be wrapped
in a Monitor wrapper to provide additional information about training.
:param wrapper_class: (gym.Wrapper or callable) Additional wrapper to use on the environment.
This can also be a function with single argument that wraps the environment in many things.
:param env_kwargs: (dict) Optional keyword argument to pass to the env constructor
:param vec_env_cls: (Type[VecEnv]) A custom ``VecEnv`` class constructor. Default: None.
:param vec_env_kwargs: (dict) Keyword arguments to pass to the ``VecEnv`` class constructor.
:return: (VecEnv) The wrapped environment
"""
env_kwargs = {} if env_kwargs is None else env_kwargs
vec_env_kwargs = {} if vec_env_kwargs is None else vec_env_kwargs
def make_env(rank):
def _init():
if isinstance(env_id, str):
env = gym.make(env_id)
if len(env_kwargs) > 0:
warnings.warn("No environment class was passed (only an env ID) so ``env_kwargs`` will be ignored")
else:
env = env_id(**env_kwargs)
if seed is not None:
env.seed(seed + rank)
env.action_space.seed(seed + rank)
# Wrap the env in a Monitor wrapper
# to have additional training information
monitor_path = os.path.join(monitor_dir, str(rank)) if monitor_dir is not None else None
# Create the monitor folder if needed
if monitor_path is not None:
os.makedirs(monitor_dir, exist_ok=True)
env = Monitor(env, filename=monitor_path)
# Optionally, wrap the environment with the provided wrapper
if wrapper_class is not None:
env = wrapper_class(env)
return env
return _init
# No custom VecEnv is passed
if vec_env_cls is None:
# Default: use a DummyVecEnv
vec_env_cls = DummyVecEnv
return vec_env_cls([make_env(i + start_index) for i in range(n_envs)], **vec_env_kwargs)
def make_atari_env(env_id: Union[str, Type[gym.Env]],
n_envs: int = 1,
seed: Optional[int] = None,
start_index: int = 0,
monitor_dir: Optional[str] = None,
wrapper_kwargs: Optional[Dict[str, Any]] = None,
env_kwargs: Optional[Dict[str, Any]] = None,
vec_env_cls: Optional[Union[DummyVecEnv, SubprocVecEnv]] = None,
vec_env_kwargs: Optional[Dict[str, Any]] = None):
"""
Create a wrapped, monitored VecEnv for Atari.
It is a wrapper around ``make_vec_env`` that includes common preprocessing for Atari games.
:param env_id: (str or Type[gym.Env]) the environment ID or the environment class
:param n_envs: (int) the number of environments you wish to have in parallel
:param seed: (int) the initial seed for the random number generator
:param start_index: (int) start rank index
:param monitor_dir: (str) Path to a folder where the monitor files will be saved.
If None, no file will be written, however, the env will still be wrapped
in a Monitor wrapper to provide additional information about training.
:param wrapper_kwargs: (Dict[str, Any]) Optional keyword argument to pass to the env constructor
:param env_kwargs: (Dict[str, Any]) Optional keyword argument to pass to the env constructor
:param vec_env_cls: (Type[VecEnv]) A custom `VecEnv` class constructor. Default: None.
:param vec_env_kwargs: (Dict[str, Any]) Keyword arguments to pass to the `VecEnv` class constructor.
:return: (VecEnv) The wrapped environment
"""
if wrapper_kwargs is None:
wrapper_kwargs = {}
def atari_wrapper(env: gym.Env) -> gym.Env:
env = AtariWrapper(env, **wrapper_kwargs)
return env
return make_vec_env(env_id, n_envs=n_envs, seed=seed, start_index=start_index,
monitor_dir=monitor_dir, wrapper_class=atari_wrapper,
env_kwargs=env_kwargs, vec_env_cls=vec_env_cls, vec_env_kwargs=vec_env_kwargs)

74
stable_baselines3/common/vec_env/base_vec_env.py
View file

@ -1,9 +1,39 @@
import inspect import inspect
import pickle import pickle
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Sequence, Optional, List, Union
import numpy as np
import cloudpickle import cloudpickle
from stable_baselines3.common import logger
def tile_images(img_nhwc: Sequence[np.ndarray]) -> np.ndarray:
"""
Tile N images into one big PxQ image
(P,Q) are chosen to be as close as possible, and if N
is square, then P=Q.
:param img_nhwc: (Sequence[np.ndarray]) list or array of images, ndim=4 once turned into array. img nhwc
n = batch index, h = height, w = width, c = channel
:return: (np.ndarray) img_HWc, ndim=3
"""
img_nhwc = np.asarray(img_nhwc)
n_images, height, width, n_channels = img_nhwc.shape
# new_height was named H before
new_height = int(np.ceil(np.sqrt(n_images)))
# new_width was named W before
new_width = int(np.ceil(float(n_images) / new_height))
img_nhwc = np.array(list(img_nhwc) + [img_nhwc[0] * 0 for _ in range(n_images, new_height * new_width)])
# img_HWhwc
out_image = img_nhwc.reshape((new_height, new_width, height, width, n_channels))
# img_HhWwc
out_image = out_image.transpose(0, 2, 1, 3, 4)
# img_Hh_Ww_c
out_image = out_image.reshape((new_height * height, new_width * width, n_channels))
return out_image
class AlreadySteppingError(Exception): class AlreadySteppingError(Exception):
""" """
@ -132,31 +162,46 @@ class VecEnv(ABC):
self.step_async(actions) self.step_async(actions)
return self.step_wait() return self.step_wait()
def get_images(self): def get_images(self, *args, **kwargs) -> Sequence[np.ndarray]:
""" """
Return RGB images from each environment Return RGB images from each environment
""" """
raise NotImplementedError raise NotImplementedError
def render(self, *args, **kwargs): def render(self, *args, mode: str = 'human', **kwargs):
""" """
Gym environment rendering Gym environment rendering
:param mode: (str) the rendering type :param mode: the rendering type
""" """
raise NotImplementedError() try:
imgs = self.get_images(*args, **kwargs)
except NotImplementedError:
logger.warn('Render not defined for {}'.format(self))
return
def seed(self, seed, indices=None): # Create a big image by tiling images from subprocesses
bigimg = tile_images(imgs)
if mode == 'human':
import cv2 # pytype:disable=import-error
cv2.imshow('vecenv', bigimg[:, :, ::-1])
cv2.waitKey(1)
elif mode == 'rgb_array':
return bigimg
else:
raise NotImplementedError
@abstractmethod
def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
""" """
:param seed: (int or [int]) Sets the random seeds for all environments, based on a given seed.
:param indices: ([int]) Each individual environment will still get its own seed, by incrementing the given seed.
:param seed: (Optional[int]) The random seed. May be None for completely random seeding.
:return: (List[Union[None, int]]) Returns a list containing the seeds for each individual env.
Note that all list elements may be None, if the env does not return anything when being seeded.
""" """
indices = self._get_indices(indices) pass
# Different seed per environment
if not hasattr(seed, 'len'):
seed = [seed + i for i in range(len(indices))]
assert len(seed) == len(indices)
return [self.env_method('seed', seed[i], indices=i) for i in indices]
@property @property
def unwrapped(self): def unwrapped(self):
@ -217,6 +262,9 @@ class VecEnvWrapper(VecEnv):
def step_wait(self): def step_wait(self):
pass pass
def seed(self, seed=None):
return self.venv.seed(seed)
def close(self): def close(self):
return self.venv.close() return self.venv.close()

24
stable_baselines3/common/vec_env/dummy_vec_env.py
View file

@ -1,5 +1,6 @@
from collections import OrderedDict from collections import OrderedDict
from copy import deepcopy from copy import deepcopy
from typing import Sequence
import numpy as np import numpy as np
@ -10,7 +11,7 @@ from stable_baselines3.common.vec_env.util import copy_obs_dict, dict_to_obs, ob
class DummyVecEnv(VecEnv): class DummyVecEnv(VecEnv):
""" """
Creates a simple vectorized wrapper for multiple environments, calling each environment in sequence on the current Creates a simple vectorized wrapper for multiple environments, calling each environment in sequence on the current
Python process. This is useful for computationally simple environment such as ``cartpole-v1``, as the overhead of Python process. This is useful for computationally simple environment such as ````cartpole-v1````, as the overhead of
multiprocess or multithread outweighs the environment computation time. This can also be used for RL methods that multiprocess or multithread outweighs the environment computation time. This can also be used for RL methods that
require a vectorized environment, but that you want a single environments to train with. require a vectorized environment, but that you want a single environments to train with.
@ -48,6 +49,12 @@ class DummyVecEnv(VecEnv):
return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones), return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones),
deepcopy(self.buf_infos)) deepcopy(self.buf_infos))
def seed(self, seed=None):
seeds = list()
for idx, env in enumerate(self.envs):
seeds.append(env.seed(seed + idx))
return seeds
def reset(self): def reset(self):
for env_idx in range(self.num_envs): for env_idx in range(self.num_envs):
obs = self.envs[env_idx].reset() obs = self.envs[env_idx].reset()
@ -58,10 +65,21 @@ class DummyVecEnv(VecEnv):
for env in self.envs: for env in self.envs:
env.close() env.close()
def get_images(self): def get_images(self, *args, **kwargs) -> Sequence[np.ndarray]:
return [env.render(mode='rgb_array') for env in self.envs] return [env.render(*args, mode='rgb_array', **kwargs) for env in self.envs]
def render(self, *args, **kwargs): def render(self, *args, **kwargs):
"""
Gym environment rendering. If there are multiple environments then
they are tiled together in one image via ``BaseVecEnv.render()``.
Otherwise (if ``self.num_envs == 1``), we pass the render call directly to the
underlying environment.
Therefore, some arguments such as ``mode`` will have values that are valid
only when ``num_envs == 1``.
:param mode: The rendering type.
"""
if self.num_envs == 1: if self.num_envs == 1:
return self.envs[0].render(*args, **kwargs) return self.envs[0].render(*args, **kwargs)
else: else:

49
stable_baselines3/common/vec_env/subproc_vec_env.py
View file

@ -1,5 +1,6 @@
import multiprocessing import multiprocessing
from collections import OrderedDict from collections import OrderedDict
from typing import Sequence
import gym import gym
import numpy as np import numpy as np
@ -43,32 +44,6 @@ def _worker(remote, parent_remote, env_fn_wrapper):
break break
def tile_images(img_nhwc):
"""
Tile N images into one big PxQ image
(P,Q) are chosen to be as close as possible, and if N
is square, then P=Q.
:param img_nhwc: (list) list or array of images, ndim=4 once turned into array. img nhwc
n = batch index, h = height, w = width, c = channel
:return: (numpy float) img_HWc, ndim=3
"""
img_nhwc = np.asarray(img_nhwc)
n_images, height, width, n_channels = img_nhwc.shape
# new_height was named H before
new_height = int(np.ceil(np.sqrt(n_images)))
# new_width was named W before
new_width = int(np.ceil(float(n_images) / new_height))
img_nhwc = np.array(list(img_nhwc) + [img_nhwc[0] * 0 for _ in range(n_images, new_height * new_width)])
# img_HWhwc
out_image = img_nhwc.reshape((new_height, new_width, height, width, n_channels))
# img_HhWwc
out_image = out_image.transpose(0, 2, 1, 3, 4)
# img_Hh_Ww_c
out_image = out_image.reshape((new_height * height, new_width * width, n_channels))
return out_image
class SubprocVecEnv(VecEnv): class SubprocVecEnv(VecEnv):
""" """
Creates a multiprocess vectorized wrapper for multiple environments, distributing each environment to its own Creates a multiprocess vectorized wrapper for multiple environments, distributing each environment to its own
@ -131,6 +106,11 @@ class SubprocVecEnv(VecEnv):
obs, rews, dones, infos = zip(*results) obs, rews, dones, infos = zip(*results)
return _flatten_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos return _flatten_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos
def seed(self, seed=None):
for idx, remote in enumerate(self.remotes):
remote.send(('seed', seed + idx))
return [remote.recv() for remote in self.remotes]
def reset(self): def reset(self):
for remote in self.remotes: for remote in self.remotes:
remote.send(('reset', None)) remote.send(('reset', None))
@ -149,27 +129,12 @@ class SubprocVecEnv(VecEnv):
process.join() process.join()
self.closed = True self.closed = True
def render(self, mode='human', *args, **kwargs): def get_images(self, *args, **kwargs) -> Sequence[np.ndarray]:
for pipe in self.remotes: for pipe in self.remotes:
# gather images from subprocesses # gather images from subprocesses
# `mode` will be taken into account later # `mode` will be taken into account later
pipe.send(('render', (args, {'mode': 'rgb_array', **kwargs}))) pipe.send(('render', (args, {'mode': 'rgb_array', **kwargs})))
imgs = [pipe.recv() for pipe in self.remotes] imgs = [pipe.recv() for pipe in self.remotes]
# Create a big image by tiling images from subprocesses
bigimg = tile_images(imgs)
if mode == 'human':
import cv2
cv2.imshow('vecenv', bigimg[:, :, ::-1])
cv2.waitKey(1)
elif mode == 'rgb_array':
return bigimg
else:
raise NotImplementedError
def get_images(self):
for pipe in self.remotes:
pipe.send(('render', {"mode": 'rgb_array'}))
imgs = [pipe.recv() for pipe in self.remotes]
return imgs return imgs
def get_attr(self, attr_name, indices=None): def get_attr(self, attr_name, indices=None):

2
stable_baselines3/version.txt
View file

@ -1 +1 @@
0.6.0a1 0.6.0a3

101
tests/test_utils.py Normal file
View file

@ -0,0 +1,101 @@
import os
import shutil
import pytest
import gym
from stable_baselines3 import A2C
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.cmd_util import make_vec_env, make_atari_env
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
@pytest.mark.parametrize("env_id", ['CartPole-v1', lambda: gym.make('CartPole-v1')])
@pytest.mark.parametrize("n_envs", [1, 2])
@pytest.mark.parametrize("vec_env_cls", [None, SubprocVecEnv])
@pytest.mark.parametrize("wrapper_class", [None, gym.wrappers.TimeLimit])
def test_make_vec_env(env_id, n_envs, vec_env_cls, wrapper_class):
env = make_vec_env(env_id, n_envs, vec_env_cls=vec_env_cls,
wrapper_class=wrapper_class, monitor_dir=None, seed=0)
assert env.num_envs == n_envs
if vec_env_cls is None:
assert isinstance(env, DummyVecEnv)
if wrapper_class is not None:
assert isinstance(env.envs[0], wrapper_class)
else:
assert isinstance(env.envs[0], Monitor)
else:
assert isinstance(env, SubprocVecEnv)
# Kill subprocesses
env.close()
@pytest.mark.parametrize("env_id", ['BreakoutNoFrameskip-v4'])
@pytest.mark.parametrize("n_envs", [1, 2])
@pytest.mark.parametrize("wrapper_kwargs", [None, dict(clip_reward=False, screen_size=60)])
def test_make_vec_env(env_id, n_envs, wrapper_kwargs):
env_id = 'BreakoutNoFrameskip-v4'
env = make_atari_env(env_id, n_envs,
wrapper_kwargs=wrapper_kwargs, monitor_dir=None, seed=0)
assert env.num_envs == n_envs
wrapped_atari_env = env.envs[0]
if wrapper_kwargs is not None:
# Wrapped into DummyVecEnv + Monitor
assert wrapped_atari_env.observation_space.shape == (60, 60, 1)
assert wrapped_atari_env.clip_reward == False
else:
assert wrapped_atari_env.observation_space.shape == (84, 84, 1)
assert wrapped_atari_env.clip_reward == True
def test_custom_vec_env(tmp_path):
"""
Stand alone test for a special case (passing a custom VecEnv class) to avoid doubling the number of tests.
"""
monitor_dir = tmp_path / 'test_make_vec_env/'
env = make_vec_env('CartPole-v1', n_envs=1,
monitor_dir=monitor_dir, seed=0,
vec_env_cls=SubprocVecEnv, vec_env_kwargs={'start_method': None})
assert env.num_envs == 1
assert isinstance(env, SubprocVecEnv)
assert os.path.isdir(monitor_dir)
# Kill subprocess
env.close()
# Cleanup folder
shutil.rmtree(monitor_dir)
# This should fail because DummyVecEnv does not have any keyword argument
with pytest.raises(TypeError):
make_vec_env('CartPole-v1', n_envs=1, vec_env_kwargs={'dummy': False})
def test_evaluate_policy():
model = A2C('MlpPolicy', 'Pendulum-v0', seed=0)
n_steps_per_episode, n_eval_episodes = 200, 2
model.n_callback_calls = 0
def dummy_callback(locals_, _globals):
locals_['model'].n_callback_calls += 1
_, episode_lengths = evaluate_policy(model, model.get_env(), n_eval_episodes, deterministic=True,
render=False, callback=dummy_callback, reward_threshold=None,
return_episode_rewards=True)
n_steps = sum(episode_lengths)
assert n_steps == n_steps_per_episode * n_eval_episodes
assert n_steps == model.n_callback_calls
# Reaching a mean reward of zero is impossible with the Pendulum env
with pytest.raises(AssertionError):
evaluate_policy(model, model.get_env(), n_eval_episodes, reward_threshold=0.0)
episode_rewards, _ = evaluate_policy(model, model.get_env(), n_eval_episodes, return_episode_rewards=True)
assert len(episode_rewards) == n_eval_episodes