Merge pull request #47 from Antonin-Raffin/feat/callbacks

Add callback support

.coveragerc

@@ -8,3 +8,4 @@ omit =
 exclude_lines =
     pragma: no cover
     raise NotImplementedError()
+    if typing.TYPE_CHECKING:

.github/ISSUE_TEMPLATE/issue-template.md

@@ -0,0 +1,52 @@
+---
+name: Issue Template
+about: How to create an issue for this repository
+
+---
+
+**Important Note: We do not do technical support, nor consulting** and don't answer personal questions per email.
+
+If you have any questions, feel free to create an issue with the tag [question].  
+If you wish to suggest an enhancement or feature request, add the tag [feature request].  
+If you are submitting a bug report, please fill in the following details.
+
+If your issue is related to a custom gym environment, please check it first using:
+
+```python
+from torchy_baselines.common.env_checker import check_env
+
+env = CustomEnv(arg1, ...)
+# It will check your custom environment and output additional warnings if needed
+check_env(env)
+```
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**Code example**
+Please try to provide a minimal example to reproduce the bug. Error messages and stack traces are also helpful.
+
+Please use the [markdown code blocks](https://help.github.com/en/articles/creating-and-highlighting-code-blocks)
+for both code and stack traces.
+
+```python
+from torchy_baselines import ...
+
+```
+
+```bash
+Traceback (most recent call last): File ...
+
+```
+
+**System Info**
+Describe the characteristic of your environment:
+ * Describe how the library was installed (pip, docker, source, ...)
+ * GPU models and configuration
+ * Python version
+ * PyTorch version
+ * Gym version
+ * Versions of any other relevant libraries
+
+**Additional context**
+Add any other context about the problem here.

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,29 @@
+<!--- Provide a general summary of your changes in the Title above -->
+
+## Description
+<!--- Describe your changes in detail -->
+
+## Motivation and Context
+<!--- Why is this change required? What problem does it solve? -->
+<!--- If it fixes an open issue, please link to the issue here. -->
+<!--- You can use the syntax `closes #100` if this solves the issue #100 -->
+- [ ] I have raised an issue to propose this change ([required](https://github.com/hill-a/stable-baselines/blob/master/CONTRIBUTING.md) for new features and bug fixes)
+
+## Types of changes
+<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->
+- [ ] Bug fix (non-breaking change which fixes an issue)
+- [ ] New feature (non-breaking change which adds functionality)
+- [ ] Breaking change (fix or feature that would cause existing functionality to change)
+- [ ] Documentation (update in the documentation)
+
+## Checklist:
+<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
+<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
+- [ ] I've read the [CONTRIBUTION](https://github.com/hill-a/stable-baselines/blob/master/CONTRIBUTING.md) guide (**required**)
+- [ ] I have updated the changelog accordingly (**required**).
+- [ ] My change requires a change to the documentation.
+- [ ] I have updated the tests accordingly (*required for a bug fix or a new feature*).
+- [ ] I have updated the documentation accordingly.
+- [ ] I have ensured `pytest` and `pytype` both pass.
+
+<!--- This Template is an edited version of the one from https://github.com/evilsocket/pwnagotchi/ -->

Makefile

@@ -6,8 +6,11 @@ pytest:
 type:
 	pytype
 
-docs:
+doc:
 	cd docs && make html
 
 spelling:
 	cd docs && make spelling
+
+clean:
+	cd docs && make clean

docs/conf.py

@@ -75,6 +75,7 @@ extensions = [
     'sphinx.ext.mathjax',
     'sphinx.ext.ifconfig',
     'sphinx.ext.viewcode',
+    # 'sphinx.ext.intersphinx',
 ]
 
 if enable_spell_check:
@@ -206,3 +207,10 @@ texinfo_documents = [
 
 
 # -- Extension configuration -------------------------------------------------
+
+# Example configuration for intersphinx: refer to the Python standard library.
+# intersphinx_mapping = {
+#     'python': ('https://docs.python.org/3/', None),
+#     'numpy': ('http://docs.scipy.org/doc/numpy/', None),
+#     'torch': ('http://pytorch.org/docs/master/', None),
+# }

docs/misc/changelog.rst

@@ -9,14 +9,17 @@ Pre-Release 0.2.0a0 (WIP)
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - Python 2 support was dropped, Torchy Baselines now requires Python 3.6 or above
+- Return type of `evaluation.evaluate_policy()` has been changed
 
 New Features:
 ^^^^^^^^^^^^^
 - Add `seed()` method to `VecEnv` class
+- Add support for Callback (cf https://github.com/hill-a/stable-baselines/pull/644)
 
 Bug Fixes:
 ^^^^^^^^^^
 - Fix loading model on CPU that were trained on GPU
+- Fix `reset_num_timesteps` that was not used
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -25,6 +28,8 @@ Others:
 ^^^^^^^
 - Add type check
 - Converted all format string to f-strings
+- Add test for `OrnsteinUhlenbeckActionNoise`
+- Add type aliases in `common.type_aliases`
 
 Documentation:
 ^^^^^^^^^^^^^^

tests/test_callbacks.py

@@ -0,0 +1,49 @@
+import os
+import shutil
+
+import pytest
+import gym
+
+from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
+from torchy_baselines.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback,
+    EveryNTimesteps, StopTrainingOnRewardThreshold)
+
+
+@pytest.mark.parametrize("model_class", [A2C, CEMRL, PPO, SAC, TD3])
+def test_callbacks(model_class):
+    log_folder = './logs/callbacks/'
+    kwargs = {}
+    if model_class == CEMRL:
+        kwargs['pop_size'] = 2
+        kwargs['n_grad'] = 1
+
+    # Create RL model
+    # Small network for fast test
+    model = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[32]), **kwargs)
+
+    checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_folder)
+
+    eval_env = gym.make('Pendulum-v0')
+    # Stop training if the performance is good enough
+    callback_on_best = StopTrainingOnRewardThreshold(reward_threshold=-1200, verbose=1)
+
+    eval_callback = EvalCallback(eval_env, callback_on_new_best=callback_on_best,
+                                 best_model_save_path=log_folder,
+                                 log_path=log_folder, eval_freq=100)
+
+    # Equivalent to the `checkpoint_callback`
+    # but here in an event-driven manner
+    checkpoint_on_event = CheckpointCallback(save_freq=1, save_path=log_folder,
+                                             name_prefix='event')
+    event_callback = EveryNTimesteps(n_steps=500, callback=checkpoint_on_event)
+
+    callback = CallbackList([checkpoint_callback, eval_callback, event_callback])
+
+    model.learn(500, callback=callback)
+    model.learn(500, callback=None)
+    # Transform callback into a callback list automatically
+    model.learn(500, callback=[checkpoint_callback, eval_callback])
+    # Automatic wrapping, old way of doing callbacks
+    model.learn(500, callback=lambda _locals, _globals : True)
+    if os.path.exists(log_folder):
+        shutil.rmtree(log_folder)

tests/test_run.py

@@ -4,12 +4,13 @@ import pytest
 import numpy as np
 
 from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
-from torchy_baselines.common.noise import NormalActionNoise
+from torchy_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
 
 action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1))
 
 
-def test_td3():
+@pytest.mark.parametrize('action_noise', [action_noise, OrnsteinUhlenbeckActionNoise(np.zeros(1), 0.1 * np.ones(1))])
+def test_td3(action_noise):
     model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]),
                 learning_starts=100, verbose=1, create_eval_env=True, action_noise=action_noise)
     model.learn(total_timesteps=1000, eval_freq=500)

tests/test_save_load.py

@@ -69,7 +69,7 @@ def test_save_load(model_class):
 
     # check if model still selects the same actions
     new_selected_actions = [model.predict(observation, deterministic=True) for observation in observations]
-    assert np.allclose(selected_actions, new_selected_actions)
+    assert np.allclose(selected_actions, new_selected_actions, 1e-4)
 
     # check if learn still works
     model.learn(total_timesteps=1000, eval_freq=500)

torchy_baselines/a2c/a2c.py

@@ -130,8 +130,10 @@ class A2C(PPO):
             logger.logkv("std", th.exp(self.policy.log_std).mean().item())
 
     def learn(self, total_timesteps, callback=None, log_interval=100,
-              eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="A2C", reset_num_timesteps=True):
+              eval_env=None, eval_freq=-1, n_eval_episodes=5,
+              tb_log_name="A2C", eval_log_path=None, reset_num_timesteps=True):
 
         return super(A2C, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval,
                                       eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
-                                      tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)
+                                      tb_log_name=tb_log_name, eval_log_path=eval_log_path,
+                                      reset_num_timesteps=reset_num_timesteps)

torchy_baselines/cem_rl/cem_rl.py

@@ -100,21 +100,21 @@ class CEMRL(TD3):
                       elitism=self.elitism)
 
     def learn(self, total_timesteps, callback=None, log_interval=4,
-              eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", reset_num_timesteps=True):
+              eval_env=None, eval_freq=-1, n_eval_episodes=5,
+              tb_log_name="CEMRL", eval_log_path=None, reset_num_timesteps=True):
 
-        timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq,
+                                                       n_eval_episodes, eval_log_path, reset_num_timesteps)
         actor_steps = 0
+        continue_training = True
+
+        callback.on_training_start(locals(), globals())
 
         while self.num_timesteps < total_timesteps:
 
             self.fitnesses = []
             self.es_params = self.es.ask(self.pop_size)
 
-            if callback is not None:
-                # Only stop training if return value is False, not when it is None.
-                if callback(locals(), globals()) is False:
-                    break
-
             if self.num_timesteps > 0:
                 # self.train(episode_timesteps)
                 # Gradient steps for half of the population
@@ -157,46 +157,35 @@ class CEMRL(TD3):
                     # Get the params back in the population
                     self.es_params[i] = self.actor.parameters_to_vector()
 
-            # Evaluate agent
-            if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-                timesteps_since_eval %= eval_freq
-
-                self.actor.load_from_vector(self.es.mu)
-                sync_envs_normalization(self.env, eval_env)
-
-                mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes)
-                evaluations.append(mean_reward)
-
-                if self.verbose > 0:
-                    print("Eval num_timesteps={}, "
-                          "episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
-                    print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
-
             actor_steps = 0
             # evaluate all actors
             for params in self.es_params:
-
                 self.actor.load_from_vector(params)
 
                 rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout,
                                                 n_steps=-1, action_noise=self.action_noise,
-                                                deterministic=False, callback=None,
+                                                deterministic=False, callback=callback,
                                                 learning_starts=self.learning_starts,
-                                                num_timesteps=self.num_timesteps,
                                                 replay_buffer=self.replay_buffer,
                                                 obs=obs, episode_num=episode_num,
                                                 log_interval=log_interval)
 
                 # Unpack
-                episode_reward, episode_timesteps, n_episodes, obs = rollout
+                episode_reward, episode_timesteps, n_episodes, obs, continue_training = rollout
+
+                if continue_training is False:
+                    break
 
                 episode_num += n_episodes
-                self.num_timesteps += episode_timesteps
-                timesteps_since_eval += episode_timesteps
                 actor_steps += episode_timesteps
                 self.fitnesses.append(episode_reward)
 
+            if continue_training is False:
+                break
+
             self._update_current_progress(self.num_timesteps, total_timesteps)
             self.es.tell(self.es_params, self.fitnesses)
-            timesteps_since_eval += actor_steps
+
+        callback.on_training_end()
+
         return self

torchy_baselines/common/base_class.py

@@ -2,8 +2,7 @@ import time
 import os
 import io
 import zipfile
-import typing
+from typing import Union, Type, Optional, Dict, Any, List, Tuple, Callable
-from typing import Union, Type, Optional, Dict, Any, List, Tuple
 from abc import ABC, abstractmethod
 from collections import deque
 
@@ -14,14 +13,12 @@ import numpy as np
 from torchy_baselines.common import logger
 from torchy_baselines.common.policies import BasePolicy, get_policy_from_name
 from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
-from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, sync_envs_normalization
+from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize
 from torchy_baselines.common.monitor import Monitor
-from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr
-
+from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict
-# TODO: define aliases, ex GymEnv = Union[gym.Env, VecEnv]
+from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback
-if typing.TYPE_CHECKING:
+from torchy_baselines.common.noise import ActionNoise
-    from torchy_baselines.common.noise import ActionNoise
 
 
 class BaseRLModel(ABC):
@@ -52,7 +49,7 @@ class BaseRLModel(ABC):
 
     def __init__(self,
                  policy: Type[BasePolicy],
-                 env: Union[gym.Env, VecEnv, str],
+                 env: Union[GymEnv, str],
                  policy_base: Type[BasePolicy],
                  policy_kwargs: Dict[str, Any] = None,
                  verbose: int = 0,
@@ -76,7 +73,7 @@ class BaseRLModel(ABC):
         if verbose > 0:
             print(f"Using {self.device} device")
 
-        self.env = None  # type: Union[gym.Env, VecEnv]
+        self.env = None  # type: GymEnv
         # get VecNormalize object if needed
         self._vec_normalize_env = unwrap_vec_normalize(env)
         self.verbose = verbose
@@ -134,16 +131,12 @@ class BaseRLModel(ABC):
     def _setup_model(self) -> None:
         """
         Setup model so state_dict can be loaded
-
         """
         raise NotImplementedError()
 
-    def _get_eval_env(self, eval_env: Union[gym.Env, VecEnv, None]) -> Union[gym.Env, VecEnv, None]:
+    def _get_eval_env(self, eval_env: Optional[GymEnv]) -> Optional[GymEnv]:
         """
         Return the environment that will be used for evaluation.
-
-        :param eval_env:
-        :return:
         """
         if eval_env is None:
             eval_env = self.eval_env
@@ -159,8 +152,7 @@ class BaseRLModel(ABC):
         Rescale the action from [low, high] to [-1, 1]
         (no need for symmetric action space)
 
-        :param action:
+        :param action: Action to scale
-        :return:
         """
         low, high = self.action_space.low, self.action_space.high
         return 2.0 * ((action - low) / (high - low)) - 1.0
@@ -170,8 +162,7 @@ class BaseRLModel(ABC):
         Rescale the action from [-1, 1] to [low, high]
         (no need for symmetric action space)
 
-        :param scaled_action:
+        :param scaled_action: Action to un-scale
-        :return:
         """
         low, high = self.action_space.low, self.action_space.high
         return low + (0.5 * (scaled_action + 1.0) * (high - low))
@@ -215,7 +206,7 @@ class BaseRLModel(ABC):
         """
         return np.nan if len(arr) == 0 else np.mean(arr)
 
-    def get_env(self) -> Union[VecEnv, None]:
+    def get_env(self) -> Optional[VecEnv]:
         """
         Returns the current environment (can be None if not defined).
 
@@ -230,7 +221,10 @@ class BaseRLModel(ABC):
         Checked parameters:
         - observation_space
         - action_space
-        :return: True if environment seems to be coherent
+
+        :param observation_space: (gym.spaces.Space)
+        :param action_space: (gym.spaces.Space)
+        :return: (bool) True if environment seems to be coherent
         """
         if observation_space != env.observation_space:
             return False
@@ -239,7 +233,7 @@ class BaseRLModel(ABC):
         # return true if no check failed
         return True
 
-    def set_env(self, env: Union[gym.Env, VecEnv]) -> None:
+    def set_env(self, env: GymEnv) -> None:
         """
         Checks the validity of the environment, and if it is coherent, set it as the current environment.
         Furthermore wrap any non vectorized env into a vectorized
@@ -276,11 +270,13 @@ class BaseRLModel(ABC):
 
     @abstractmethod
     def learn(self, total_timesteps: int,
-              callback=None, log_interval: int = 100,
+              callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None,
+              log_interval: int = 100,
               tb_log_name: str = "run",
-              eval_env: Union[gym.Env, VecEnv, None] = None,
+              eval_env: Optional[GymEnv] = None,
               eval_freq: int = -1,
               n_eval_episodes: int = 5,
+              eval_log_path: Optional[str] = None,
               reset_num_timesteps: bool = True):
         """
         Return a trained model.
@@ -294,6 +290,8 @@ class BaseRLModel(ABC):
         :param eval_env: (gym.Env) Environment that will be used to evaluate the agent
         :param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
         :param n_eval_episodes: (int) Number of episode to evaluate the agent
+        :param eval_log_path: (Optional[str]) Path to a folder where the evaluations will be saved
+        :param reset_num_timesteps: (bool)
         :return: (BaseRLModel) the trained model
         """
         raise NotImplementedError()
@@ -315,7 +313,7 @@ class BaseRLModel(ABC):
         raise NotImplementedError()
 
     @classmethod
-    def load(cls, load_path: str, env: Union[gym.Env, VecEnv, None] = None, **kwargs):
+    def load(cls, load_path: str, env: Optional[GymEnv] = None, **kwargs):
         """
         Load the model from a zip-file
 
@@ -364,7 +362,9 @@ class BaseRLModel(ABC):
         return model
 
     @staticmethod
-    def _load_from_file(load_path: str, load_data: bool = True):
+    def _load_from_file(load_path: str, load_data: bool = True) -> (Tuple[Optional[Dict[str, Any]],
+                                                                          Optional[TensorDict],
+                                                                          Optional[TensorDict]]):
         """ Load model data from a .zip archive
 
         :param load_path: Where to load the model from
@@ -467,12 +467,52 @@ class BaseRLModel(ABC):
         if self.eval_env is not None:
             self.eval_env.seed(seed)
 
-    def _setup_learn(self, eval_env):
+    def _init_callback(self,
+                       callback: Union[None, Callable, List[BaseCallback], BaseCallback],
+                       eval_env: Optional[VecEnv] = None,
+                       eval_freq: int = 10000,
+                       n_eval_episodes: int = 5,
+                       log_path: Optional[str] = None) -> BaseCallback:
+        """
+        :param callback: (Union[callable, [BaseCallback], BaseCallback, None])
+        :return: (BaseCallback)
+        """
+        # Convert a list of callbacks into a callback
+        if isinstance(callback, list):
+            callback = CallbackList(callback)
+
+        # Convert functional callback to object
+        if not isinstance(callback, BaseCallback):
+            callback = ConvertCallback(callback)
+
+        # Create eval callback in charge of the evaluation
+        if eval_env is not None:
+            eval_callback = EvalCallback(eval_env,
+                                         best_model_save_path=log_path,
+                                         log_path=log_path, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes)
+            callback = CallbackList([callback, eval_callback])
+
+        callback.init_callback(self)
+        return callback
+
+    def _setup_learn(self,
+                     eval_env: Optional[GymEnv],
+                     callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None,
+                     eval_freq: int = 10000,
+                     n_eval_episodes: int = 5,
+                     log_path: Optional[str] = None,
+                     reset_num_timesteps: bool = True,
+                     ) -> Tuple[int, np.ndarray, BaseCallback]:
         """
         Initialize different variables needed for training.
 
-        :param eval_env: (gym.Env or VecEnv)
+        :param eval_env: (Optional[GymEnv])
-        :return: (int, int, [float], np.ndarray, VecEnv)
+        :param callback: (Union[None, BaseCallback, List[BaseCallback, Callable]])
+        :param eval_freq: (int)
+        :param n_eval_episodes: (int)
+        :param log_path (Optional[str]): Path to a log folder
+        :param reset_num_timesteps: (bool) Whether to reset or not the `num_timesteps` attribute
+        :return: (Tuple[int, np.ndarray, BaseCallback])
         """
         self.start_time = time.time()
         self.ep_info_buffer = deque(maxlen=100)
@@ -481,16 +521,22 @@ class BaseRLModel(ABC):
             self.action_noise.reset()
 
         timesteps_since_eval, episode_num = 0, 0
-        evaluations = []
+
+        if reset_num_timesteps:
+            self.num_timesteps = 0
 
         if eval_env is not None and self.seed is not None:
             eval_env.seed(self.seed)
 
         eval_env = self._get_eval_env(eval_env)
-        obs = self.env.reset()  # type: Union[gym.Env, VecEnv]
+        obs = self.env.reset()
-        return timesteps_since_eval, episode_num, evaluations, obs, eval_env
 
-    def _update_info_buffer(self, infos):
+        # Create eval callback if needed
+        callback = self._init_callback(callback, eval_env, eval_freq, n_eval_episodes, log_path)
+
+        return episode_num, obs, callback
+
+    def _update_info_buffer(self, infos: List[Dict[str, Any]]) -> None:
         """
         Retrieve reward and episode length and update the buffer
         if using Monitor wrapper.
@@ -502,11 +548,18 @@ class BaseRLModel(ABC):
             if maybe_ep_info is not None:
                 self.ep_info_buffer.extend([maybe_ep_info])
 
-    def collect_rollouts(self, env, n_episodes=1, n_steps=-1, action_noise=None,
+    def collect_rollouts(self,
-                         deterministic=False, callback=None,
+                         env: VecEnv,
-                         learning_starts=0, num_timesteps=0,
+                         callback: 'BaseCallback',  # Type hint as string to avoid circular import
-                         replay_buffer=None, obs=None,
+                         n_episodes: int = 1,
-                         episode_num=0, log_interval=None):
+                         n_steps: int = -1,
+                         action_noise: Optional[ActionNoise] = None,
+                         deterministic: bool = False,
+                         learning_starts: int = 0,
+                         replay_buffer=None,
+                         obs: Optional[np.ndarray] = None,
+                         episode_num: int = 0,
+                         log_interval: Optional[int] = None) -> Tuple[float, int, int, Optional[np.ndarray], bool]:
         """
         Collect rollout using the current policy (and possibly fill the replay buffer)
         TODO: move this method to off-policy base class.
@@ -516,9 +569,8 @@ class BaseRLModel(ABC):
         :param n_steps: (int)
         :param action_noise: (ActionNoise)
         :param deterministic: (bool)
-        :param callback: (callable)
+        :param callback: (BaseCallback)
         :param learning_starts: (int)
-        :param num_timesteps: (int)
         :param replay_buffer: (ReplayBuffer)
         :param obs: (np.ndarray)
         :param episode_num: (int)
@@ -541,6 +593,9 @@ class BaseRLModel(ABC):
             if self.on_policy_exploration:
                 self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones', 'values']}
 
+        callback.on_rollout_start()
+        continue_training = True
+
         while total_steps < n_steps or total_episodes < n_episodes:
             done = False
             # Reset environment: not needed for VecEnv
@@ -548,6 +603,12 @@ class BaseRLModel(ABC):
             episode_reward, episode_timesteps = 0.0, 0
 
             while not done:
+
+                # Only stop training if return value is False, not when it is None.
+                if callback() is False:
+                    continue_training = False
+                    return 0.0, total_steps, total_episodes, None, continue_training
+
                 if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
                     # Sample a new noise matrix
                     self.actor.reset_noise()
@@ -555,7 +616,7 @@ class BaseRLModel(ABC):
                 # Select action randomly or according to policy
                 # TODO: use action from policy when using SDE during the warmup phase?
                 # if num_timesteps < learning_starts and not self.use_sde:
-                if num_timesteps < learning_starts:
+                if self.num_timesteps < learning_starts:
                     # Warmup phase
                     unscaled_action = np.array([self.action_space.sample()])
                 else:
@@ -614,7 +675,7 @@ class BaseRLModel(ABC):
                 if self._vec_normalize_env is not None:
                     obs_ = new_obs_
 
-                num_timesteps += 1
+                self.num_timesteps += 1
                 episode_timesteps += 1
                 total_steps += 1
                 if 0 < n_steps <= total_steps:
@@ -630,8 +691,8 @@ class BaseRLModel(ABC):
 
                 # Display training infos
                 if self.verbose >= 1 and log_interval is not None and (
-                            episode_num + total_episodes) % log_interval == 0:
+                        episode_num + total_episodes) % log_interval == 0:
-                    fps = int(num_timesteps / (time.time() - self.start_time))
+                    fps = int(self.num_timesteps / (time.time() - self.start_time))
                     logger.logkv("episodes", episode_num + total_episodes)
                     if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
                         logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
@@ -639,7 +700,7 @@ class BaseRLModel(ABC):
                     # logger.logkv("n_updates", n_updates)
                     logger.logkv("fps", fps)
                     logger.logkv('time_elapsed', int(time.time() - self.start_time))
-                    logger.logkv("total timesteps", num_timesteps)
+                    logger.logkv("total timesteps", self.num_timesteps)
                     if self.use_sde:
                         logger.logkv("std", (self.actor.get_std()).mean().item())
                     logger.dumpkvs()
@@ -667,7 +728,9 @@ class BaseRLModel(ABC):
                 self.rollout_data['returns'][step] = last_return
             self.rollout_data['advantage'] = self.rollout_data['returns'] - self.rollout_data['values']
 
-        return mean_reward, total_steps, total_episodes, obs
+        callback.on_rollout_end()
+
+        return mean_reward, total_steps, total_episodes, obs, continue_training
 
     @staticmethod
     def _save_to_file_zip(save_path: str, data: Dict[str, Any] = None,
@@ -767,27 +830,3 @@ class BaseRLModel(ABC):
             params_to_save[name] = attr.state_dict()
 
         self._save_to_file_zip(path, data=data, params=params_to_save, tensors=tensors)
-
-    def _eval_policy(self, eval_freq: int, eval_env: int, n_eval_episodes: int,
-                     timesteps_since_eval: int, render: bool = False, deterministic: bool = True) -> int:
-        """
-        Evaluate the current policy on a test environment.
-
-        :param eval_freq: Evaluate the agent every `eval_freq` timesteps (this may vary a little)
-        :param n_eval_episodes: Number of episode to evaluate the agent
-        :parma timesteps_since_eval: Number of timesteps since last evaluation
-        :param deterministic: Whether to use deterministic or stochastic actions
-        :param render: Whether to render the eval env or not
-        :return: Number of timesteps since last evaluation
-        """
-        if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
-            timesteps_since_eval %= eval_freq
-            # Synchronise the normalization stats if needed
-            sync_envs_normalization(self.env, eval_env)
-            mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes,
-                                                      render=render, deterministic=deterministic)
-            if self.verbose > 0:
-                print(f"Eval num_timesteps={self.num_timesteps}, "
-                      f"episode_reward={mean_reward:.2f} +/- {std_reward:.2f}")
-                print(f"FPS: {self.num_timesteps / (time.time() - self.start_time):.2f}")
-        return timesteps_since_eval

torchy_baselines/common/callbacks.py

@@ -0,0 +1,345 @@
+import os
+from abc import ABC, abstractmethod
+import typing
+from typing import Union, List, Dict, Any, Optional
+
+import gym
+import numpy as np
+
+from torchy_baselines.common.vec_env import VecEnv, sync_envs_normalization
+from torchy_baselines.common.evaluation import evaluate_policy
+from torchy_baselines.common.logger import Logger
+
+if typing.TYPE_CHECKING:
+    from torchy_baselines.common.base_class import BaseRLModel  # pytype: disable=pyi-error
+
+
+class BaseCallback(ABC):
+    """
+    Base class for callback.
+
+    :param verbose: (int)
+    """
+    def __init__(self, verbose: int = 0):
+        super(BaseCallback, self).__init__()
+        self.model = None  # type: BaseRLModel
+        self.training_env = None  # type: Union[gym.Env, VecEnv, None]
+        self.n_calls = 0  # type: int
+        self.num_timesteps = 0  # type: int
+        self.verbose = verbose
+        self.locals = None  # type: Dict[str, Any]
+        self.globals = None  # type: Dict[str, Any]
+        self.logger = None  # type: Logger
+        # Sometimes, for event callback, it is useful
+        # to have access to the parent object
+        self.parent = None  # type: Optional[BaseCallback]
+
+    # Type hint as string to avoid circular import
+    def init_callback(self, model: 'BaseRLModel') -> None:
+        """
+        Initialize the callback by saving references to the
+        RL model and the training environment for convenience.
+        """
+        self.model = model
+        self.training_env = model.get_env()
+        self.logger = Logger.CURRENT
+        self._init_callback()
+
+    def _init_callback(self) -> None:
+        pass
+
+    def on_training_start(self, locals_: Dict[str, Any], globals_: Dict[str, Any]) -> None:
+        # Those are reference and will be updated automatically
+        self.locals = locals_
+        self.globals = globals_
+        self._on_training_start()
+
+    def _on_training_start(self) -> None:
+        pass
+
+    def on_rollout_start(self) -> None:
+        self._on_rollout_start()
+
+    def _on_rollout_start(self) -> None:
+        pass
+
+    @abstractmethod
+    def _on_step(self) -> bool:
+        """
+        :return: (bool) If the callback returns False, training is aborted early.
+        """
+        return True
+
+    def __call__(self) -> bool:
+        """
+        This method will be called by the model. This is the equivalent to the callback function.
+        :return: (bool) If the callback returns False, training is aborted early.
+        """
+        self.n_calls += 1
+        # timesteps start at zero
+        self.num_timesteps = self.model.num_timesteps + 1
+
+        return self._on_step()
+
+    def on_training_end(self) -> None:
+        self._on_training_end()
+
+    def _on_training_end(self) -> None:
+        pass
+
+    def on_rollout_end(self) -> None:
+        self._on_rollout_end()
+
+    def _on_rollout_end(self) -> None:
+        pass
+
+
+class EventCallback(BaseCallback):
+    """
+    Base class for triggering callback on event.
+
+    :param callback: (Optional[BaseCallback]) Callback that will be called
+        when an event is triggered.
+    :param verbose: (int)
+    """
+    def __init__(self, callback: Optional[BaseCallback] = None, verbose: int = 0):
+        super(EventCallback, self).__init__(verbose=verbose)
+        self.callback = callback
+        # Give access to the parent
+        if callback is not None:
+            self.callback.parent = self
+
+    def init_callback(self, model: 'BaseRLModel') -> None:
+        super(EventCallback, self).init_callback(model)
+        if self.callback is not None:
+            self.callback.init_callback(self.model)
+
+    def _on_training_start(self) -> None:
+        if self.callback is not None:
+            self.callback.on_training_start(self.locals, self.globals)
+
+    def _on_event(self) -> bool:
+        if self.callback is not None:
+            return self.callback()
+        return True
+
+    def _on_step(self) -> bool:
+        return True
+
+
+class CallbackList(BaseCallback):
+    def __init__(self, callbacks: List[BaseCallback]):
+        super(CallbackList, self).__init__()
+        assert isinstance(callbacks, list)
+        self.callbacks = callbacks
+
+    def _init_callback(self) -> None:
+        for callback in self.callbacks:
+            callback.init_callback(self.model)
+
+    def _on_training_start(self) -> None:
+        for callback in self.callbacks:
+            callback.on_training_start(self.locals, self.globals)
+
+    def _on_step(self) -> bool:
+        continue_training = True
+        for callback in self.callbacks:
+            # # Update variables
+            # callback.num_timesteps = self.num_timesteps
+            # callback.n_calls = self.n_calls
+            # Return False (stop training) if at least one callback returns False
+            continue_training = callback() and continue_training
+        return continue_training
+
+    def _on_training_end(self) -> None:
+        for callback in self.callbacks:
+            callback.on_training_end()
+
+
+class CheckpointCallback(BaseCallback):
+    """
+    Callback for saving a model every `save_freq` steps
+
+    :param save_freq: (int)
+    :param save_path: (str) Path to the folder where the model will be saved.
+    :param name_prefix: (str) Common prefix to the saved models
+    """
+    def __init__(self, save_freq: int, save_path: str, name_prefix='rl_model', verbose=0):
+        super(CheckpointCallback, self).__init__(verbose)
+        self.save_freq = save_freq
+        self.save_path = save_path
+        self.name_prefix = name_prefix
+
+    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.save_freq == 0:
+            path = os.path.join(self.save_path, f'{self.name_prefix}_{self.num_timesteps}_steps')
+            self.model.save(path)
+            if self.verbose > 1:
+                print(f"Saving model checkpoint to {path}")
+        return True
+
+
+class ConvertCallback(BaseCallback):
+    """
+    Convert functional callback (old-style) to object.
+
+    :param callback: (callable)
+    :param verbose: (int)
+    """
+    def __init__(self, callback, verbose=0):
+        super(ConvertCallback, self).__init__(verbose)
+        self.callback = callback
+
+    def _on_step(self) -> bool:
+        if self.callback is not None:
+            return self.callback(self.locals, self.globals)
+        return True
+
+
+class EvalCallback(EventCallback):
+    """
+    Callback for evaluating an agent.
+
+    :param eval_env: (Union[gym.Env, VecEnv]) The environment used for initialization
+    :param callback_on_new_best: (Optional[BaseCallback]) Callback to trigger
+        when there is a new best model according to the `mean_reward`
+    :param n_eval_episodes: (int) The number of episodes to test the agent
+    :param eval_freq: (int) Evaluate the agent every eval_freq call of the callback.
+    :param log_path: (str) Path to a folder where the evaluations (`evaluations.npz`)
+        will be saved. It will be updated at each evaluation.
+    :param best_model_save_path: (str) Path to a folder where the best model
+        according to performance on the eval env will be saved.
+    :param deterministic: (bool) Whether the evaluation should
+        use a stochastic or deterministic actions.
+    :param deterministic: (bool) Whether to render or not the environment during evaluation
+    :param verbose: (int)
+    """
+    def __init__(self, eval_env: Union[gym.Env, VecEnv],
+                 callback_on_new_best: Optional[BaseCallback] = None,
+                 n_eval_episodes: int = 5,
+                 eval_freq: int = 10000,
+                 log_path: str = None,
+                 best_model_save_path: str = None,
+                 deterministic: bool = True,
+                 render: bool = False,
+                 verbose: int = 1):
+        super(EvalCallback, self).__init__(callback_on_new_best, verbose=verbose)
+        self.n_eval_episodes = n_eval_episodes
+        self.eval_freq = eval_freq
+        self.best_mean_reward = -np.inf
+        self.deterministic = deterministic
+        self.render = render
+
+        if isinstance(eval_env, VecEnv):
+            assert eval_env.num_envs == 1, "You must pass only one environment for evaluation"
+
+        self.eval_env = eval_env
+        self.best_model_save_path = best_model_save_path
+        # Logs will be written in `evaluations.npz`
+        if log_path is not None:
+            os.path.join(log_path, 'evaluations')
+        self.log_path = log_path
+        self.evaluations_results = []
+        self.evaluations_timesteps = []
+        self.evaluations_length = []
+
+    def _init_callback(self):
+        # Does not work when eval_env is a gym.Env and training_env is a VecEnv
+        # assert type(self.training_env) is type(self.eval_env), ("training and eval env are not of the same type",
+        #                                                         "{} != {}".format(self.training_env, self.eval_env))
+
+        # Create folders if needed
+        if self.best_model_save_path is not None:
+            os.makedirs(self.best_model_save_path, exist_ok=True)
+        if self.log_path is not None:
+            os.makedirs(os.path.dirname(self.log_path), exist_ok=True)
+
+    def _on_step(self) -> bool:
+
+        if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0:
+            # Sync training and eval env if there is VecNormalize
+            sync_envs_normalization(self.training_env, self.eval_env)
+
+            episode_rewards, episode_lengths = evaluate_policy(self.model, self.eval_env,
+                                                               n_eval_episodes=self.n_eval_episodes,
+                                                               render=self.render,
+                                                               deterministic=self.deterministic,
+                                                               return_episode_rewards=True)
+
+            if self.log_path is not None:
+                self.evaluations_timesteps.append(self.num_timesteps)
+                self.evaluations_results.append(episode_rewards)
+                self.evaluations_length.append(episode_lengths)
+                np.savez(self.log_path, timesteps=self.evaluations_timesteps,
+                         results=self.evaluations_results, ep_lengths=self.evaluations_length)
+
+            mean_reward, std_reward = np.mean(episode_rewards), np.std(episode_rewards)
+            mean_ep_length, std_ep_length = np.mean(episode_lengths), np.std(episode_lengths)
+
+            if self.verbose > 0:
+                print(f"Eval num_timesteps={self.num_timesteps}, "
+                      f"episode_reward={mean_reward:.2f} +/- {std_reward:.2f}")
+                print(f"Episode length: {mean_ep_length:.2f} +/- {std_ep_length:.2f}")
+
+            if mean_reward > self.best_mean_reward:
+                if self.verbose > 0:
+                    print("New best mean reward!")
+                if self.best_model_save_path is not None:
+                    self.model.save(os.path.join(self.best_model_save_path, 'best_model'))
+                self.best_mean_reward = mean_reward
+                # Trigger callback if needed
+                if self.callback is not None:
+                    return self._on_event()
+
+        return True
+
+
+class StopTrainingOnRewardThreshold(BaseCallback):
+    """
+    Stop the training once a threshold in episodic reward
+    has been reached (i.e. when the model is good enough).
+
+    It must be used with the `EvalCallback`.
+
+    :param reward_threshold: (float)  Minimum expected reward per episode
+        to stop training.
+    :param verbose: (int)
+    """
+    def __init__(self, reward_threshold: float, verbose: int = 0):
+        super(StopTrainingOnRewardThreshold, self).__init__(verbose=verbose)
+        self.reward_threshold = reward_threshold
+
+    def _on_step(self) -> bool:
+        assert self.parent is not None, ("`StopTrainingOnMinimumReward` callback must be used "
+                                         "with an `EvalCallback`")
+        # Convert np.bool to bool, otherwise callback() is False won't work
+        continue_training = bool(self.parent.best_mean_reward < self.reward_threshold)
+        if self.verbose > 0 and not continue_training:
+            print(f"Stopping training because the mean reward {self.parent.best_mean_reward:.2f} "
+                  f" is above the threshold {self.reward_threshold}")
+        return continue_training
+
+
+class EveryNTimesteps(EventCallback):
+    """
+    Trigger a callback every `n_steps` timesteps
+
+    :param n_steps: (int) Number of timesteps between two trigger.
+    :param callback: (BaseCallback) Callback that will be called
+        when the event is triggered.
+    """
+    def __init__(self, n_steps: int, callback: BaseCallback):
+        super(EveryNTimesteps, self).__init__(callback)
+        self.n_steps = n_steps
+        self.last_time_trigger = 0
+
+    def _on_step(self) -> bool:
+        if (self.num_timesteps - self.last_time_trigger) >= self.n_steps:
+            self.last_time_trigger = self.num_timesteps
+            return self._on_event()
+        return True

torchy_baselines/common/evaluation.py

@@ -24,31 +24,33 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
     :param return_episode_rewards: (bool) If True, a list of reward per episode
         will be returned instead of the mean.
     :return: (float, float) Mean reward per episode, std of reward per episode
-        returns ([float], int) when `return_episode_rewards` is True
+        returns ([float], [int]) when `return_episode_rewards` is True
     """
     if isinstance(env, VecEnv):
         assert env.num_envs == 1, "You must pass only one environment when using this function"
 
-    episode_rewards, n_steps = [], 0
+    episode_rewards, episode_lengths = [], []
     for _ in range(n_eval_episodes):
         obs = env.reset()
         done = False
         episode_reward = 0.0
+        episode_length = 0
         while not done:
             action = model.predict(obs, deterministic=deterministic)
             obs, reward, done, _info = env.step(action)
             episode_reward += reward
             if callback is not None:
                 callback(locals(), globals())
-            n_steps += 1
+            episode_length += 1
             if render:
                 env.render()
         episode_rewards.append(episode_reward)
+        episode_lengths.append(episode_length)
     mean_reward = np.mean(episode_rewards)
     std_reward = np.std(episode_rewards)
     if reward_threshold is not None:
         assert mean_reward > reward_threshold, (f'Mean reward below threshold: '
                                                 '{mean_reward:.2f} < {reward_threshold:.2f}')
     if return_episode_rewards:
-        return episode_rewards, n_steps
+        return episode_rewards, episode_lengths
     return mean_reward, std_reward

torchy_baselines/common/noise.py

@@ -1,19 +1,26 @@
 """
 Taken from stable-baselines
 """
+from abc import ABC, abstractmethod
 import numpy as np
 
 
-class ActionNoise(object):
+class ActionNoise(ABC):
     """
     The action noise base class
     """
+    def __init__(self):
+        super(ActionNoise, self).__init__()
+
     def reset(self):
         """
         call end of episode reset for the noise
         """
         pass
 
+    @abstractmethod
+    def __call__(self):
+        pass
 
 class NormalActionNoise(ActionNoise):
     """
@@ -25,6 +32,7 @@ class NormalActionNoise(ActionNoise):
     def __init__(self, mean, sigma):
         self._mu = mean
         self._sigma = sigma
+        super(NormalActionNoise, self).__init__()
 
     def __call__(self):
         return np.random.normal(self._mu, self._sigma)
@@ -54,6 +62,7 @@ class OrnsteinUhlenbeckActionNoise(ActionNoise):
         self.initial_noise = initial_noise
         self.noise_prev = None
         self.reset()
+        super(OrnsteinUhlenbeckActionNoise, self).__init__()
 
     def __call__(self):
         noise = self.noise_prev + self._theta * (self._mu - self.noise_prev) * self._dt + \

torchy_baselines/common/type_aliases.py

@@ -0,0 +1,14 @@
+"""
+Common aliases for type hing
+"""
+from typing import Union, Type, Optional, Dict, Any, List, Tuple
+
+import torch
+import gym
+
+from torchy_baselines.common.vec_env import VecEnv
+
+
+GymEnv = Union[gym.Env, VecEnv]
+TensorDict = Dict[str, torch.Tensor]
+OptimizerStateDict = Dict[str, Any]

torchy_baselines/ppo/ppo.py

@@ -1,6 +1,6 @@
 import os
 import time
-from typing import List, Tuple
+from typing import Optional, Tuple, List
 
 import gym
 from gym import spaces
@@ -17,6 +17,8 @@ import numpy as np
 from torchy_baselines.common.base_class import BaseRLModel
 from torchy_baselines.common.buffers import RolloutBuffer
 from torchy_baselines.common.utils import explained_variance, get_schedule_fn
+from torchy_baselines.common.vec_env import VecEnv
+from torchy_baselines.common.callbacks import BaseCallback
 from torchy_baselines.common import logger
 from torchy_baselines.ppo.policies import PPOPolicy
 
@@ -151,17 +153,30 @@ class PPO(BaseRLModel):
             clipped_actions = np.clip(clipped_actions, self.action_space.low, self.action_space.high)
         return clipped_actions
 
-    def collect_rollouts(self, env, rollout_buffer, n_rollout_steps=256, callback=None,
+    def collect_rollouts(self,
-                         obs=None):
+                        env: VecEnv,
+                        callback: BaseCallback,
+                        rollout_buffer: RolloutBuffer,
+                        n_rollout_steps: int = 256,
+                        obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]:
 
         n_steps = 0
+        continue_training = True
         rollout_buffer.reset()
         # Sample new weights for the state dependent exploration
         # TODO: ensure episodic setting?
         if self.use_sde:
             self.policy.reset_noise(env.num_envs)
 
+        callback.on_rollout_start()
+
         while n_steps < n_rollout_steps:
+
+            if callback() is False:
+                continue_training = False
+                return None, continue_training
+
+
             if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
                 # Sample a new noise matrix
                 self.policy.reset_noise(env.num_envs)
@@ -179,6 +194,8 @@ class PPO(BaseRLModel):
 
             self._update_info_buffer(infos)
             n_steps += 1
+            self.num_timesteps += env.num_envs
+
             if isinstance(self.action_space, gym.spaces.Discrete):
                 # Reshape in case of discrete action
                 actions = actions.reshape(-1, 1)
@@ -187,7 +204,9 @@ class PPO(BaseRLModel):
 
         rollout_buffer.compute_returns_and_advantage(values, dones=dones)
 
-        return obs
+        callback.on_rollout_end()
+
+        return obs, continue_training
 
     def train(self, gradient_steps, batch_size=64):
         # Update optimizer learning rate
@@ -268,25 +287,29 @@ class PPO(BaseRLModel):
             logger.logkv("std", th.exp(self.policy.log_std).mean().item())
 
     def learn(self, total_timesteps, callback=None, log_interval=1,
-              eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", reset_num_timesteps=True):
+              eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO",
+              eval_log_path=None, reset_num_timesteps=True):
 
-        timesteps_since_eval, iteration, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq,
+                                                       n_eval_episodes, eval_log_path, reset_num_timesteps)
+        iteration = 0
 
         if self.tensorboard_log is not None and SummaryWriter is not None:
             self.tb_writer = SummaryWriter(log_dir=os.path.join(self.tensorboard_log, tb_log_name))
 
+        callback.on_training_start(locals(), globals())
+
         while self.num_timesteps < total_timesteps:
 
-            if callback is not None:
+            obs, continue_training = self.collect_rollouts(self.env, callback,
-                # Only stop training if return value is False, not when it is None.
+                                                           self.rollout_buffer,
-                if callback(locals(), globals()) is False:
+                                                           n_rollout_steps=self.n_steps,
-                    break
+                                                           obs=obs)
+
+            if continue_training is False:
+                break
 
-            obs = self.collect_rollouts(self.env, self.rollout_buffer, n_rollout_steps=self.n_steps,
-                                        obs=obs)
             iteration += 1
-            self.num_timesteps += self.n_steps * self.n_envs
-            timesteps_since_eval += self.n_steps * self.n_envs
             self._update_current_progress(self.num_timesteps, total_timesteps)
 
             # Display training infos
@@ -303,13 +326,12 @@ class PPO(BaseRLModel):
 
             self.train(self.n_epochs, batch_size=self.batch_size)
 
-            # Evaluate the agent
-            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
-                                                     timesteps_since_eval, deterministic=True)
             # For tensorboard integration
             # if self.tb_writer is not None:
             #     self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
 
+        callback.on_training_end()
+
         return self
 
     def get_torch_variables(self) -> Tuple[List[str], List[str]]:

torchy_baselines/sac/sac.py

@@ -259,31 +259,27 @@ class SAC(BaseRLModel):
 
     def learn(self, total_timesteps, callback=None, log_interval=4,
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC",
-              reset_num_timesteps=True):
+              eval_log_path=None, reset_num_timesteps=True):
 
-        timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq,
+                                                       n_eval_episodes, eval_log_path, reset_num_timesteps)
+        callback.on_training_start(locals(), globals())
 
         while self.num_timesteps < total_timesteps:
-
-            if callback is not None:
-                # Only stop training if return value is False, not when it is None.
-                if callback(locals(), globals()) is False:
-                    break
-
             rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout,
                                             n_steps=self.train_freq, action_noise=self.action_noise,
-                                            deterministic=False, callback=None,
+                                            deterministic=False, callback=callback,
                                             learning_starts=self.learning_starts,
-                                            num_timesteps=self.num_timesteps,
                                             replay_buffer=self.replay_buffer,
                                             obs=obs, episode_num=episode_num,
                                             log_interval=log_interval)
             # Unpack
-            episode_reward, episode_timesteps, n_episodes, obs = rollout
+            episode_reward, episode_timesteps, n_episodes, obs, continue_training = rollout
+
+            if continue_training is False:
+                break
 
-            self.num_timesteps += episode_timesteps
             episode_num += n_episodes
-            timesteps_since_eval += episode_timesteps
             self._update_current_progress(self.num_timesteps, total_timesteps)
 
             if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
@@ -291,9 +287,7 @@ class SAC(BaseRLModel):
 
                 self.train(gradient_steps, batch_size=self.batch_size)
 
-            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
+        callback.on_training_end()
-                                                     timesteps_since_eval, deterministic=True)
-
         return self
 
     def excluded_save_params(self) -> List[str]:

torchy_baselines/td3/policies.py

@@ -1,3 +1,6 @@
+from typing import List, Tuple, Callable, Optional
+
+import torch
 import torch as th
 import torch.nn as nn
 
@@ -27,9 +30,18 @@ class Actor(BaseNetwork):
         a positive standard deviation (cf paper). It allows to keep variance
         above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
     """
-    def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
+    def __init__(self,
-                 use_sde=False, log_std_init=-3, clip_noise=None,
+                 obs_dim: int,
-                 lr_sde=3e-4, full_std=False, sde_net_arch=None, use_expln=False):
+                 action_dim: int,
+                 net_arch: List[int],
+                 activation_fn: nn.Module = nn.ReLU,
+                 use_sde: bool = False,
+                 log_std_init: float = -3,
+                 clip_noise: Optional[float] = None,
+                 lr_sde: float = 3e-4,
+                 full_std: bool = False,
+                 sde_net_arch: Optional[List[int]] = None,
+                 use_expln: bool = False):
         super(Actor, self).__init__()
 
         self.latent_pi, self.log_std = None, None
@@ -65,7 +77,7 @@ class Actor(BaseNetwork):
             actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True)
             self.mu = nn.Sequential(*actor_net)
 
-    def get_std(self):
+    def get_std(self) -> torch.Tensor:
         """
         Retrieve the standard deviation of the action distribution.
         Only useful when using SDE.
@@ -81,7 +93,7 @@ class Actor(BaseNetwork):
         mean_actions = self.mu(latent_pi)
         return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde)
 
-    def _get_latent(self, obs):
+    def _get_latent(self, obs) -> Tuple[torch.Tensor, torch.Tensor]:
         latent_pi = self.latent_pi(obs)
 
         if self.sde_feature_extractor is not None:
@@ -90,7 +102,7 @@ class Actor(BaseNetwork):
             latent_sde = latent_pi
         return latent_pi, latent_sde
 
-    def evaluate_actions(self, obs, action):
+    def evaluate_actions(self, obs: torch.Tensor, action: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Evaluate actions according to the current policy,
         given the observations. Only useful when using SDE.
@@ -106,13 +118,13 @@ class Actor(BaseNetwork):
         # value = self.value_net(latent_vf)
         return log_prob, distribution.entropy()
 
-    def reset_noise(self):
+    def reset_noise(self) -> None:
         """
         Sample new weights for the exploration matrix, when using SDE.
         """
         self.action_dist.sample_weights(self.log_std)
 
-    def forward(self, obs, deterministic=True):
+    def forward(self, obs: torch.Tensor, deterministic: bool = True) -> torch.Tensor:
         if self.use_sde:
             latent_pi, latent_sde = self._get_latent(obs)
             if deterministic:
@@ -141,8 +153,8 @@ class Critic(BaseNetwork):
     :param net_arch: ([int]) Network architecture
     :param activation_fn: (nn.Module) Activation function
     """
-    def __init__(self, obs_dim, action_dim,
+    def __init__(self, obs_dim: int, action_dim: int,
-                 net_arch, activation_fn=nn.ReLU):
+                 net_arch: List[int], activation_fn: nn.Module = nn.ReLU):
         super(Critic, self).__init__()
 
         q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
@@ -151,14 +163,12 @@ class Critic(BaseNetwork):
         q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
         self.q2_net = nn.Sequential(*q2_net)
 
-        self.q_networks = [self.q1_net, self.q2_net]
+    def forward(self, obs: torch.Tensor, action: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-
-    def forward(self, obs, action):
         qvalue_input = th.cat([obs, action], dim=1)
-        return [q_net(qvalue_input) for q_net in self.q_networks]
+        return self.q1_net(qvalue_input), self.q2_net(qvalue_input)
 
-    def q1_forward(self, obs, action):
+    def q1_forward(self, obs: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
-        return self.q_networks[0](th.cat([obs, action], dim=1))
+        return self.q1_net(th.cat([obs, action], dim=1))
 
 
 class ValueFunction(BaseNetwork):

torchy_baselines/td3/td3.py

@@ -251,31 +251,30 @@ class TD3(BaseRLModel):
         del self.rollout_data
 
     def learn(self, total_timesteps, callback=None, log_interval=4,
-              eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True):
+              eval_env=None, eval_freq=-1, n_eval_episodes=5,
+              tb_log_name="TD3", eval_log_path=None, reset_num_timesteps=True):
 
-        timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq,
+                                                       n_eval_episodes, eval_log_path, reset_num_timesteps)
+
+        callback.on_training_start(locals(), globals())
 
         while self.num_timesteps < total_timesteps:
 
-            if callback is not None:
-                # Only stop training if return value is False, not when it is None.
-                if callback(locals(), globals()) is False:
-                    break
-
             rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout,
                                             n_steps=self.train_freq, action_noise=self.action_noise,
-                                            deterministic=False, callback=None,
+                                            deterministic=False, callback=callback,
                                             learning_starts=self.learning_starts,
-                                            num_timesteps=self.num_timesteps,
                                             replay_buffer=self.replay_buffer,
                                             obs=obs, episode_num=episode_num,
                                             log_interval=log_interval)
             # Unpack
-            episode_reward, episode_timesteps, n_episodes, obs = rollout
+            episode_reward, episode_timesteps, n_episodes, obs, continue_training = rollout
+
+            if continue_training is False:
+                break
 
             episode_num += n_episodes
-            self.num_timesteps += episode_timesteps
-            timesteps_since_eval += episode_timesteps
             self._update_current_progress(self.num_timesteps, total_timesteps)
 
             if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
@@ -292,9 +291,7 @@ class TD3(BaseRLModel):
                 gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
                 self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay)
 
-            # Evaluate the agent
+        callback.on_training_end()
-            timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
-                                                     timesteps_since_eval, deterministic=True)
 
         return self