Merge branch 'master' into feat/mps-support

.gitignore

@@ -4,6 +4,7 @@
 *.py~
 *.bak
 .pytest_cache
+.mypy_cache
 .DS_Store
 .idea
 .vscode

.gitlab-ci.yml

@@ -2,7 +2,7 @@ image: stablebaselines/stable-baselines3-cpu:1.4.1a0
 
 type-check:
   script:
-  - pip install pytype --upgrade
+  - pip install pytype mypy --upgrade
   - make type
 
 pytest:

CONTRIBUTING.md

@@ -78,7 +78,7 @@ To run tests with `pytest`:
 make pytest
 ```
 
-Type checking with `pytype`:
+Type checking with `pytype` and `mypy`:
 
 ```
 make type
@@ -91,7 +91,7 @@ make check-codestyle
 make lint
 ```
 
-To run `pytype`, `format` and `lint` in one command:
+To run `type`, `format` and `lint` in one command:
 ```
 make commit-checks
 ```

Makefile

@@ -4,9 +4,14 @@ LINT_PATHS=stable_baselines3/ tests/ docs/conf.py setup.py
 pytest:
 	./scripts/run_tests.sh
 
-type:
+pytype:
 	pytype -j auto
 
+mypy:
+	MYPY_FORCE_COLOR=1 mypy ${LINT_PATHS}
+
+type: pytype mypy
+
 lint:
 	# stop the build if there are Python syntax errors or undefined names
 	# see https://lintlyci.github.io/Flake8Rules/

README.md

@@ -15,7 +15,7 @@ You can read a detailed presentation of Stable Baselines3 in the [v1.0 blog post
 
 These algorithms will make it easier for the research community and industry to replicate, refine, and identify new ideas, and will create good baselines to build projects on top of. We expect these tools will be used as a base around which new ideas can be added, and as a tool for comparing a new approach against existing ones. We also hope that the simplicity of these tools will allow beginners to experiment with a more advanced toolset, without being buried in implementation details.
 
-**Note: despite its simplicity of use, Stable Baselines3 (SB3) assumes you have some knowledge about Reinforcement Learning (RL).** You should not utilize this library without some practice. To that extent, we provide good resources in the [documentation](https://stable-baselines3.readthedocs.io/en/master/guide/rl.html) to get started with RL.
+**Note: Despite its simplicity of use, Stable Baselines3 (SB3) assumes you have some knowledge about Reinforcement Learning (RL).** You should not utilize this library without some practice. To that extent, we provide good resources in the [documentation](https://stable-baselines3.readthedocs.io/en/master/guide/rl.html) to get started with RL.
 
 ## Main Features
 
@@ -113,7 +113,7 @@ Please read the [documentation](https://stable-baselines3.readthedocs.io/) for m
 
 ## Example
 
-Most of the library tries to follow a sklearn-like syntax for the Reinforcement Learning algorithms.
+Most of the code in the library tries to follow a sklearn-like syntax for the Reinforcement Learning algorithms.
 
 Here is a quick example of how to train and run PPO on a cartpole environment:
 ```python
@@ -150,7 +150,7 @@ Please read the [documentation](https://stable-baselines3.readthedocs.io/) for m
 
 ## 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)
 - [All Notebooks](https://github.com/Stable-Baselines-Team/rl-colab-notebooks/tree/sb3)
@@ -198,9 +198,9 @@ pip install pytest pytest-cov
 make pytest
 ```
 
-You can also do a static type check using `pytype`:
+You can also do a static type check using `pytype` and `mypy`:
 ```
-pip install pytype
+pip install pytype mypy
 make type
 ```
 
@@ -212,8 +212,8 @@ make lint
 
 ## Projects Using Stable-Baselines3
 
-We try to maintain a list of project using stable-baselines3 in the [documentation](https://stable-baselines3.readthedocs.io/en/master/misc/projects.html),
-please tell us when if you want your project to appear on this page ;)
+We try to maintain a list of projects using stable-baselines3 in the [documentation](https://stable-baselines3.readthedocs.io/en/master/misc/projects.html),
+please tell us if you want your project to appear on this page ;)
 
 ## Citing the Project
 
@@ -236,8 +236,8 @@ To cite this repository in publications:
 
 Stable-Baselines3 is currently maintained by [Ashley Hill](https://github.com/hill-a) (aka @hill-a), [Antonin Raffin](https://araffin.github.io/) (aka [@araffin](https://github.com/araffin)), [Maximilian Ernestus](https://github.com/ernestum) (aka @ernestum), [Adam Gleave](https://github.com/adamgleave) (@AdamGleave), [Anssi Kanervisto](https://github.com/Miffyli) (@Miffyli) and [Quentin Gallouédec](https://gallouedec.com/) (@qgallouedec).
 
-**Important Note: We do not do technical support, nor consulting** and don't answer personal questions per email.
-Please post your question on the [RL Discord](https://discord.com/invite/xhfNqQv), [Reddit](https://www.reddit.com/r/reinforcementlearning/) or [Stack Overflow](https://stackoverflow.com/) in that case.
+**Important Note: We do not provide technical support, or consulting** and do not answer personal questions via email.
+Please post your question on the [RL Discord](https://discord.com/invite/xhfNqQv), [Reddit](https://www.reddit.com/r/reinforcementlearning/), or [Stack Overflow](https://stackoverflow.com/) in that case.
 
 
 ## How To Contribute

docs/conf.py

@@ -13,6 +13,7 @@
 #
 import os
 import sys
+from typing import Dict, List
 from unittest.mock import MagicMock
 
 # We CANNOT enable 'sphinxcontrib.spelling' because ReadTheDocs.org does not support
@@ -37,7 +38,7 @@ sys.path.insert(0, os.path.abspath(".."))
 
 
 class Mock(MagicMock):
-    __subclasses__ = []
+    __subclasses__ = []  # type: ignore
 
     @classmethod
     def __getattr__(cls, name):
@@ -48,7 +49,7 @@ class Mock(MagicMock):
 # Note: because of that we cannot test examples using CI
 # 'torch', 'torch.nn', 'torch.nn.functional',
 # DO not mock modules for now, we will need to do that for read the docs later
-MOCK_MODULES = []
+MOCK_MODULES: List[str] = []
 sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)
 
 # Read version from file
@@ -171,7 +172,7 @@ htmlhelp_basename = "StableBaselines3doc"
 
 # -- Options for LaTeX output ------------------------------------------------
 
-latex_elements = {
+latex_elements: Dict[str, str] = {
     # The paper size ('letterpaper' or 'a4paper').
     #
     # 'papersize': 'letterpaper',

docs/guide/checking_nan.rst

@@ -7,7 +7,7 @@ corrupted when a NaN or an inf is given or returned from the RL model.
 How and why?
 ------------
 
-The issue arises then NaNs or infs do not crash, but simply get propagated through the training,
+The issue arises when NaNs or infs do not crash, but simply get propagated through the training,
 until all the floating point number converge to NaN or inf. This is in line with the
 `IEEE Standard for Floating-Point Arithmetic (IEEE 754) <https://ieeexplore.ieee.org/document/4610935>`_ standard, as it says:
 

docs/misc/changelog.rst

@@ -16,6 +16,7 @@ Breaking Changes:
 
 New Features:
 ^^^^^^^^^^^^^
+- Introduced mypy type checking
 
 SB3-Contrib
 ^^^^^^^^^^^
@@ -25,6 +26,7 @@ Bug Fixes:
 - Fix return type of ``evaluate_actions`` in ``ActorCritcPolicy`` to reflect that entropy is an optional tensor (@Rocamonde)
 - Fix type annotation of ``policy`` in ``BaseAlgorithm`` and ``OffPolicyAlgorithm``
 - Allowed model trained with Python 3.7 to be loaded with Python 3.8+ without the ``custom_objects`` workaround
+- Fix type annotation of ``model`` in ``evaluate_policy``
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -1102,7 +1104,7 @@ In random order...
 Thanks to the maintainers of V2: @hill-a @enerijunior @AdamGleave @Miffyli
 
 And all the contributors:
-@bjmuld @iambenzo @iandanforth @r7vme @brendenpetersen @huvar @abhiskk @JohannesAck
+@taymuur @bjmuld @iambenzo @iandanforth @r7vme @brendenpetersen @huvar @abhiskk @JohannesAck
 @EliasHasle @mrakgr @Bleyddyn @antoine-galataud @junhyeokahn @AdamGleave @keshaviyengar @tperol
 @XMaster96 @kantneel @Pastafarianist @GerardMaggiolino @PatrickWalter214 @yutingsz @sc420 @Aaahh @billtubbs
 @Miffyli @dwiel @miguelrass @qxcv @jaberkow @eavelardev @ruifeng96150 @pedrohbtp @srivatsankrishnan @evilsocket

setup.cfg

@@ -25,6 +25,60 @@ markers =
 inputs = stable_baselines3
 disable = pyi-error
 
+[mypy]
+ignore_missing_imports = True
+follow_imports = silent
+show_error_codes = True
+exclude = (?x)(
+    stable_baselines3/a2c/a2c.py$
+    | stable_baselines3/common/atari_wrappers.py$
+    | stable_baselines3/common/base_class.py$
+    | stable_baselines3/common/buffers.py$
+    | stable_baselines3/common/callbacks.py$
+    | stable_baselines3/common/distributions.py$
+    | stable_baselines3/common/env_util.py$
+    | stable_baselines3/common/envs/bit_flipping_env.py$
+    | stable_baselines3/common/envs/identity_env.py$
+    | stable_baselines3/common/envs/multi_input_envs.py$
+    | stable_baselines3/common/logger.py$
+    | stable_baselines3/common/monitor.py$
+    | stable_baselines3/common/off_policy_algorithm.py$
+    | stable_baselines3/common/on_policy_algorithm.py$
+    | stable_baselines3/common/policies.py$
+    | stable_baselines3/common/preprocessing.py$
+    | stable_baselines3/common/save_util.py$
+    | stable_baselines3/common/sb2_compat/rmsprop_tf_like.py$
+    | stable_baselines3/common/torch_layers.py$
+    | stable_baselines3/common/type_aliases.py$
+    | stable_baselines3/common/utils.py$
+    | stable_baselines3/common/vec_env/__init__.py$
+    | stable_baselines3/common/vec_env/base_vec_env.py$
+    | stable_baselines3/common/vec_env/dummy_vec_env.py$
+    | stable_baselines3/common/vec_env/stacked_observations.py$
+    | stable_baselines3/common/vec_env/subproc_vec_env.py$
+    | stable_baselines3/common/vec_env/util.py$
+    | stable_baselines3/common/vec_env/vec_check_nan.py$
+    | stable_baselines3/common/vec_env/vec_extract_dict_obs.py$
+    | stable_baselines3/common/vec_env/vec_frame_stack.py$
+    | stable_baselines3/common/vec_env/vec_monitor.py$
+    | stable_baselines3/common/vec_env/vec_normalize.py$
+    | stable_baselines3/common/vec_env/vec_transpose.py$
+    | stable_baselines3/common/vec_env/vec_video_recorder.py$
+    | stable_baselines3/dqn/dqn.py$
+    | stable_baselines3/dqn/policies.py$
+    | stable_baselines3/her/her_replay_buffer.py$
+    | stable_baselines3/ppo/ppo.py$
+    | stable_baselines3/sac/policies.py$
+    | stable_baselines3/sac/sac.py$
+    | stable_baselines3/td3/policies.py$
+    | stable_baselines3/td3/td3.py$
+    | tests/test_distributions.py$
+    | tests/test_logger.py$
+    | tests/test_tensorboard.py$
+    | tests/test_train_eval_mode.py$
+    | tests/test_vec_normalize.py$
+  )
+
 [flake8]
 ignore = W503,W504,E203,E231  # line breaks before and after binary operators
 # Ignore import not used when aliases are defined

setup.py

@@ -76,6 +76,7 @@ setup(
         "gym==0.21",  # Fixed version due to breaking changes in 0.22
         "numpy",
         "torch>=1.11",
+        'typing_extensions>=4.0,<5; python_version < "3.8.0"',
         # For saving models
         "cloudpickle",
         # For reading logs
@@ -94,6 +95,7 @@ setup(
             "pytest-xdist",
             # Type check
             "pytype",
+            "mypy",
             # Lint code
             "flake8>=3.8",
             # Find likely bugs

stable_baselines3/common/base_class.py

@@ -513,7 +513,7 @@ class BaseAlgorithm(ABC):
 
     def predict(
         self,
-        observation: np.ndarray,
+        observation: Union[np.ndarray, Dict[str, np.ndarray]],
         state: Optional[Tuple[np.ndarray, ...]] = None,
         episode_start: Optional[np.ndarray] = None,
         deterministic: bool = False,

stable_baselines3/common/evaluation.py

@@ -4,12 +4,12 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import gym
 import numpy as np
 
-from stable_baselines3.common import base_class
+from stable_baselines3.common import type_aliases
 from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, VecMonitor, is_vecenv_wrapped
 
 
 def evaluate_policy(
-    model: "base_class.BaseAlgorithm",
+    model: "type_aliases.PolicyPredictor",
     env: Union[gym.Env, VecEnv],
     n_eval_episodes: int = 10,
     deterministic: bool = True,
@@ -34,7 +34,9 @@ def evaluate_policy(
         results as well. You can avoid this by wrapping environment with ``Monitor``
         wrapper before anything else.
 
-    :param model: The RL agent you want to evaluate.
+    :param model: The RL agent you want to evaluate. This can be any object
+        that implements a `predict` method, such as an RL algorithm (``BaseAlgorithm``)
+        or policy (``BasePolicy``).
     :param env: The gym environment or ``VecEnv`` environment.
     :param n_eval_episodes: Number of episode to evaluate the agent
     :param deterministic: Whether to use deterministic or stochastic actions

stable_baselines3/common/type_aliases.py

@@ -1,12 +1,18 @@
 """Common aliases for type hints"""
 
+import sys
 from enum import Enum
-from typing import Any, Callable, Dict, List, NamedTuple, Tuple, Union
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Union
 
 import gym
 import numpy as np
 import torch as th
 
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
 from stable_baselines3.common import callbacks, vec_env
 
 GymEnv = Union[gym.Env, vec_env.VecEnv]
@@ -69,3 +75,26 @@ class TrainFrequencyUnit(Enum):
 class TrainFreq(NamedTuple):
     frequency: int
     unit: TrainFrequencyUnit  # either "step" or "episode"
+
+
+class PolicyPredictor(Protocol):
+    def predict(
+        self,
+        observation: Union[np.ndarray, Dict[str, np.ndarray]],
+        state: Optional[Tuple[np.ndarray, ...]] = None,
+        episode_start: Optional[np.ndarray] = None,
+        deterministic: bool = False,
+    ) -> Tuple[np.ndarray, Optional[Tuple[np.ndarray, ...]]]:
+        """
+        Get the policy action from an observation (and optional hidden state).
+        Includes sugar-coating to handle different observations (e.g. normalizing images).
+
+        :param observation: the input observation
+        :param state: The last hidden states (can be None, used in recurrent policies)
+        :param episode_start: The last masks (can be None, used in recurrent policies)
+            this correspond to beginning of episodes,
+            where the hidden states of the RNN must be reset.
+        :param deterministic: Whether or not to return deterministic actions.
+        :return: the model's action and the next hidden state
+            (used in recurrent policies)
+        """

stable_baselines3/dqn/dqn.py

@@ -224,7 +224,7 @@ class DQN(OffPolicyAlgorithm):
 
     def predict(
         self,
-        observation: np.ndarray,
+        observation: Union[np.ndarray, Dict[str, np.ndarray]],
         state: Optional[Tuple[np.ndarray, ...]] = None,
         episode_start: Optional[np.ndarray] = None,
         deterministic: bool = False,
@@ -241,7 +241,7 @@ class DQN(OffPolicyAlgorithm):
         """
         if not deterministic and np.random.rand() < self.exploration_rate:
             if is_vectorized_observation(maybe_transpose(observation, self.observation_space), self.observation_space):
-                if isinstance(self.observation_space, gym.spaces.Dict):
+                if isinstance(observation, dict):
                     n_batch = observation[list(observation.keys())[0]].shape[0]
                 else:
                     n_batch = observation.shape[0]

tests/test_utils.py

@@ -141,16 +141,19 @@ def test_custom_vec_env(tmp_path):
         make_vec_env("CartPole-v1", n_envs=1, vec_env_kwargs={"dummy": False})
 
 
-def test_evaluate_policy():
+@pytest.mark.parametrize("direct_policy", [False, True])
+def test_evaluate_policy(direct_policy: bool):
     model = A2C("MlpPolicy", "Pendulum-v1", 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
 
+    assert model.policy is not None
+    policy = model.policy if direct_policy else model
+    policy.n_callback_calls = 0
     _, episode_lengths = evaluate_policy(
-        model,
+        policy,
         model.get_env(),
         n_eval_episodes,
         deterministic=True,
@@ -162,19 +165,19 @@ def test_evaluate_policy():
 
     n_steps = sum(episode_lengths)
     assert n_steps == n_steps_per_episode * n_eval_episodes
-    assert n_steps == model.n_callback_calls
+    assert n_steps == policy.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)
+        evaluate_policy(policy, 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)
+    episode_rewards, _ = evaluate_policy(policy, model.get_env(), n_eval_episodes, return_episode_rewards=True)
     assert len(episode_rewards) == n_eval_episodes
 
     # Test that warning is given about no monitor
     eval_env = gym.make("Pendulum-v1")
     with pytest.warns(UserWarning):
-        _ = evaluate_policy(model, eval_env, n_eval_episodes)
+        _ = evaluate_policy(policy, eval_env, n_eval_episodes)
 
 
 class ZeroRewardWrapper(gym.RewardWrapper):