Merge branch 'master' into feat/crr

.gitlab-ci.yml

@@ -1,4 +1,4 @@
-image: stablebaselines/stable-baselines3-cpu:0.8.0a4
+image: stablebaselines/stable-baselines3-cpu:0.9.0a1
 
 type-check:
   script:

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)
 ------------------------------
@@ -394,4 +398,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/buffers.py

@@ -171,12 +171,12 @@ class ReplayBuffer(BaseBuffer):
             mem_available = psutil.virtual_memory().available
 
         self.optimize_memory_usage = optimize_memory_usage
-        self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=observation_space.dtype)
+        self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=observation_space.dtype)
         if optimize_memory_usage:
             # `observations` contains also the next observation
             self.next_observations = None
         else:
-            self.next_observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=observation_space.dtype)
+            self.next_observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=observation_space.dtype)
         self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=action_space.dtype)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.dones = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
@@ -284,7 +284,7 @@ class RolloutBuffer(BaseBuffer):
         self.reset()
 
     def reset(self) -> None:
-        self.observations = np.zeros((self.buffer_size, self.n_envs,) + self.obs_shape, dtype=np.float32)
+        self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=np.float32)
         self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)

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

@@ -410,6 +410,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.
@@ -438,9 +442,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

stable_baselines3/common/policies.py

@@ -699,7 +699,10 @@ class ContinuousCritic(BaseModel):
         n_critics: int = 2,
     ):
         super().__init__(
-            observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images,
+            observation_space,
+            action_space,
+            features_extractor=features_extractor,
+            normalize_images=normalize_images,
         )
 
         action_dim = get_action_dim(self.action_space)

stable_baselines3/common/save_util.py

@@ -350,7 +350,9 @@ def load_from_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], verbose=0)
 
 
 def load_from_zip_file(
-    load_path: Union[str, pathlib.Path, io.BufferedIOBase], load_data: bool = True, verbose=0,
+    load_path: Union[str, pathlib.Path, io.BufferedIOBase],
+    load_data: bool = True,
+    verbose=0,
 ) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]):
     """
     Load model data from a .zip archive

stable_baselines3/dqn/policies.py

@@ -31,7 +31,10 @@ class QNetwork(BasePolicy):
         normalize_images: bool = True,
     ):
         super(QNetwork, self).__init__(
-            observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images,
+            observation_space,
+            action_space,
+            features_extractor=features_extractor,
+            normalize_images=normalize_images,
         )
 
         if net_arch is None:

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)
 

tests/test_custom_policy.py

@@ -22,7 +22,7 @@ def test_flexible_mlp(model_class, net_arch):
     _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
 
 
-@pytest.mark.parametrize("net_arch", [[4], [4, 4],])
+@pytest.mark.parametrize("net_arch", [[4], [4, 4]])
 @pytest.mark.parametrize("model_class", [SAC, TD3])
 def test_custom_offpolicy(model_class, net_arch):
     _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=dict(net_arch=net_arch)).learn(1000)

tests/test_distributions.py

@@ -67,7 +67,11 @@ def test_sde_distribution():
 
 # TODO: analytical form for squashed Gaussian?
 @pytest.mark.parametrize(
-    "dist", [DiagGaussianDistribution(N_ACTIONS), StateDependentNoiseDistribution(N_ACTIONS, squash_output=False),]
+    "dist",
+    [
+        DiagGaussianDistribution(N_ACTIONS),
+        StateDependentNoiseDistribution(N_ACTIONS, squash_output=False),
+    ],
 )
 def test_entropy(dist):
     # The entropy can be approximated by averaging the negative log likelihood

tests/test_vec_envs.py

@@ -225,7 +225,7 @@ def check_vecenv_spaces(vec_env_class, space, obs_assert):
 
 def check_vecenv_obs(obs, space):
     """Helper method to check observations from multiple environments each belong to
-       the appropriate observation space."""
+    the appropriate observation space."""
     assert obs.shape[0] == N_ENVS
     for value in obs:
         assert space.contains(value)