saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-30 20:18:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge branch 'master' into feat/crr

Browse source
This commit is contained in:
Antonin Raffin 2020-10-13 12:05:15 +02:00
commit b1a737e055
38 changed files with 599 additions and 195 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
README.md
View file

@ -50,9 +50,9 @@ Planned features:
- [ ] TRPO
## Migration guide
## Migration guide: from Stable-Baselines (SB2) to Stable-Baselines3 (SB3)
**TODO: migration guide from Stable-Baselines in the documentation**
A migration guide from SB2 to SB3 can be found in the [documentation](https://stable-baselines3.readthedocs.io/en/master/guide/migration.html).
## Documentation

34
docs/guide/custom_policy.rst
View file

@ -258,9 +258,31 @@ If your task requires even more granular control over the policy/value architect
.. TODO (see https://github.com/DLR-RM/stable-baselines3/issues/113)
.. Off-Policy Algorithms
.. ^^^^^^^^^^^^^^^^^^^^^
..
.. If you need a network architecture that is different for the actor and the critic when using ``SAC``, ``DDPG`` or ``TD3``,
.. you can easily redefine the actor class for instance.
Off-Policy Algorithms
^^^^^^^^^^^^^^^^^^^^^
If you need a network architecture that is different for the actor and the critic when using ``SAC``, ``DDPG`` or ``TD3``,
you can pass a dictionary of the following structure: ``dict(qf=[<critic network architecture>], pi=[<actor network architecture>])``.
For example, if you want a different architecture for the actor (aka ``pi``) and the critic (Q-function aka ``qf``) networks,
then you can specify ``net_arch=dict(qf=[400, 300], pi=[64, 64])``.
Otherwise, to have actor and critic that share the same network architecture,
you only need to specify ``net_arch=[256, 256]`` (here, two hidden layers of 256 units each).
.. note::
Compared to their on-policy counterparts, no shared layers (other than the feature extractor)
between the actor and the critic are allowed (to prevent issues with target networks).
.. code-block:: python
from stable_baselines3 import SAC
# Custom actor architecture with two layers of 64 units each
# Custom critic architecture with two layers of 400 and 300 units
policy_kwargs = dict(net_arch=dict(pi=[64, 64], qf=[400, 300]))
# Create the agent
model = SAC("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs, verbose=1)
model.learn(5000)

2
docs/guide/install.rst
View file

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

190
docs/guide/migration.rst
View file

@ -5,8 +5,194 @@ Migrating from Stable-Baselines
================================
This is a guide to migrate from Stable-Baselines to Stable-Baselines3.
This is a guide to migrate from Stable-Baselines (SB2) to Stable-Baselines3 (SB3).
It also references the main changes.
**TODO**
Overview
========
Overall Stable-Baselines3 (SB3) keeps the high-level API of Stable-Baselines (SB2).
Most of the changes are to ensure more consistency and are internal ones.
Because of the backend change, from Tensorflow to PyTorch, the internal code is much much readable and easy to debug
at the cost of some speed (dynamic graph vs static graph., see `Issue #90 <https://github.com/DLR-RM/stable-baselines3/issues/90>`_)
However, the algorithms were extensively benchmarked on Atari games and continuous control PyBullet envs
(see `Issue #48 <https://github.com/DLR-RM/stable-baselines3/issues/48>`_ and `Issue #49 <https://github.com/DLR-RM/stable-baselines3/issues/49>`_)
so you should not expect performance drop when switching from SB2 to SB3.
How to migrate?
===============
In most cases, replacing ``from stable_baselines`` by ``from stable_baselines3`` will be sufficient.
Some files were moved to the common folder (cf below) and could result to import errors.
Some algorithms were removed because of their complexity to improve the maintainability of the project.
We recommend reading this guide carefully to understand all the changes that were made.
You can also take a look at the `rl-zoo3 <https://github.com/DLR-RM/rl-baselines3-zoo>`_ and compare the imports
to the `rl-zoo <https://github.com/araffin/rl-baselines-zoo>`_ of SB2 to have a concrete example of successful migration.
Breaking Changes
================
- SB3 requires python 3.6+ (instead of python 3.5+ for SB2)
- Dropped MPI support
- Dropped layer normalized policies (e.g. ``LnMlpPolicy``)
- Dropped parameter noise for DDPG and DQN
- PPO is now closer to the original implementation (no clipping of the value function by default), cf PPO section below
- Orthogonal initialization is only used by A2C/PPO
- The features extractor (CNN extractor) is shared between policy and q-networks for DDPG/SAC/TD3 and only the policy loss used to update it (much faster)
- Tensorboard legacy logging was dropped in favor of having one logger for the terminal and Tensorboard (cf :ref:`Tensorboard integration <tensorboard>`)
- We dropped ACKTR/ACER support because of their complexity compared to simpler alternatives (PPO, SAC, TD3) performing as good.
- We dropped GAIL support as we are focusing on model-free RL only, you can however take a look at the `Imitation Learning Baseline Implementations <https://github.com/HumanCompatibleAI/imitation>`_
which are based on SB3.
- ``action_probability`` is currently not implemented in the base class
You can take a look at the `issue about SB3 implementation design <https://github.com/hill-a/stable-baselines/issues/576>`_ for more details.
Moved Files
-----------
- ``bench/monitor.py`` -> ``common/monitor.py``
- ``logger.py`` -> ``common/logger.py``
- ``results_plotter.py`` -> ``common/results_plotter.py``
Utility functions are no longer exported from ``common`` module, you should import them with their absolute path, e.g.:
.. code-block:: python
from stable_baselines3.common.cmd_util import make_atari_env, make_vec_env
from stable_baselines3.common.utils import set_random_seed
instead of ``from stable_baselines3.common import make_atari_env``
Changes and renaming in parameters
----------------------------------
Base-class (all algorithms)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
- ``load_parameters`` -> ``set_parameters``
- ``get/set_parameters`` return a dictionary mapping object names
to their respective PyTorch tensors and other objects representing
their parameters, instead of simpler mapping of parameter name to
a NumPy array. These functions also return PyTorch tensors rather
than NumPy arrays.
Policies
^^^^^^^^
- ``cnn_extractor`` -> ``feature_extractor``, as ``feature_extractor`` in now used with ``MlpPolicy`` too
A2C
^^^
- ``epsilon`` -> ``rms_prop_eps``
- ``lr_schedule`` is part of ``learning_rate`` (it can be a callable).
- ``alpha``, ``momentum`` are modifiable through ``policy_kwargs`` key ``optimizer_kwargs``.
.. warning::
PyTorch implementation of RMSprop `differs from Tensorflow's <https://github.com/pytorch/pytorch/issues/23796>`_,
which leads to `different and potentially more unstable results <https://github.com/DLR-RM/stable-baselines3/pull/110#issuecomment-663255241>`_.
Use ``stable_baselines3.common.sb2_compat.rmsprop_tf_like.RMSpropTFLike`` optimizer to match the results
with Tensorflow's implementation. This can be done through ``policy_kwargs``: ``A2C(policy_kwargs=dict(optimizer_class=RMSpropTFLike))``
PPO
^^^
- ``cliprange`` -> ``clip_range``
- ``cliprange_vf`` -> ``clip_range_vf``
- ``nminibatches`` -> ``batch_size``
.. warning::
``nminibatches`` gave different batch size depending on the number of environments: ``batch_size = (n_steps * n_envs) // nminibatches``
- ``clip_range_vf`` behavior for PPO is slightly different: Set it to ``None`` (default) to deactivate clipping (in SB2, you had to pass ``-1``, ``None`` meant to use ``clip_range`` for the clipping)
- ``lam`` -> ``gae_lambda``
- ``noptepochs`` -> ``n_epochs``
PPO default hyperparameters are the one tuned for continuous control environment.
We recommend taking a look at the :ref:`RL Zoo <rl_zoo>` for hyperparameters tuned for Atari games.
DQN
^^^
Only the vanilla DQN is implemented right now but extensions will follow.
Default hyperparameters are taken from the nature paper, except for the optimizer and learning rate that were taken from Stable Baselines defaults.
DDPG
^^^^
DDPG now follows the same interface as SAC/TD3.
For state/reward normalization, you should use ``VecNormalize`` as for all other algorithms.
SAC/TD3
^^^^^^^
SAC/TD3 now accept any number of critics, e.g. ``policy_kwargs=dict(n_critics=3)``, instead of only two before.
.. note::
SAC/TD3 default hyperparameters (including network architecture) now match the ones from the original papers.
DDPG is using TD3 defaults.
SAC
^^^
SAC implementation matches the latest version of the original implementation: it uses two Q function networks and two target Q function networks
instead of two Q function networks and one Value function network (SB2 implementation, first version of the original implementation).
Despite this change, no change in performance should be expected.
.. note::
SAC ``predict()`` method has now ``deterministic=False`` by default for consistency.
To match SB2 behavior, you need to explicitly pass ``deterministic=True``
New logger API
--------------
- Methods were renamed in the logger:
- ``logkv`` -> ``record``, ``writekvs`` -> ``write``, ``writeseq`` -> ``write_sequence``,
- ``logkvs`` -> ``record_dict``, ``dumpkvs`` -> ``dump``,
- ``getkvs`` -> ``get_log_dict``, ``logkv_mean`` -> ``record_mean``,
Internal Changes
----------------
Please read the :ref:`Developer Guide <developer>` section.
New Features (SB3 vs SB2)
=========================
- Much cleaner and consistent base code (and no more warnings =D!) and static type checks
- Independent saving/loading/predict for policies
- A2C now supports Generalized Advantage Estimation (GAE) and advantage normalization (both are deactivated by default)
- Generalized State-Dependent Exploration (gSDE) exploration is available for A2C/PPO/SAC. It allows to use RL directly on real robots (cf https://arxiv.org/abs/2005.05719)
- Proper evaluation (using separate env) is included in the base class (using ``EvalCallback``),
if you pass the environment as a string, you can pass ``create_eval_env=True`` to the algorithm constructor.
- Better saving/loading: optimizers are now included in the saved parameters and there is two new methods ``save_replay_buffer`` and ``load_replay_buffer`` for the replay buffer when using off-policy algorithms (DQN/DDPG/SAC/TD3)
- You can pass ``optimizer_class`` and ``optimizer_kwargs`` to ``policy_kwargs`` in order to easily
customize optimizers
- Seeding now works properly to have deterministic results
- Replay buffer does not grow, allocate everything at build time (faster)
- We added a memory efficient replay buffer variant (pass ``optimize_memory_usage=True`` to the constructor), it reduces drastically the memory used especially when using images
- You can specify an arbitrary number of critics for SAC/TD3 (e.g. ``policy_kwargs=dict(n_critics=3)``)

40
docs/misc/changelog.rst
View file

@ -3,14 +3,43 @@
Changelog
==========
Pre-Release 0.9.0a2 (WIP)
Pre-Release 0.10.0a0 (WIP)
------------------------------
Breaking Changes:
^^^^^^^^^^^^^^^^^
New Features:
^^^^^^^^^^^^^
- Allow custom actor/critic network architectures using ``net_arch=dict(qf=[400, 300], pi=[64, 64])`` for off-policy algorithms (SAC, TD3, DDPG)
Bug Fixes:
^^^^^^^^^^
- Fix GAE computation for on-policy algorithms (off-by one for the last value) (thanks @Wovchena)
Deprecations:
^^^^^^^^^^^^^
Others:
^^^^^^^
- Improved typing coverage
- Improved error messages for unsupported spaces
Documentation:
^^^^^^^^^^^^^^
- Added first draft of migration guide
Pre-Release 0.9.0 (2020-10-03)
------------------------------
**Bug fixes, get/set parameters and improved docs**
Breaking Changes:
^^^^^^^^^^^^^^^^^
- Removed ``device`` keyword argument of policies; use ``policy.to(device)`` instead. (@qxcv)
- Rename ``BaseClass.get_torch_variables`` -> ``BaseClass._get_torch_save_params`` and
``BaseClass.excluded_save_params`` -> ``BaseClass._excluded_save_params``
- Rename ``BaseClass.get_torch_variables`` -> ``BaseClass._get_torch_save_params`` and ``BaseClass.excluded_save_params`` -> ``BaseClass._excluded_save_params``
- Renamed saved items ``tensors`` to ``pytorch_variables`` for clarity
- ``make_atari_env``, ``make_vec_env`` and ``set_random_seed`` must be imported with (and not directly from ``stable_baselines3.common``):
@ -46,10 +75,11 @@ Others:
- Removed ``AlreadySteppingError`` and ``NotSteppingError`` that were not used
- Fixed typos in SAC and TD3
- Reorganized functions for clarity in ``BaseClass`` (save/load functions close to each other, private
functions at top)
functions at top)
- Clarified docstrings on what is saved and loaded to/from files
- Simplified ``save_to_zip_file`` function by removing duplicate code
- Store library version along with the saved models
- DQN loss is now logged
Documentation:
^^^^^^^^^^^^^^
@ -72,7 +102,7 @@ Breaking Changes:
- Refactored ``Critic`` class for ``TD3`` and ``SAC``, it is now called ``ContinuousCritic``
and has an additional parameter ``n_critics``
- ``SAC`` and ``TD3`` now accept an arbitrary number of critics (e.g. ``policy_kwargs=dict(n_critics=3)``)
instead of only 2 previously
instead of only 2 previously
New Features:
^^^^^^^^^^^^^

4
stable_baselines3/common/atari_wrappers.py
View file

@ -9,7 +9,7 @@ try:
except ImportError:
cv2 = None
from stable_baselines3.common.type_aliases import GymStepReturn
from stable_baselines3.common.type_aliases import GymObs, GymStepReturn
class NoopResetEnv(gym.Wrapper):
@ -146,7 +146,7 @@ class MaxAndSkipEnv(gym.Wrapper):
return max_frame, total_reward, done, info
def reset(self, **kwargs):
def reset(self, **kwargs) -> GymObs:
return self.env.reset(**kwargs)

4
stable_baselines3/common/base_class.py
View file

@ -478,7 +478,7 @@ class BaseAlgorithm(ABC):
load_path_or_dict: Union[str, Dict[str, Dict]],
exact_match: bool = True,
device: Union[th.device, str] = "auto",
):
) -> None:
"""
Load parameters from a given zip-file or a nested dictionary containing parameters for
different modules (see ``get_parameters``).
@ -610,7 +610,7 @@ class BaseAlgorithm(ABC):
model.policy.reset_noise() # pytype: disable=attribute-error
return model
def get_parameters(self):
def get_parameters(self) -> Dict[str, Dict]:
"""
Return the parameters of the agent. This includes parameters from different networks, e.g.
critics (value functions) and policies (pi functions).

20
stable_baselines3/common/buffers.py
View file

@ -1,4 +1,5 @@
import warnings
from abc import ABC, abstractmethod
from typing import Generator, Optional, Union
import numpy as np
@ -16,7 +17,7 @@ from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBu
from stable_baselines3.common.vec_env import VecNormalize
class BaseBuffer(object):
class BaseBuffer(ABC):
"""
Base class that represent a buffer (rollout or replay)
@ -102,7 +103,10 @@ class BaseBuffer(object):
batch_inds = np.random.randint(0, upper_bound, size=batch_size)
return self._get_samples(batch_inds, env=env)
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None):
@abstractmethod
def _get_samples(
self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None
) -> Union[ReplayBufferSamples, RolloutBufferSamples]:
"""
:param batch_inds:
:param env:
@ -295,7 +299,7 @@ class RolloutBuffer(BaseBuffer):
self.generator_ready = False
super(RolloutBuffer, self).reset()
def compute_returns_and_advantage(self, last_value: th.Tensor, dones: np.ndarray) -> None:
def compute_returns_and_advantage(self, last_values: th.Tensor, dones: np.ndarray) -> None:
"""
Post-processing step: compute the returns (sum of discounted rewards)
and GAE advantage.
@ -306,22 +310,22 @@ class RolloutBuffer(BaseBuffer):
where R is the discounted reward with value bootstrap,
set ``gae_lambda=1.0`` during initialization.
:param last_value:
:param last_values:
:param dones:
"""
# convert to numpy
last_value = last_value.clone().cpu().numpy().flatten()
last_values = last_values.clone().cpu().numpy().flatten()
last_gae_lam = 0
for step in reversed(range(self.buffer_size)):
if step == self.buffer_size - 1:
next_non_terminal = 1.0 - dones
next_value = last_value
next_values = last_values
else:
next_non_terminal = 1.0 - self.dones[step + 1]
next_value = self.values[step + 1]
delta = self.rewards[step] + self.gamma * next_value * next_non_terminal - self.values[step]
next_values = self.values[step + 1]
delta = self.rewards[step] + self.gamma * next_values * next_non_terminal - self.values[step]
last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam
self.advantages[step] = last_gae_lam
self.returns = self.advantages + self.values

11
stable_baselines3/common/callbacks.py
View file

@ -1,7 +1,7 @@
import os
import warnings
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Union
from typing import Any, Callable, Dict, List, Optional, Union
import gym
import numpy as np
@ -217,9 +217,10 @@ class CheckpointCallback(BaseCallback):
:param save_freq:
:param save_path: Path to the folder where the model will be saved.
:param name_prefix: Common prefix to the saved models
:param verbose:
"""
def __init__(self, save_freq: int, save_path: str, name_prefix="rl_model", verbose=0):
def __init__(self, save_freq: int, save_path: str, name_prefix: str = "rl_model", verbose: int = 0):
super(CheckpointCallback, self).__init__(verbose)
self.save_freq = save_freq
self.save_path = save_path
@ -247,7 +248,7 @@ class ConvertCallback(BaseCallback):
:param verbose:
"""
def __init__(self, callback, verbose=0):
def __init__(self, callback: Callable, verbose: int = 0):
super(ConvertCallback, self).__init__(verbose)
self.callback = callback
@ -314,7 +315,7 @@ class EvalCallback(EventCallback):
self.evaluations_timesteps = []
self.evaluations_length = []
def _init_callback(self):
def _init_callback(self) -> None:
# Does not work in some corner cases, where the wrapper is not the same
if not isinstance(self.training_env, type(self.eval_env)):
warnings.warn("Training and eval env are not of the same type" f"{self.training_env} != {self.eval_env}")
@ -450,7 +451,7 @@ class StopTrainingOnMaxEpisodes(BaseCallback):
self._total_max_episodes = max_episodes
self.n_episodes = 0
def _init_callback(self):
def _init_callback(self) -> None:
# At start set total max according to number of envirnments
self._total_max_episodes = self.max_episodes * self.training_env.num_envs

6
stable_baselines3/common/cmd_util.py
View file

@ -6,7 +6,7 @@ import gym
from stable_baselines3.common.atari_wrappers import AtariWrapper
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecEnv
def make_vec_env(
@ -19,7 +19,7 @@ def make_vec_env(
env_kwargs: Optional[Dict[str, Any]] = None,
vec_env_cls: Optional[Type[Union[DummyVecEnv, SubprocVecEnv]]] = None,
vec_env_kwargs: Optional[Dict[str, Any]] = None,
):
) -> VecEnv:
"""
Create a wrapped, monitored ``VecEnv``.
By default it uses a ``DummyVecEnv`` which is usually faster
@ -85,7 +85,7 @@ def make_atari_env(
env_kwargs: Optional[Dict[str, Any]] = None,
vec_env_cls: Optional[Union[DummyVecEnv, SubprocVecEnv]] = None,
vec_env_kwargs: Optional[Dict[str, Any]] = None,
):
) -> VecEnv:
"""
Create a wrapped, monitored VecEnv for Atari.
It is a wrapper around ``make_vec_env`` that includes common preprocessing for Atari games.

4
stable_baselines3/common/distributions.py
View file

@ -1,7 +1,7 @@
"""Probability distributions."""
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple
from typing import Any, Dict, List, Optional, Tuple, Union
import gym
import torch as th
@ -19,7 +19,7 @@ class Distribution(ABC):
super(Distribution, self).__init__()
@abstractmethod
def proba_distribution_net(self, *args, **kwargs):
def proba_distribution_net(self, *args, **kwargs) -> Union[nn.Module, Tuple[nn.Module, nn.Parameter]]:
"""Create the layers and parameters that represent the distribution.
Subclasses must define this, but the arguments and return type vary between

2
stable_baselines3/common/logger.py
View file

@ -50,7 +50,7 @@ class SeqWriter(object):
sequence writer
"""
def write_sequence(self, sequence: List):
def write_sequence(self, sequence: List) -> None:
"""
write_sequence an array to file

10
stable_baselines3/common/monitor.py
View file

@ -5,12 +5,14 @@ import json
import os
import time
from glob import glob
from typing import Any, Dict, List, Optional, Tuple
from typing import List, Optional, Tuple, Union
import gym
import numpy as np
import pandas
from stable_baselines3.common.type_aliases import GymObs, GymStepReturn
class Monitor(gym.Wrapper):
"""
@ -62,7 +64,7 @@ class Monitor(gym.Wrapper):
self.total_steps = 0
self.current_reset_info = {} # extra info about the current episode, that was passed in during reset()
def reset(self, **kwargs) -> np.ndarray:
def reset(self, **kwargs) -> GymObs:
"""
Calls the Gym environment reset. Can only be called if the environment is over, or if allow_early_resets is True
@ -83,7 +85,7 @@ class Monitor(gym.Wrapper):
self.current_reset_info[key] = value
return self.env.reset(**kwargs)
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, Dict[Any, Any]]:
def step(self, action: Union[np.ndarray, int]) -> GymStepReturn:
"""
Step the environment with the given action
@ -112,7 +114,7 @@ class Monitor(gym.Wrapper):
self.total_steps += 1
return observation, reward, done, info
def close(self):
def close(self) -> None:
"""
Closes the environment
"""

2
stable_baselines3/common/noise.py
View file

@ -139,7 +139,7 @@ class VectorizedActionNoise(ActionNoise):
return self._base_noise
@base_noise.setter
def base_noise(self, base_noise: ActionNoise):
def base_noise(self, base_noise: ActionNoise) -> None:
if base_noise is None:
raise ValueError("Expected base_noise to be an instance of ActionNoise, not None", ActionNoise)
if not isinstance(base_noise, ActionNoise):

7
stable_baselines3/common/on_policy_algorithm.py
View file

@ -179,7 +179,12 @@ class OnPolicyAlgorithm(BaseAlgorithm):
self._last_obs = new_obs
self._last_dones = dones
rollout_buffer.compute_returns_and_advantage(values, dones=dones)
with th.no_grad():
# Compute value for the last timestep
obs_tensor = th.as_tensor(new_obs).to(self.device)
_, values, _ = self.policy.forward(obs_tensor)
rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)
callback.on_rollout_end()

2
stable_baselines3/common/policies.py
View file

@ -145,7 +145,7 @@ class BaseModel(nn.Module, ABC):
model.to(device)
return model
def load_from_vector(self, vector: np.ndarray):
def load_from_vector(self, vector: np.ndarray) -> None:
"""
Load parameters from a 1D vector.

6
stable_baselines3/common/preprocessing.py
View file

@ -78,7 +78,7 @@ def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space, normalize_im
return obs.float()
else:
raise NotImplementedError()
raise NotImplementedError(f"Preprocessing not implemented for {observation_space}")
def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]:
@ -100,7 +100,7 @@ def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]:
# Number of binary features
return (int(observation_space.n),)
else:
raise NotImplementedError()
raise NotImplementedError(f"{observation_space} observation space is not supported")
def get_flattened_obs_dim(observation_space: spaces.Space) -> int:
@ -139,4 +139,4 @@ def get_action_dim(action_space: spaces.Space) -> int:
# Number of binary actions
return int(action_space.n)
else:
raise NotImplementedError()
raise NotImplementedError(f"{action_space} action space is not supported")

17
stable_baselines3/common/save_util.py
View file

@ -176,7 +176,7 @@ def json_to_data(json_string: str, custom_objects: Optional[Dict[str, Any]] = No
@functools.singledispatch
def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verbose=0, suffix=None):
def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verbose: int = 0, suffix: Optional[str] = None):
"""
Opens a path for reading or writing with a preferred suffix and raises debug information.
If the provided path is a derivative of io.BufferedIOBase it ensures that the file
@ -197,6 +197,7 @@ def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verb
:param suffix: The preferred suffix. If mode is "w" then the opened file has the suffix.
If mode is "r" then we attempt to open the path. If an error is raised and the suffix
is not None, we attempt to open the path with the suffix.
:return:
"""
if not isinstance(path, io.BufferedIOBase):
raise TypeError("Path parameter has invalid type.", io.BufferedIOBase)
@ -214,7 +215,7 @@ def open_path(path: Union[str, pathlib.Path, io.BufferedIOBase], mode: str, verb
@open_path.register(str)
def open_path_str(path: str, mode: str, verbose=0, suffix=None) -> io.BufferedIOBase:
def open_path_str(path: str, mode: str, verbose: int = 0, suffix: Optional[str] = None) -> io.BufferedIOBase:
"""
Open a path given by a string. If writing to the path, the function ensures
that the path exists.
@ -226,12 +227,13 @@ def open_path_str(path: str, mode: str, verbose=0, suffix=None) -> io.BufferedIO
:param suffix: The preferred suffix. If mode is "w" then the opened file has the suffix.
If mode is "r" then we attempt to open the path. If an error is raised and the suffix
is not None, we attempt to open the path with the suffix.
:return:
"""
return open_path(pathlib.Path(path), mode, verbose, suffix)
@open_path.register(pathlib.Path)
def open_path_pathlib(path: pathlib.Path, mode: str, verbose=0, suffix=None) -> io.BufferedIOBase:
def open_path_pathlib(path: pathlib.Path, mode: str, verbose: int = 0, suffix: Optional[str] = None) -> io.BufferedIOBase:
"""
Open a path given by a string. If writing to the path, the function ensures
that the path exists.
@ -244,6 +246,7 @@ def open_path_pathlib(path: pathlib.Path, mode: str, verbose=0, suffix=None) ->
:param suffix: The preferred suffix. If mode is "w" then the opened file has the suffix.
If mode is "r" then we attempt to open the path. If an error is raised and the suffix
is not None, we attempt to open the path with the suffix.
:return:
"""
if mode not in ("w", "r"):
raise ValueError("Expected mode to be either 'w' or 'r'.")
@ -286,7 +289,7 @@ def save_to_zip_file(
data: Dict[str, Any] = None,
params: Dict[str, Any] = None,
pytorch_variables: Dict[str, Any] = None,
verbose=0,
verbose: int = 0,
) -> None:
"""
Save model data to a zip archive.
@ -321,7 +324,7 @@ def save_to_zip_file(
archive.writestr("_stable_baselines3_version", stable_baselines3.__version__)
def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj, verbose=0) -> None:
def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj: Any, verbose: int = 0) -> None:
"""
Save an object to path creating the necessary folders along the way.
If the path exists and is a directory, it will raise a warning and rename the path.
@ -337,7 +340,7 @@ def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj, verbose=
pickle.dump(obj, file_handler)
def load_from_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], verbose=0) -> Any:
def load_from_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], verbose: int = 0) -> Any:
"""
Load an object from the path. If a suffix is provided in the path, it will use that suffix.
If the path does not exist, it will attempt to load using the .pkl suffix.
@ -355,7 +358,7 @@ def load_from_zip_file(
load_path: Union[str, pathlib.Path, io.BufferedIOBase],
load_data: bool = True,
device: Union[th.device, str] = "auto",
verbose=0,
verbose: int = 0,
) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]):
"""
Load model data from a .zip archive

44
stable_baselines3/common/sb2_compat/rmsprop_tf_like.py
View file

@ -1,3 +1,5 @@
from typing import Any, Callable, Dict, Iterable, Optional
import torch
from torch.optim import Optimizer
@ -28,21 +30,29 @@ class RMSpropTFLike(Optimizer):
is the scheduled learning rate and :math:`v` is the weighted moving average
of the squared gradient.
Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-2)
momentum (float, optional): momentum factor (default: 0)
alpha (float, optional): smoothing constant (default: 0.99)
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
centered (bool, optional) : if ``True``, compute the centered RMSProp,
the gradient is normalized by an estimation of its variance
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
:params: iterable of parameters to optimize or dicts defining
parameter groups
:param lr: learning rate (default: 1e-2)
:param momentum: momentum factor (default: 0)
:param alpha: smoothing constant (default: 0.99)
:param eps: term added to the denominator to improve
numerical stability (default: 1e-8)
:param centered: if ``True``, compute the centered RMSProp,
the gradient is normalized by an estimation of its variance
:param weight_decay: weight decay (L2 penalty) (default: 0)
"""
def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False):
def __init__(
self,
params: Iterable[torch.nn.Parameter],
lr: float = 1e-2,
alpha: float = 0.99,
eps: float = 1e-8,
weight_decay: float = 0,
momentum: float = 0,
centered: bool = False,
):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
@ -57,19 +67,19 @@ class RMSpropTFLike(Optimizer):
defaults = dict(lr=lr, momentum=momentum, alpha=alpha, eps=eps, centered=centered, weight_decay=weight_decay)
super(RMSpropTFLike, self).__init__(params, defaults)
def __setstate__(self, state):
def __setstate__(self, state: Dict[str, Any]) -> None:
super(RMSpropTFLike, self).__setstate__(state)
for group in self.param_groups:
group.setdefault("momentum", 0)
group.setdefault("centered", False)
@torch.no_grad()
def step(self, closure=None):
def step(self, closure: Optional[Callable] = None) -> Optional[torch.Tensor]:
"""Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
:param closure: A closure that reevaluates the model
and returns the loss.
:return: loss
"""
loss = None
if closure is not None:

40
stable_baselines3/common/torch_layers.py
View file

@ -219,3 +219,43 @@ class MlpExtractor(nn.Module):
"""
shared_latent = self.shared_net(features)
return self.policy_net(shared_latent), self.value_net(shared_latent)
def get_actor_critic_arch(net_arch: Union[List[int], Dict[str, List[int]]]) -> Tuple[List[int], List[int]]:
"""
Get the actor and critic network architectures for off-policy actor-critic algorithms (SAC, TD3, DDPG).
The ``net_arch`` parameter allows to specify the amount and size of the hidden layers,
which can be different for the actor and the critic.
It is assumed to be a list of ints or a dict.
1. If it is a list, actor and critic networks will have the same architecture.
The architecture is represented by a list of integers (of arbitrary length (zero allowed))
each specifying the number of units per layer.
If the number of ints is zero, the network will be linear.
2. If it is a dict, it should have the following structure:
``dict(qf=[<critic network architecture>], pi=[<actor network architecture>])``.
where the network architecture is a list as described in 1.
For example, to have actor and critic that share the same network architecture,
you only need to specify ``net_arch=[256, 256]`` (here, two hidden layers of 256 units each).
If you want a different architecture for the actor and the critic,
then you can specify ``net_arch=dict(qf=[400, 300], pi=[64, 64])``.
.. note::
Compared to their on-policy counterparts, no shared layers (other than the feature extractor)
between the actor and the critic are allowed (to prevent issues with target networks).
:param net_arch: The specification of the actor and critic networks.
See above for details on its formatting.
:return: The network architectures for the actor and the critic
"""
if isinstance(net_arch, list):
actor_arch, critic_arch = net_arch, net_arch
else:
assert isinstance(net_arch, dict), "Error: the net_arch can only contain be a list of ints or a dict"
assert "pi" in net_arch, "Error: no key 'pi' was provided in net_arch for the actor network"
assert "qf" in net_arch, "Error: no key 'qf' was provided in net_arch for the critic network"
actor_arch, critic_arch = net_arch["pi"], net_arch["qf"]
return actor_arch, critic_arch

24
stable_baselines3/common/utils.py
View file

@ -2,6 +2,7 @@ import glob
import os
import random
from collections import deque
from itertools import zip_longest
from typing import Callable, Iterable, Optional, Union
import gym
@ -191,7 +192,7 @@ def configure_logger(
logger.configure(format_strings=[""])
def check_for_correct_spaces(env: GymEnv, observation_space: gym.spaces.Space, action_space: gym.spaces.Space):
def check_for_correct_spaces(env: GymEnv, observation_space: gym.spaces.Space, action_space: gym.spaces.Space) -> None:
"""
Checks that the environment has same spaces as provided ones. Used by BaseAlgorithm to check if
spaces match after loading the model with given env.
@ -286,6 +287,24 @@ def safe_mean(arr: Union[np.ndarray, list, deque]) -> np.ndarray:
return np.nan if len(arr) == 0 else np.mean(arr)
def zip_strict(*iterables: Iterable) -> Iterable:
r"""
``zip()`` function but enforces that iterables are of equal length.
Raises ``ValueError`` if iterables not of equal length.
Code inspired by Stackoverflow answer for question #32954486.
:param \*iterables: iterables to ``zip()``
"""
# As in Stackoverflow #32954486, use
# new object for "empty" in case we have
# Nones in iterable.
sentinel = object()
for combo in zip_longest(*iterables, fillvalue=sentinel):
if sentinel in combo:
raise ValueError("Iterables have different lengths")
yield combo
def polyak_update(params: Iterable[th.nn.Parameter], target_params: Iterable[th.nn.Parameter], tau: float) -> None:
"""
Perform a Polyak average update on ``target_params`` using ``params``:
@ -303,6 +322,7 @@ def polyak_update(params: Iterable[th.nn.Parameter], target_params: Iterable[th.
:param tau: the soft update coefficient ("Polyak update", between 0 and 1)
"""
with th.no_grad():
for param, target_param in zip(params, target_params):
# zip does not raise an exception if length of parameters does not match.
for param, target_param in zip_strict(params, target_params):
target_param.data.mul_(1 - tau)
th.add(target_param.data, param.data, alpha=tau, out=target_param.data)

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

@ -13,7 +13,7 @@ from stable_baselines3.common import logger
VecEnvIndices = Union[None, int, Iterable[int]]
# VecEnvObs is what is returned by the reset() method
# it contains the observation for each env
VecEnvObs = Union[np.ndarray, Dict[str, Any]]
VecEnvObs = Union[np.ndarray, Dict[str, np.ndarray], Tuple[np.ndarray, ...]]
# VecEnvStepReturn is what is returned by the step() method
# it contains the observation, reward, done, info for each env
VecEnvStepReturn = Tuple[VecEnvObs, np.ndarray, np.ndarray, List[Dict]]
@ -76,7 +76,7 @@ class VecEnv(ABC):
raise NotImplementedError()
@abstractmethod
def step_async(self, actions: np.ndarray):
def step_async(self, actions: np.ndarray) -> None:
"""
Tell all the environments to start taking a step
with the given actions.
@ -104,7 +104,7 @@ class VecEnv(ABC):
raise NotImplementedError()
@abstractmethod
def get_attr(self, attr_name: str, indices: "VecEnvIndices" = None) -> List[Any]:
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
"""
Return attribute from vectorized environment.
@ -115,7 +115,7 @@ class VecEnv(ABC):
raise NotImplementedError()
@abstractmethod
def set_attr(self, attr_name: str, value: Any, indices: "VecEnvIndices" = None) -> None:
def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
"""
Set attribute inside vectorized environments.
@ -127,7 +127,7 @@ class VecEnv(ABC):
raise NotImplementedError()
@abstractmethod
def env_method(self, method_name: str, *method_args, indices: "VecEnvIndices" = None, **method_kwargs) -> List[Any]:
def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
"""
Call instance methods of vectorized environments.
@ -210,7 +210,7 @@ class VecEnv(ABC):
else:
return None
def _get_indices(self, indices: "VecEnvIndices") -> Iterable[int]:
def _get_indices(self, indices: VecEnvIndices) -> Iterable[int]:
"""
Convert a flexibly-typed reference to environment indices to an implied list of indices.
@ -248,7 +248,7 @@ class VecEnvWrapper(VecEnv):
)
self.class_attributes = dict(inspect.getmembers(self.__class__))
def step_async(self, actions: np.ndarray):
def step_async(self, actions: np.ndarray) -> None:
self.venv.step_async(actions)
@abstractmethod
@ -259,7 +259,7 @@ class VecEnvWrapper(VecEnv):
def step_wait(self) -> VecEnvStepReturn:
pass
def seed(self, seed: Optional[int] = None):
def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
return self.venv.seed(seed)
def close(self) -> None:
@ -271,13 +271,13 @@ class VecEnvWrapper(VecEnv):
def get_images(self) -> Sequence[np.ndarray]:
return self.venv.get_images()
def get_attr(self, attr_name, indices=None):
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
return self.venv.get_attr(attr_name, indices)
def set_attr(self, attr_name, value, indices=None):
def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
return self.venv.set_attr(attr_name, value, indices)
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
return self.venv.env_method(method_name, *method_args, indices=indices, **method_kwargs)
def __getattr__(self, name: str) -> Any:
@ -305,7 +305,7 @@ class VecEnvWrapper(VecEnv):
all_attributes.update(self.class_attributes)
return all_attributes
def getattr_recursive(self, name: str):
def getattr_recursive(self, name: str) -> Any:
"""Recursively check wrappers to find attribute.
:param name: name of attribute to look for
@ -323,7 +323,7 @@ class VecEnvWrapper(VecEnv):
return attr
def getattr_depth_check(self, name: str, already_found: bool):
def getattr_depth_check(self, name: str, already_found: bool) -> str:
"""See base class.
:return: name of module whose attribute is being shadowed, if any.

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

@ -1,11 +1,11 @@
from collections import OrderedDict
from copy import deepcopy
from typing import Callable, List, Optional, Sequence
from typing import Any, Callable, List, Optional, Sequence, Union
import gym
import numpy as np
from stable_baselines3.common.vec_env.base_vec_env import VecEnv
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvIndices, VecEnvObs, VecEnvStepReturn
from stable_baselines3.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info
@ -35,10 +35,10 @@ class DummyVecEnv(VecEnv):
self.actions = None
self.metadata = env.metadata
def step_async(self, actions: np.ndarray):
def step_async(self, actions: np.ndarray) -> None:
self.actions = actions
def step_wait(self):
def step_wait(self) -> VecEnvStepReturn:
for env_idx in range(self.num_envs):
obs, self.buf_rews[env_idx], self.buf_dones[env_idx], self.buf_infos[env_idx] = self.envs[env_idx].step(
self.actions[env_idx]
@ -50,26 +50,26 @@ class DummyVecEnv(VecEnv):
self._save_obs(env_idx, obs)
return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones), deepcopy(self.buf_infos))
def seed(self, seed: Optional[int] = None) -> List[int]:
def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
seeds = list()
for idx, env in enumerate(self.envs):
seeds.append(env.seed(seed + idx))
return seeds
def reset(self):
def reset(self) -> VecEnvObs:
for env_idx in range(self.num_envs):
obs = self.envs[env_idx].reset()
self._save_obs(env_idx, obs)
return self._obs_from_buf()
def close(self):
def close(self) -> None:
for env in self.envs:
env.close()
def get_images(self) -> Sequence[np.ndarray]:
return [env.render(mode="rgb_array") for env in self.envs]
def render(self, mode: str = "human"):
def render(self, mode: str = "human") -> Optional[np.ndarray]:
"""
Gym environment rendering. If there are multiple environments then
they are tiled together in one image via ``BaseVecEnv.render()``.
@ -86,32 +86,32 @@ class DummyVecEnv(VecEnv):
else:
return super().render(mode=mode)
def _save_obs(self, env_idx, obs):
def _save_obs(self, env_idx: int, obs: VecEnvObs) -> None:
for key in self.keys:
if key is None:
self.buf_obs[key][env_idx] = obs
else:
self.buf_obs[key][env_idx] = obs[key]
def _obs_from_buf(self):
def _obs_from_buf(self) -> VecEnvObs:
return dict_to_obs(self.observation_space, copy_obs_dict(self.buf_obs))
def get_attr(self, attr_name, indices=None):
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
"""Return attribute from vectorized environment (see base class)."""
target_envs = self._get_target_envs(indices)
return [getattr(env_i, attr_name) for env_i in target_envs]
def set_attr(self, attr_name, value, indices=None):
def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
"""Set attribute inside vectorized environments (see base class)."""
target_envs = self._get_target_envs(indices)
for env_i in target_envs:
setattr(env_i, attr_name, value)
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
"""Call instance methods of vectorized environments."""
target_envs = self._get_target_envs(indices)
return [getattr(env_i, method_name)(*method_args, **method_kwargs) for env_i in target_envs]
def _get_target_envs(self, indices):
def _get_target_envs(self, indices: VecEnvIndices) -> List[gym.Env]:
indices = self._get_indices(indices)
return [self.envs[i] for i in indices]

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

@ -1,14 +1,22 @@
import multiprocessing
import multiprocessing as mp
from collections import OrderedDict
from typing import Sequence
from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
import gym
import numpy as np
from stable_baselines3.common.vec_env.base_vec_env import CloudpickleWrapper, VecEnv
from stable_baselines3.common.vec_env.base_vec_env import (
CloudpickleWrapper,
VecEnv,
VecEnvIndices,
VecEnvObs,
VecEnvStepReturn,
)
def _worker(remote, parent_remote, env_fn_wrapper):
def _worker(
remote: mp.connection.Connection, parent_remote: mp.connection.Connection, env_fn_wrapper: CloudpickleWrapper
) -> None:
parent_remote.close()
env = env_fn_wrapper.var()
while True:
@ -71,7 +79,7 @@ class SubprocVecEnv(VecEnv):
Defaults to 'forkserver' on available platforms, and 'spawn' otherwise.
"""
def __init__(self, env_fns, start_method=None):
def __init__(self, env_fns: List[Callable[[], gym.Env]], start_method: Optional[str] = None):
self.waiting = False
self.closed = False
n_envs = len(env_fns)
@ -80,9 +88,9 @@ class SubprocVecEnv(VecEnv):
# Fork is not a thread safe method (see issue #217)
# but is more user friendly (does not require to wrap the code in
# a `if __name__ == "__main__":`)
forkserver_available = "forkserver" in multiprocessing.get_all_start_methods()
forkserver_available = "forkserver" in mp.get_all_start_methods()
start_method = "forkserver" if forkserver_available else "spawn"
ctx = multiprocessing.get_context(start_method)
ctx = mp.get_context(start_method)
self.remotes, self.work_remotes = zip(*[ctx.Pipe() for _ in range(n_envs)])
self.processes = []
@ -98,29 +106,29 @@ class SubprocVecEnv(VecEnv):
observation_space, action_space = self.remotes[0].recv()
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
def step_async(self, actions):
def step_async(self, actions: np.ndarray) -> None:
for remote, action in zip(self.remotes, actions):
remote.send(("step", action))
self.waiting = True
def step_wait(self):
def step_wait(self) -> VecEnvStepReturn:
results = [remote.recv() for remote in self.remotes]
self.waiting = False
obs, rews, dones, infos = zip(*results)
return _flatten_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos
def seed(self, seed=None):
def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
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) -> VecEnvObs:
for remote in self.remotes:
remote.send(("reset", None))
obs = [remote.recv() for remote in self.remotes]
return _flatten_obs(obs, self.observation_space)
def close(self):
def close(self) -> None:
if self.closed:
return
if self.waiting:
@ -140,14 +148,14 @@ class SubprocVecEnv(VecEnv):
imgs = [pipe.recv() for pipe in self.remotes]
return imgs
def get_attr(self, attr_name, indices=None):
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
"""Return attribute from vectorized environment (see base class)."""
target_remotes = self._get_target_remotes(indices)
for remote in target_remotes:
remote.send(("get_attr", attr_name))
return [remote.recv() for remote in target_remotes]
def set_attr(self, attr_name, value, indices=None):
def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
"""Set attribute inside vectorized environments (see base class)."""
target_remotes = self._get_target_remotes(indices)
for remote in target_remotes:
@ -155,14 +163,14 @@ class SubprocVecEnv(VecEnv):
for remote in target_remotes:
remote.recv()
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
"""Call instance methods of vectorized environments."""
target_remotes = self._get_target_remotes(indices)
for remote in target_remotes:
remote.send(("env_method", (method_name, method_args, method_kwargs)))
return [remote.recv() for remote in target_remotes]
def _get_target_remotes(self, indices):
def _get_target_remotes(self, indices: VecEnvIndices) -> List[Any]:
"""
Get the connection object needed to communicate with the wanted
envs that are in subprocesses.
@ -174,14 +182,14 @@ class SubprocVecEnv(VecEnv):
return [self.remotes[i] for i in indices]
def _flatten_obs(obs, space):
def _flatten_obs(obs: Union[List[VecEnvObs], Tuple[VecEnvObs]], space: gym.spaces.Space) -> VecEnvObs:
"""
Flatten observations, depending on the observation space.
:param obs: (list<X> or tuple<X> where X is dict<ndarray>, tuple<ndarray> or ndarray) observations.
:param obs: observations.
A list or tuple of observations, one per environment.
Each environment observation may be a NumPy array, or a dict or tuple of NumPy arrays.
:return (OrderedDict<ndarray>, tuple<ndarray> or ndarray) flattened observations.
:return: flattened observations.
A flattened NumPy array or an OrderedDict or tuple of flattened numpy arrays.
Each NumPy array has the environment index as its first axis.
"""

8
stable_baselines3/common/vec_env/util.py
View file

@ -2,11 +2,13 @@
Helpers for dealing with vectorized environments.
"""
from collections import OrderedDict
from typing import Any, Dict, List, Tuple, Union
from typing import Any, Dict, List, Tuple
import gym
import numpy as np
from stable_baselines3.common.vec_env.base_vec_env import VecEnvObs
def copy_obs_dict(obs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
"""
@ -19,9 +21,7 @@ def copy_obs_dict(obs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
return OrderedDict([(k, np.copy(v)) for k, v in obs.items()])
def dict_to_obs(
space: gym.spaces.Space, obs_dict: Dict[Any, np.ndarray]
) -> Union[Dict[Any, np.ndarray], Tuple[np.ndarray, ...], np.ndarray]:
def dict_to_obs(space: gym.spaces.Space, obs_dict: Dict[Any, np.ndarray]) -> VecEnvObs:
"""
Convert an internal representation raw_obs into the appropriate type
specified by space.

18
stable_baselines3/common/vec_env/vec_check_nan.py
View file

@ -2,7 +2,7 @@ import warnings
import numpy as np
from stable_baselines3.common.vec_env.base_vec_env import VecEnvWrapper
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, VecEnvWrapper
class VecCheckNan(VecEnvWrapper):
@ -16,7 +16,7 @@ class VecCheckNan(VecEnvWrapper):
:param check_inf: Whether or not to check for +inf or -inf as well
"""
def __init__(self, venv, raise_exception=False, warn_once=True, check_inf=True):
def __init__(self, venv: VecEnv, raise_exception: bool = False, warn_once: bool = True, check_inf: bool = True):
VecEnvWrapper.__init__(self, venv)
self.raise_exception = raise_exception
self.warn_once = warn_once
@ -25,13 +25,13 @@ class VecCheckNan(VecEnvWrapper):
self._observations = None
self._user_warned = False
def step_async(self, actions):
def step_async(self, actions: np.ndarray) -> None:
self._check_val(async_step=True, actions=actions)
self._actions = actions
self.venv.step_async(actions)
def step_wait(self):
def step_wait(self) -> VecEnvStepReturn:
observations, rewards, news, infos = self.venv.step_wait()
self._check_val(async_step=False, observations=observations, rewards=rewards, news=news)
@ -39,7 +39,7 @@ class VecCheckNan(VecEnvWrapper):
self._observations = observations
return observations, rewards, news, infos
def reset(self):
def reset(self) -> VecEnvObs:
observations = self.venv.reset()
self._actions = None
@ -48,7 +48,7 @@ class VecCheckNan(VecEnvWrapper):
self._observations = observations
return observations
def _check_val(self, *, async_step, **kwargs):
def _check_val(self, *, async_step: bool, **kwargs) -> None:
# 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
@ -66,7 +66,7 @@ class VecCheckNan(VecEnvWrapper):
self._user_warned = True
msg = ""
for i, (name, type_val) in enumerate(found):
msg += "found {} in {}".format(type_val, name)
msg += f"found {type_val} in {name}"
if i != len(found) - 1:
msg += ", "
@ -76,9 +76,9 @@ class VecCheckNan(VecEnvWrapper):
if self._actions is None:
msg += "environment observation (at reset)"
else:
msg += "environment, Last given value was: \r\n\taction={}".format(self._actions)
msg += f"environment, Last given value was: \r\n\taction={self._actions}"
else:
msg += "RL model, Last given value was: \r\n\tobservations={}".format(self._observations)
msg += f"RL model, Last given value was: \r\n\tobservations={self._observations}"
if self.raise_exception:
raise ValueError(msg)

37
stable_baselines3/common/vec_env/vec_normalize.py
View file

@ -1,9 +1,10 @@
import pickle
from typing import Any, Dict
import numpy as np
from stable_baselines3.common.running_mean_std import RunningMeanStd
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvWrapper
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper
class VecNormalize(VecEnvWrapper):
@ -22,7 +23,15 @@ class VecNormalize(VecEnvWrapper):
"""
def __init__(
self, venv, training=True, norm_obs=True, norm_reward=True, clip_obs=10.0, clip_reward=10.0, gamma=0.99, epsilon=1e-8
self,
venv: VecEnv,
training: bool = True,
norm_obs: bool = True,
norm_reward: bool = True,
clip_obs: float = 10.0,
clip_reward: float = 10.0,
gamma: float = 0.99,
epsilon: float = 1e-8,
):
VecEnvWrapper.__init__(self, venv)
self.obs_rms = RunningMeanStd(shape=self.observation_space.shape)
@ -39,7 +48,7 @@ class VecNormalize(VecEnvWrapper):
self.old_obs = np.array([])
self.old_reward = np.array([])
def __getstate__(self):
def __getstate__(self) -> Dict[str, Any]:
"""
Gets state for pickling.
@ -52,7 +61,7 @@ class VecNormalize(VecEnvWrapper):
del state["ret"]
return state
def __setstate__(self, state):
def __setstate__(self, state: Dict[str, Any]) -> None:
"""
Restores pickled state.
@ -63,7 +72,7 @@ class VecNormalize(VecEnvWrapper):
assert "venv" not in state
self.venv = None
def set_venv(self, venv):
def set_venv(self, venv: VecEnv) -> None:
"""
Sets the vector environment to wrap to venv.
@ -78,7 +87,7 @@ class VecNormalize(VecEnvWrapper):
raise ValueError("venv is incompatible with current statistics.")
self.ret = np.zeros(self.num_envs)
def step_wait(self):
def step_wait(self) -> VecEnvStepReturn:
"""
Apply sequence of actions to sequence of environments
actions -> (observations, rewards, news)
@ -100,12 +109,12 @@ class VecNormalize(VecEnvWrapper):
self.ret[news] = 0
return obs, rews, news, infos
def _update_reward(self, reward):
def _update_reward(self, reward: np.ndarray) -> None:
"""Update reward normalization statistics."""
self.ret = self.ret * self.gamma + reward
self.ret_rms.update(self.ret)
def normalize_obs(self, obs):
def normalize_obs(self, obs: np.ndarray) -> np.ndarray:
"""
Normalize observations using this VecNormalize's observations statistics.
Calling this method does not update statistics.
@ -114,7 +123,7 @@ class VecNormalize(VecEnvWrapper):
obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs)
return obs
def normalize_reward(self, reward):
def normalize_reward(self, reward: np.ndarray) -> np.ndarray:
"""
Normalize rewards using this VecNormalize's rewards statistics.
Calling this method does not update statistics.
@ -123,30 +132,30 @@ class VecNormalize(VecEnvWrapper):
reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
return reward
def unnormalize_obs(self, obs):
def unnormalize_obs(self, obs: np.ndarray) -> np.ndarray:
if self.norm_obs:
return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean
return obs
def unnormalize_reward(self, reward):
def unnormalize_reward(self, reward: np.ndarray) -> np.ndarray:
if self.norm_reward:
return reward * np.sqrt(self.ret_rms.var + self.epsilon)
return reward
def get_original_obs(self):
def get_original_obs(self) -> np.ndarray:
"""
Returns an unnormalized version of the observations from the most recent
step or reset.
"""
return self.old_obs.copy()
def get_original_reward(self):
def get_original_reward(self) -> np.ndarray:
"""
Returns an unnormalized version of the rewards from the most recent step.
"""
return self.old_reward.copy()
def reset(self):
def reset(self) -> np.ndarray:
"""
Reset all environments
"""

28
stable_baselines3/common/vec_env/vec_video_recorder.py
View file

@ -1,13 +1,12 @@
import os
from typing import Callable
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.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, 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):
@ -24,7 +23,14 @@ class VecVideoRecorder(VecEnvWrapper):
:param name_prefix: Prefix to the video name
"""
def __init__(self, venv, video_folder, record_video_trigger, video_length=200, name_prefix="rl-video"):
def __init__(
self,
venv: VecEnv,
video_folder: str,
record_video_trigger: Callable[[int], bool],
video_length: int = 200,
name_prefix: str = "rl-video",
):
VecEnvWrapper.__init__(self, venv)
@ -34,7 +40,7 @@ class VecVideoRecorder(VecEnvWrapper):
# Unwrap to retrieve metadata dict
# that will be used by gym recorder
while isinstance(temp_env, VecNormalize) or isinstance(temp_env, VecFrameStack):
while isinstance(temp_env, VecEnvWrapper):
temp_env = temp_env.venv
if isinstance(temp_env, DummyVecEnv) or isinstance(temp_env, SubprocVecEnv):
@ -58,12 +64,12 @@ class VecVideoRecorder(VecEnvWrapper):
self.recording = False
self.recorded_frames = 0
def reset(self):
def reset(self) -> VecEnvObs:
obs = self.venv.reset()
self.start_video_recorder()
return obs
def start_video_recorder(self):
def start_video_recorder(self) -> None:
self.close_video_recorder()
video_name = f"{self.name_prefix}-step-{self.step_id}-to-step-{self.step_id + self.video_length}"
@ -76,10 +82,10 @@ class VecVideoRecorder(VecEnvWrapper):
self.recorded_frames = 1
self.recording = True
def _video_enabled(self):
def _video_enabled(self) -> bool:
return self.record_video_trigger(self.step_id)
def step_wait(self):
def step_wait(self) -> VecEnvStepReturn:
obs, rews, dones, infos = self.venv.step_wait()
self.step_id += 1
@ -94,13 +100,13 @@ class VecVideoRecorder(VecEnvWrapper):
return obs, rews, dones, infos
def close_video_recorder(self):
def close_video_recorder(self) -> None:
if self.recording:
self.video_recorder.close()
self.recording = False
self.recorded_frames = 1
def close(self):
def close(self) -> None:
VecEnvWrapper.close(self)
self.close_video_recorder()

5
stable_baselines3/dqn/dqn.py
View file

@ -137,7 +137,7 @@ class DQN(OffPolicyAlgorithm):
self.q_net = self.policy.q_net
self.q_net_target = self.policy.q_net_target
def _on_step(self):
def _on_step(self) -> None:
"""
Update the exploration rate and target network if needed.
This method is called in ``collect_rollout()`` after each step in the environment.
@ -152,6 +152,7 @@ class DQN(OffPolicyAlgorithm):
# Update learning rate according to schedule
self._update_learning_rate(self.policy.optimizer)
losses = []
for gradient_step in range(gradient_steps):
# Sample replay buffer
replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
@ -174,6 +175,7 @@ class DQN(OffPolicyAlgorithm):
# Compute Huber loss (less sensitive to outliers)
loss = F.smooth_l1_loss(current_q, target_q)
losses.append(loss.item())
# Optimize the policy
self.policy.optimizer.zero_grad()
@ -186,6 +188,7 @@ class DQN(OffPolicyAlgorithm):
self._n_updates += gradient_steps
logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
logger.record("train/loss", np.mean(losses))
def predict(
self,

22
stable_baselines3/sac/policies.py
View file

@ -1,4 +1,4 @@
from typing import Any, Callable, Dict, List, Optional, Tuple, Type
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
import gym
import torch as th
@ -7,7 +7,13 @@ from torch import nn
from stable_baselines3.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
from stable_baselines3.common.policies import BasePolicy, ContinuousCritic, create_sde_features_extractor, register_policy
from stable_baselines3.common.preprocessing import get_action_dim
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp
from stable_baselines3.common.torch_layers import (
BaseFeaturesExtractor,
FlattenExtractor,
NatureCNN,
create_mlp,
get_actor_critic_arch,
)
# CAP the standard deviation of the actor
LOG_STD_MAX = 2
@ -232,7 +238,7 @@ class SACPolicy(BasePolicy):
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[int]] = None,
net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
activation_fn: Type[nn.Module] = nn.ReLU,
use_sde: bool = False,
log_std_init: float = -3,
@ -262,6 +268,8 @@ class SACPolicy(BasePolicy):
else:
net_arch = []
actor_arch, critic_arch = get_actor_critic_arch(net_arch)
# Create shared features extractor
self.features_extractor = features_extractor_class(self.observation_space, **self.features_extractor_kwargs)
self.features_dim = self.features_extractor.features_dim
@ -273,7 +281,7 @@ class SACPolicy(BasePolicy):
"action_space": self.action_space,
"features_extractor": self.features_extractor,
"features_dim": self.features_dim,
"net_arch": self.net_arch,
"net_arch": actor_arch,
"activation_fn": self.activation_fn,
"normalize_images": normalize_images,
}
@ -287,7 +295,7 @@ class SACPolicy(BasePolicy):
}
self.actor_kwargs.update(sde_kwargs)
self.critic_kwargs = self.net_args.copy()
self.critic_kwargs.update({"n_critics": n_critics})
self.critic_kwargs.update({"n_critics": n_critics, "net_arch": critic_arch})
self.actor, self.actor_target = None, None
self.critic, self.critic_target = None, None
@ -312,7 +320,7 @@ class SACPolicy(BasePolicy):
data.update(
dict(
net_arch=self.net_args["net_arch"],
net_arch=self.net_arch,
activation_fn=self.net_args["activation_fn"],
use_sde=self.actor_kwargs["use_sde"],
log_std_init=self.actor_kwargs["log_std_init"],
@ -386,7 +394,7 @@ class CnnPolicy(SACPolicy):
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[int]] = None,
net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
activation_fn: Type[nn.Module] = nn.ReLU,
use_sde: bool = False,
log_std_init: float = -3,

27
stable_baselines3/td3/policies.py
View file

@ -1,4 +1,4 @@
from typing import Any, Callable, Dict, List, Optional, Type
from typing import Any, Callable, Dict, List, Optional, Type, Union
import gym
import torch as th
@ -6,7 +6,13 @@ from torch import nn
from stable_baselines3.common.policies import BasePolicy, ContinuousCritic, register_policy
from stable_baselines3.common.preprocessing import get_action_dim
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp
from stable_baselines3.common.torch_layers import (
BaseFeaturesExtractor,
FlattenExtractor,
NatureCNN,
create_mlp,
get_actor_critic_arch,
)
class Actor(BasePolicy):
@ -101,7 +107,7 @@ class TD3Policy(BasePolicy):
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[int]] = None,
net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
activation_fn: Type[nn.Module] = nn.ReLU,
features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
features_extractor_kwargs: Optional[Dict[str, Any]] = None,
@ -127,6 +133,8 @@ class TD3Policy(BasePolicy):
else:
net_arch = []
actor_arch, critic_arch = get_actor_critic_arch(net_arch)
self.features_extractor = features_extractor_class(self.observation_space, **self.features_extractor_kwargs)
self.features_dim = self.features_extractor.features_dim
@ -137,12 +145,13 @@ class TD3Policy(BasePolicy):
"action_space": self.action_space,
"features_extractor": self.features_extractor,
"features_dim": self.features_dim,
"net_arch": self.net_arch,
"net_arch": actor_arch,
"activation_fn": self.activation_fn,
"normalize_images": normalize_images,
}
self.actor_kwargs = self.net_args.copy()
self.critic_kwargs = self.net_args.copy()
self.critic_kwargs.update({"n_critics": n_critics})
self.critic_kwargs.update({"n_critics": n_critics, "net_arch": critic_arch})
self.actor, self.actor_target = None, None
self.critic, self.critic_target = None, None
@ -163,7 +172,7 @@ class TD3Policy(BasePolicy):
data.update(
dict(
net_arch=self.net_args["net_arch"],
net_arch=self.net_arch,
activation_fn=self.net_args["activation_fn"],
n_critics=self.critic_kwargs["n_critics"],
lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone
@ -176,12 +185,12 @@ class TD3Policy(BasePolicy):
return data
def make_actor(self) -> Actor:
return Actor(**self.net_args).to(self.device)
return Actor(**self.actor_kwargs).to(self.device)
def make_critic(self) -> ContinuousCritic:
return ContinuousCritic(**self.critic_kwargs).to(self.device)
def forward(self, observation: th.Tensor, deterministic: bool = False):
def forward(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
return self._predict(observation, deterministic=deterministic)
def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
@ -217,7 +226,7 @@ class CnnPolicy(TD3Policy):
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[int]] = None,
net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
activation_fn: Type[nn.Module] = nn.ReLU,
features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN,
features_extractor_kwargs: Optional[Dict[str, Any]] = None,

2
stable_baselines3/version.txt
View file

@ -1 +1 @@
0.9.0a2
0.10.0a0

23
tests/test_custom_policy.py
View file

@ -1,7 +1,7 @@
import pytest
import torch as th
from stable_baselines3 import A2C, PPO, SAC, TD3
from stable_baselines3 import A2C, DQN, PPO, SAC, TD3
from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike
@ -19,22 +19,33 @@ from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike
)
@pytest.mark.parametrize("model_class", [A2C, PPO])
def test_flexible_mlp(model_class, net_arch):
_ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
_ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(300)
@pytest.mark.parametrize("net_arch", [[4], [4, 4]])
@pytest.mark.parametrize("net_arch", [[], [4], [4, 4], dict(qf=[8], pi=[8, 4])])
@pytest.mark.parametrize("model_class", [SAC, TD3])
def test_custom_offpolicy(model_class, net_arch):
_ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=dict(net_arch=net_arch)).learn(1000)
_ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=dict(net_arch=net_arch), learning_starts=100).learn(300)
@pytest.mark.parametrize("model_class", [A2C, PPO, SAC, TD3])
@pytest.mark.parametrize("optimizer_kwargs", [None, dict(weight_decay=0.0)])
def test_custom_optimizer(model_class, optimizer_kwargs):
kwargs = {}
if model_class in {DQN, SAC, TD3}:
kwargs = dict(learning_starts=100)
elif model_class in {A2C, PPO}:
kwargs = dict(n_steps=100)
policy_kwargs = dict(optimizer_class=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32])
_ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000)
_ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs, **kwargs).learn(300)
def test_tf_like_rmsprop_optimizer():
policy_kwargs = dict(optimizer_class=RMSpropTFLike, net_arch=[32])
_ = A2C("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000)
_ = A2C("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(500)
def test_dqn_custom_policy():
policy_kwargs = dict(optimizer_class=RMSpropTFLike, net_arch=[32])
_ = DQN("MlpPolicy", "CartPole-v1", policy_kwargs=policy_kwargs, learning_starts=100).learn(300)

2
tests/test_identity.py
View file

@ -23,7 +23,7 @@ def test_discrete(model_class, env):
if isinstance(env, (IdentityEnvMultiDiscrete, IdentityEnvMultiBinary)):
return
model = model_class("MlpPolicy", env_, gamma=0.5, seed=1, **kwargs).learn(n_steps)
model = model_class("MlpPolicy", env_, gamma=0.4, seed=1, **kwargs).learn(n_steps)
evaluate_policy(model, env_, n_eval_episodes=20, reward_threshold=90)
obs = env.reset()

6
tests/test_run.py
View file

@ -80,7 +80,7 @@ def test_n_critics(n_critics):
model = SAC(
"MlpPolicy", "Pendulum-v0", policy_kwargs=dict(net_arch=[64, 64], n_critics=n_critics), learning_starts=100, verbose=1
)
model.learn(total_timesteps=1000)
model.learn(total_timesteps=500)
# "CartPole-v1"
@ -122,10 +122,10 @@ def test_dqn():
"MlpPolicy",
"CartPole-v1",
policy_kwargs=dict(net_arch=[64, 64]),
learning_starts=500,
learning_starts=100,
buffer_size=500,
learning_rate=3e-4,
verbose=1,
create_eval_env=True,
)
model.learn(total_timesteps=1000, eval_freq=500)
model.learn(total_timesteps=500, eval_freq=250)

20
tests/test_save_load.py
View file

@ -45,7 +45,7 @@ def test_save_load(tmp_path, model_class):
# create model
model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16]), verbose=1)
model.learn(total_timesteps=500, eval_freq=250)
model.learn(total_timesteps=500)
env.reset()
observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0)
@ -154,7 +154,7 @@ def test_save_load(tmp_path, model_class):
assert np.allclose(selected_actions, new_selected_actions, 1e-4)
# check if learn still works
model.learn(total_timesteps=1000, eval_freq=500)
model.learn(total_timesteps=500)
del model
@ -174,20 +174,26 @@ def test_set_env(model_class):
env2 = DummyVecEnv([lambda: select_env(model_class)])
env3 = select_env(model_class)
kwargs = {}
if model_class in {DQN, DDPG, SAC, TD3}:
kwargs = dict(learning_starts=100)
elif model_class in {A2C, PPO}:
kwargs = dict(n_steps=100)
# create model
model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16]))
model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16]), **kwargs)
# learn
model.learn(total_timesteps=1000, eval_freq=500)
model.learn(total_timesteps=300)
# change env
model.set_env(env2)
# learn again
model.learn(total_timesteps=1000, eval_freq=500)
model.learn(total_timesteps=300)
# change env test wrapping
model.set_env(env3)
# learn again
model.learn(total_timesteps=1000, eval_freq=500)
model.learn(total_timesteps=300)
@pytest.mark.parametrize("model_class", MODEL_LIST)
@ -309,7 +315,7 @@ def test_save_load_policy(tmp_path, model_class, policy_str):
# create model
model = model_class(policy_str, env, policy_kwargs=dict(net_arch=[16]), verbose=1, **kwargs)
model.learn(total_timesteps=500, eval_freq=250)
model.learn(total_timesteps=500)
env.reset()
observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0)

3
tests/test_sde.py
View file

@ -68,3 +68,6 @@ def test_state_dependent_offpolicy_noise(model_class, sde_net_arch, use_expln):
policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch, use_expln=use_expln),
)
model.learn(total_timesteps=int(500), eval_freq=250)
model.policy.reset_noise()
if model_class == SAC:
model.policy.actor.get_std()

20
tests/test_utils.py
View file

@ -12,7 +12,7 @@ from stable_baselines3.common.cmd_util import make_atari_env, make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise, VectorizedActionNoise
from stable_baselines3.common.utils import polyak_update
from stable_baselines3.common.utils import polyak_update, zip_strict
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
@ -167,3 +167,21 @@ def test_polyak():
assert th.allclose(param1, target1)
assert th.allclose(param2, target2)
def test_zip_strict():
# Iterables with different lengths
list_a = [0, 1]
list_b = [1, 2, 3]
# zip does not raise any error
for _, _ in zip(list_a, list_b):
pass
# zip_strict does raise an error
with pytest.raises(ValueError):
for _, _ in zip_strict(list_a, list_b):
pass
# same length, should not raise an error
for _, _ in zip_strict(list_a, list_b[: len(list_a)]):
pass