Add StopTrainingOnMaxEpisodes to callback collection (#147)

* Add StopTrainingOnMaxEpisodes class to pre-made callback collection

* Adjust instant when counters are incremented for both OnPolicy and OffPolicy algorithms

* Improv to StopTrainingOnMaxEpisodes including output, tests and doc

* Improv StopTrainingOnMaxEpisodes callback running _init_callback

* Update callbacks.py

* Update test_callbacks.py

* Fix style

* Update changelog.rst

* Fix test

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
Co-authored-by: Antonin Raffin <antonin.raffin@dlr.de>

docs/guide/callbacks.rst

@@ -292,5 +292,35 @@ An :ref:`EventCallback` that will trigger its child callback every ``n_steps`` t
   model.learn(int(2e4), callback=event_callback)
 
 
+.. _StopTrainingOnMaxEpisodes:
+
+StopTrainingOnMaxEpisodes
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Stop the training upon reaching the maximum number of episodes, regardless of the model's ``total_timesteps`` value.
+Also, presumes that, for multiple environments, the desired behavior is that the agent trains on each env for ``max_episodes``
+and in total for ``max_episodes * n_envs`` episodes.
+
+
+.. note::
+    For multiple environments, the agent will train for a total of ``max_episodes * n_envs`` episodes.
+    However, it can't be guaranteed that this training will occur for an exact number of ``max_episodes`` per environment.
+    Thus, there is an assumption that, on average, each environment ran for ``max_episodes``.
+
+
+.. code-block:: python
+
+    from stable_baselines3 import A2C
+    from stable_baselines3.common.callbacks import StopTrainingOnMaxEpisodes
+
+    # Stops training when the model reaches the maximum number of episodes
+    callback_max_episodes = StopTrainingOnMaxEpisodes(max_episodes=5, verbose=1)
+
+    model = A2C('MlpPolicy', 'Pendulum-v0', verbose=1)
+    # Almost infinite number of timesteps, but the training will stop
+    # early as soon as the max number of episodes is reached
+    model.learn(int(1e10), callback=callback_max_episodes)
+
+
 .. automodule:: stable_baselines3.common.callbacks
   :members:

docs/misc/changelog.rst

@@ -13,6 +13,7 @@ Breaking Changes:
 New Features:
 ^^^^^^^^^^^^^
 - Added ``unwrap_vec_wrapper()`` to ``common.vec_env`` to extract ``VecEnvWrapper`` if needed
+- Added ``StopTrainingOnMaxEpisodes`` to callback collection (@xicocaio)
 
 Bug Fixes:
 ^^^^^^^^^^
@@ -29,6 +30,9 @@ Others:
 
 Documentation:
 ^^^^^^^^^^^^^^
+- Added ``StopTrainingOnMaxEpisodes`` details and example (@xicocaio)
+
+
 
 Pre-Release 0.8.0 (2020-08-03)
 ------------------------------
@@ -393,4 +397,4 @@ And all the contributors:
 @Miffyli @dwiel @miguelrass @qxcv @jaberkow @eavelardev @ruifeng96150 @pedrohbtp @srivatsankrishnan @evilsocket
 @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching
 @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3
-@tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag
+@tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio

stable_baselines3/common/callbacks.py

@@ -87,7 +87,7 @@ class BaseCallback(ABC):
         """
         self.n_calls += 1
         # timesteps start at zero
-        self.num_timesteps = self.model.num_timesteps + 1
+        self.num_timesteps = self.model.num_timesteps
 
         return self._on_step()
 
@@ -435,3 +435,48 @@ class EveryNTimesteps(EventCallback):
             self.last_time_trigger = self.num_timesteps
             return self._on_event()
         return True
+
+
+class StopTrainingOnMaxEpisodes(BaseCallback):
+    """
+    Stop the training once a maximum number of episodes are played.
+
+    For multiple environments presumes that, the desired behavior is that the agent trains on each env for ``max_episodes``
+    and in total for ``max_episodes * n_envs`` episodes.
+
+    :param max_episodes: (int) Maximum number of episodes to stop training.
+    :param verbose: (int) Select whether to print information about when training ended by reaching ``max_episodes``
+    """
+
+    def __init__(self, max_episodes: int, verbose: int = 0):
+        super(StopTrainingOnMaxEpisodes, self).__init__(verbose=verbose)
+        self.max_episodes = max_episodes
+        self._total_max_episodes = max_episodes
+        self.n_episodes = 0
+
+    def _init_callback(self):
+        # At start set total max according to number of envirnments
+        self._total_max_episodes = self.max_episodes * self.training_env.num_envs
+
+    def _on_step(self) -> bool:
+        # Checking for both 'done' and 'dones' keywords because:
+        # Some models use keyword 'done' (e.g.,: SAC, TD3, DQN, DDPG)
+        # While some models use keyword 'dones' (e.g.,: A2C, PPO)
+        done_array = np.array(self.locals.get("done") if self.locals.get("done") is not None else self.locals.get("dones"))
+        self.n_episodes += np.sum(done_array).item()
+
+        continue_training = self.n_episodes < self._total_max_episodes
+
+        if self.verbose > 0 and not continue_training:
+            mean_episodes_per_env = self.n_episodes / self.training_env.num_envs
+            mean_ep_str = (
+                "with an average of {mean_episodes_per_env:.2f} episodes per env" if self.training_env.num_envs > 1 else ""
+            )
+
+            print(
+                f"Stopping training with a total of {self.num_timesteps} steps because the "
+                f"{self.locals.get('tb_log_name')} model reached max_episodes={self.max_episodes}, "
+                f"by playing for {self.n_episodes} episodes "
+                f"{mean_ep_str}"
+            )
+        return continue_training

stable_baselines3/common/off_policy_algorithm.py

@@ -401,6 +401,10 @@ class OffPolicyAlgorithm(BaseAlgorithm):
                 # Rescale and perform action
                 new_obs, reward, done, infos = env.step(action)
 
+                self.num_timesteps += 1
+                episode_timesteps += 1
+                total_steps += 1
+
                 # Give access to local variables
                 callback.update_locals(locals())
                 # Only stop training if return value is False, not when it is None.
@@ -429,9 +433,6 @@ class OffPolicyAlgorithm(BaseAlgorithm):
                 if self._vec_normalize_env is not None:
                     self._last_original_obs = new_obs_
 
-                self.num_timesteps += 1
-                episode_timesteps += 1
-                total_steps += 1
                 self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
 
                 # For DQN, check if the target network should be updated

stable_baselines3/common/on_policy_algorithm.py

@@ -162,6 +162,8 @@ class OnPolicyAlgorithm(BaseAlgorithm):
 
             new_obs, rewards, dones, infos = env.step(clipped_actions)
 
+            self.num_timesteps += env.num_envs
+
             # Give access to local variables
             callback.update_locals(locals())
             if callback.on_step() is False:
@@ -169,7 +171,6 @@ class OnPolicyAlgorithm(BaseAlgorithm):
 
             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

tests/test_callbacks.py

@@ -10,15 +10,17 @@ from stable_baselines3.common.callbacks import (
     CheckpointCallback,
     EvalCallback,
     EveryNTimesteps,
+    StopTrainingOnMaxEpisodes,
     StopTrainingOnRewardThreshold,
 )
+from stable_baselines3.common.cmd_util import make_vec_env
 
 
 @pytest.mark.parametrize("model_class", [A2C, PPO, SAC, TD3, DQN, DDPG])
 def test_callbacks(tmp_path, model_class):
     log_folder = tmp_path / "logs/callbacks/"
 
-    # Dyn only support discrete actions
+    # DQN only support discrete actions
     env_name = select_env(model_class)
     # Create RL model
     # Small network for fast test
@@ -39,7 +41,10 @@ def test_callbacks(tmp_path, model_class):
 
     event_callback = EveryNTimesteps(n_steps=500, callback=checkpoint_on_event)
 
-    callback = CallbackList([checkpoint_callback, eval_callback, event_callback])
+    # Stop training if max number of episodes is reached
+    callback_max_episodes = StopTrainingOnMaxEpisodes(max_episodes=100, verbose=1)
+
+    callback = CallbackList([checkpoint_callback, eval_callback, event_callback, callback_max_episodes])
     model.learn(500, callback=callback)
 
     # Check access to local variables
@@ -48,12 +53,36 @@ def test_callbacks(tmp_path, model_class):
     assert checkpoint_callback.locals["new_obs"] is callback.locals["new_obs"]
     assert event_callback.locals["new_obs"] is callback.locals["new_obs"]
     assert checkpoint_on_event.locals["new_obs"] is callback.locals["new_obs"]
+    # Check that internal callback counters match models' counters
+    assert event_callback.num_timesteps == model.num_timesteps
+    assert event_callback.n_calls == model.num_timesteps
 
     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)
+
+    # Testing models that support multiple envs
+    if model_class in [A2C, PPO]:
+        max_episodes = 1
+        n_envs = 2
+        # Pendulum-v0 has a timelimit of 200 timesteps
+        max_episode_length = 200
+        envs = make_vec_env(env_name, n_envs=n_envs, seed=0)
+
+        model = model_class("MlpPolicy", envs, policy_kwargs=dict(net_arch=[32]))
+
+        callback_max_episodes = StopTrainingOnMaxEpisodes(max_episodes=max_episodes, verbose=1)
+        callback = CallbackList([callback_max_episodes])
+        model.learn(1000, callback=callback)
+
+        # Check that the actual number of episodes and timesteps per env matches the expected one
+        episodes_per_env = callback_max_episodes.n_episodes // n_envs
+        assert episodes_per_env == max_episodes
+        timesteps_per_env = model.num_timesteps // n_envs
+        assert timesteps_per_env == max_episode_length
+
     if os.path.exists(log_folder):
         shutil.rmtree(log_folder)