Add callback support

.coveragerc

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

docs/misc/changelog.rst

@@ -13,6 +13,7 @@ Breaking Changes:
 New Features:
 ^^^^^^^^^^^^^
 - Add `seed()` method to `VecEnv` class
+- Add support for Callback (cf https://github.com/hill-a/stable-baselines/pull/644)
 
 Bug Fixes:
 ^^^^^^^^^^
@@ -24,6 +25,7 @@ Others:
 ^^^^^^^
 - Add type check
 - Converted all format string to f-strings
+- Add test for `OrnsteinUhlenbeckActionNoise`
 
 Documentation:
 ^^^^^^^^^^^^^^

tests/test_callbacks.py

@@ -0,0 +1,37 @@
+import pytest
+
+from torchy_baselines import SAC
+from torchy_baselines.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback,
+    EveryNTimesteps, StopTrainingOnRewardThreshold)
+
+
+@pytest.mark.parametrize("model_class", [SAC])
+def test_callbacks(model_class):
+    # Create RL model
+    model = model_class('MlpPolicy', 'Pendulum-v0')
+
+    checkpoint_callback = CheckpointCallback(save_freq=1000, save_path='./logs/')
+
+    # For testing: use the same training env
+    eval_env = model.get_env()
+    # 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='./logs/best_model',
+                                 log_path='./logs/results', eval_freq=100)
+
+    # Equivalent to the `checkpoint_callback`
+    # but here in an event-driven manner
+    checkpoint_on_event = CheckpointCallback(save_freq=1, save_path='./logs/',
+                                             name_prefix='event')
+    event_callback = EveryNTimesteps(n_steps=1000, callback=checkpoint_on_event)
+
+    callback = CallbackList([checkpoint_callback, eval_callback, event_callback])
+
+    model.learn(1000, 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)

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)

torchy_baselines/cem_rl/cem_rl.py

@@ -102,19 +102,17 @@ class CEMRL(TD3):
     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):
 
-        timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback)
         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
@@ -180,7 +178,7 @@ class CEMRL(TD3):
 
                 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,
@@ -188,7 +186,10 @@ class CEMRL(TD3):
                                                 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
@@ -196,7 +197,13 @@ class CEMRL(TD3):
                 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

@@ -18,10 +18,11 @@ from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_norm
 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
+from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict
+from torchy_baselines.common.noise import ActionNoise
 
-# TODO: define aliases, ex GymEnv = Union[gym.Env, VecEnv]
 if typing.TYPE_CHECKING:
-    from torchy_baselines.common.noise import ActionNoise
+    from torchy_baselines.common.callbacks import BaseCallback
 
 
 class BaseRLModel(ABC):
@@ -51,7 +52,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,
@@ -75,7 +76,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
@@ -129,7 +130,7 @@ class BaseRLModel(ABC):
                 raise ValueError("Error: the model does not support multiple envs requires a single vectorized"
                                  " environment.")
 
-    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.
 
@@ -145,6 +146,27 @@ class BaseRLModel(ABC):
             assert eval_env.num_envs == 1
         return eval_env
 
+    # Type hint as string to avoid circular import
+    def _init_callback(self, callback) -> 'BaseCallback':
+        """
+        Note: we cannot use type hint here because of circular import.
+
+        :param callback: (Union[callable, [BaseCallback], BaseCallback, None])
+        :return: (BaseCallback)
+        """
+        # Avoid circular import
+        from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback
+
+        # 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)
+
+        callback.init_callback(self)
+        return callback
+
     def scale_action(self, action: np.ndarray) -> np.ndarray:
         """
         Rescale the action from [low, high] to [-1, 1]
@@ -206,7 +228,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 +252,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
@@ -252,7 +274,7 @@ class BaseRLModel(ABC):
         self.n_envs = env.num_envs
         self.env = env
 
-    def get_parameters(self) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    def get_parameters(self) -> Tuple[TensorDict, OptimizerStateDict]:
         """
         Returns policy and optimizer parameters as a tuple
 
@@ -260,7 +282,7 @@ class BaseRLModel(ABC):
         """
         return self.get_policy_parameters(), self.get_opt_parameters()
 
-    def get_policy_parameters(self) -> Dict[str, Any]:
+    def get_policy_parameters(self) -> TensorDict:
         """
         Get current model policy parameters as dictionary of variable name -> tensors.
 
@@ -269,7 +291,7 @@ class BaseRLModel(ABC):
         return self.policy.state_dict()
 
     @abstractmethod
-    def get_opt_parameters(self)-> Dict[str, Any]:
+    def get_opt_parameters(self)-> OptimizerStateDict:
         """
         Get current model optimizer parameters as dictionary of variable names -> tensors
         :return: (dict) Dictionary of variable name -> tensor of model's optimizer parameters
@@ -280,7 +302,7 @@ class BaseRLModel(ABC):
     def learn(self, total_timesteps: int,
               callback=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,
               reset_num_timesteps: bool = True):
@@ -316,7 +338,7 @@ class BaseRLModel(ABC):
         """
         raise NotImplementedError()
 
-    def load_parameters(self, load_dict: Dict[str, Any], opt_params: Dict[str, Any]) -> None:
+    def load_parameters(self, load_dict: TensorDict, opt_params: OptimizerStateDict) -> None:
         """
         Load model parameters from a dictionary
         load_dict should contain all keys from torch.model.state_dict()
@@ -325,14 +347,14 @@ class BaseRLModel(ABC):
 
 
         :param load_dict: dict of parameters from model.state_dict()
-        :param opt_params: dict of optimizer state_dicts should be handled in child_class
+        :param opt_params: dict of optimizer state_dicts should be handled in child class
         """
         if opt_params is not None:
             raise ValueError("Optimizer Parameters where given but no overloaded load function exists for this class")
         self.policy.load_state_dict(load_dict)
 
     @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
 
@@ -368,7 +390,8 @@ 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[OptimizerStateDict]]):
         """ Load model data from a .zip archive
 
         :param load_path: Where to load the model from
@@ -450,12 +473,14 @@ class BaseRLModel(ABC):
         if self.eval_env is not None:
             self.eval_env.seed(seed)
 
-    def _setup_learn(self, eval_env):
+    def _setup_learn(self, eval_env: Optional[GymEnv], callback=None) -> (Tuple[int, int,
+            List[Any], np.ndarray, Optional[VecEnv], Any]):
         """
         Initialize different variables needed for training.
 
-        :param eval_env: (gym.Env or VecEnv)
-        :return: (int, int, [float], np.ndarray, VecEnv)
+        :param eval_env: (Optional[GymEnv])
+        :param callback: (Union[None, BaseCallback, List[BaseCallback, Callable]])
+        :return: (int, int, [float], np.ndarray, VecEnv, BaseCallback)
         """
         self.start_time = time.time()
         self.ep_info_buffer = deque(maxlen=100)
@@ -463,6 +488,8 @@ class BaseRLModel(ABC):
         if self.action_noise is not None:
             self.action_noise.reset()
 
+        callback = self._init_callback(callback)
+
         timesteps_since_eval, episode_num = 0, 0
         evaluations = []
 
@@ -470,10 +497,11 @@ class BaseRLModel(ABC):
             eval_env.seed(self.seed)
 
         eval_env = self._get_eval_env(eval_env)
-        obs = self.env.reset()  # type: Union[gym.Env, VecEnv]
-        return timesteps_since_eval, episode_num, evaluations, obs, eval_env
+        obs = self.env.reset()  # type: GymEnv
 
-    def _update_info_buffer(self, infos):
+        return timesteps_since_eval, episode_num, evaluations, obs, eval_env, 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.
@@ -485,11 +513,19 @@ 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,
-                         deterministic=False, callback=None,
-                         learning_starts=0, num_timesteps=0,
-                         replay_buffer=None, obs=None,
-                         episode_num=0, log_interval=None):
+    def collect_rollouts(self,
+                         env: VecEnv,
+                         callback: 'BaseCallback',  # Type hint as string to avoid circular import
+                         n_episodes: int = 1,
+                         n_steps: int = -1,
+                         action_noise: Optional[ActionNoise] = None,
+                         deterministic: bool = False,
+                         learning_starts: int = 0,
+                         num_timesteps: 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.
@@ -499,7 +535,7 @@ 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)
@@ -524,6 +560,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
@@ -531,6 +570,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()
@@ -650,11 +695,13 @@ 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,
-                          params:  Dict[str, Any] = None, opt_params:  Dict[str, Any] = None) -> None:
+                          params: TensorDict = None, opt_params:  OptimizerStateDict = None) -> None:
         """
         Save model to a zip archive.
 
@@ -730,7 +777,7 @@ class BaseRLModel(ABC):
         opt_params_to_save = self.get_opt_parameters()
         self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
 
-    def _eval_policy(self, eval_freq: int, eval_env: int, n_eval_episodes: int,
+    def _eval_policy(self, eval_freq: int, eval_env: GymEnv, n_eval_episodes: int,
                      timesteps_since_eval: int, render: bool = False, deterministic: bool = True) -> int:
         """
         Evaluate the current policy on a test environment.

torchy_baselines/common/callbacks.py

@@ -0,0 +1,320 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Union, List, Dict, Any, Optional
+
+import gym
+import numpy as np
+
+from torchy_baselines.common.base_class import BaseRLModel # pytype: disable=pyi-error
+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
+
+
+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]
+
+    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)
+        self.callback.init_callback(self.model)
+
+    def _on_training_start(self) -> 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
+
+
+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 on_step: (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 log file (.npz) where the evaluations
+        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 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,
+                 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
+        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
+        self.log_path = log_path
+        self.evaluations_results = []
+        self.evaluations_timesteps = []
+
+    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.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, _ = evaluate_policy(self.model, self.eval_env, n_eval_episodes=self.n_eval_episodes,
+                                                 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)
+                np.savez(self.log_path, timesteps=self.evaluations_timesteps, results=self.evaluations_results)
+
+            mean_reward, std_reward = np.mean(episode_rewards), np.std(episode_rewards)
+            if self.verbose > 0:
+                print(f"Eval num_timesteps={self.num_timesteps}, "
+                      f"episode_reward={mean_reward:.2f} +/- {std_reward:.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/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,5 +1,6 @@
 import os
 import time
+from typing import Optional, Tuple
 
 import gym
 from gym import spaces
@@ -16,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
 
@@ -149,17 +152,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,
-                         obs=None):
+    def collect_rollouts(self,
+                        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)
@@ -185,7 +201,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,20 +286,21 @@ class PPO(BaseRLModel):
     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):
 
-        timesteps_since_eval, iteration, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        timesteps_since_eval, iteration, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback)
 
         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:
-                # Only stop training if return value is False, not when it is None.
-                if callback(locals(), globals()) is False:
+            obs, continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps,
+                                                           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
@@ -308,6 +327,8 @@ class PPO(BaseRLModel):
             # 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_opt_parameters(self):

torchy_baselines/sac/sac.py

@@ -259,25 +259,24 @@ class SAC(BaseRLModel):
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC",
               reset_num_timesteps=True):
 
-        timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback)
+
+        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
@@ -292,6 +291,7 @@ class SAC(BaseRLModel):
             timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
                                                      timesteps_since_eval, deterministic=True)
 
+        callback.on_training_end()
         return self
 
     def get_opt_parameters(self):

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,
-                 use_sde=False, log_std_init=-3, clip_noise=None,
-                 lr_sde=3e-4, full_std=False, sde_net_arch=None, use_expln=False):
+    def __init__(self,
+                 obs_dim: int,
+                 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,
-                 net_arch, activation_fn=nn.ReLU):
+    def __init__(self, obs_dim: int, action_dim: int,
+                 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, action):
+    def forward(self, obs: torch.Tensor, action: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         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):
-        return self.q_networks[0](th.cat([obs, action], dim=1))
+    def q1_forward(self, obs: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
+        return self.q1_net(th.cat([obs, action], dim=1))
 
 
 class ValueFunction(BaseNetwork):

torchy_baselines/td3/td3.py

@@ -251,25 +251,25 @@ class TD3(BaseRLModel):
     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):
 
-        timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        timesteps_since_eval, episode_num, evaluations, obs, eval_env, callback = self._setup_learn(eval_env, callback)
+
+        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
@@ -294,6 +294,8 @@ class TD3(BaseRLModel):
             timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
                                                      timesteps_since_eval, deterministic=True)
 
+        callback.on_training_end()
+
         return self
 
     def get_opt_parameters(self):