diff --git a/docs/guide/callbacks.rst b/docs/guide/callbacks.rst index 84687e9..7206c7c 100644 --- a/docs/guide/callbacks.rst +++ b/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: diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 364d17f..39f9cb6 100644 --- a/docs/misc/changelog.rst +++ b/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 diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index e0d326c..8f8b572 100644 --- a/stable_baselines3/common/callbacks.py +++ b/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 diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 45c735d..31ec5ef 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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 diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 9230dd7..721bf40 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/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 diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 55c7f61..110983b 100644 --- a/tests/test_callbacks.py +++ b/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)