From 15d32c6a4a0ddc19e64baaa16e4afd5a045e23ce Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Thu, 27 Aug 2020 23:02:59 +0200 Subject: [PATCH 1/2] Update black version + update docker image (#151) * Update docker image * Update black and reformat --- .gitlab-ci.yml | 2 +- stable_baselines3/common/buffers.py | 6 +++--- stable_baselines3/common/policies.py | 5 ++++- stable_baselines3/common/save_util.py | 4 +++- stable_baselines3/dqn/policies.py | 5 ++++- tests/test_custom_policy.py | 2 +- tests/test_distributions.py | 6 +++++- tests/test_vec_envs.py | 2 +- 8 files changed, 22 insertions(+), 10 deletions(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 695c145..71826e9 100644 --- a/.gitlab-ci.yml +++ b/.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: diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 4534063..6c58953 100644 --- a/stable_baselines3/common/buffers.py +++ b/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) diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 41280a8..0478342 100644 --- a/stable_baselines3/common/policies.py +++ b/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) diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index 51fa8cd..326db1e 100644 --- a/stable_baselines3/common/save_util.py +++ b/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 diff --git a/stable_baselines3/dqn/policies.py b/stable_baselines3/dqn/policies.py index f5001c7..ebbcd34 100644 --- a/stable_baselines3/dqn/policies.py +++ b/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: diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py index c1e08df..95f4a7c 100644 --- a/tests/test_custom_policy.py +++ b/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) diff --git a/tests/test_distributions.py b/tests/test_distributions.py index a73b81e..490f80e 100644 --- a/tests/test_distributions.py +++ b/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 diff --git a/tests/test_vec_envs.py b/tests/test_vec_envs.py index 8c33341..141ca6a 100644 --- a/tests/test_vec_envs.py +++ b/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) From 5fc90a7f7dd7f3b5d1bcea25a8fb6d9853947609 Mon Sep 17 00:00:00 2001 From: Francisco Caio Date: Fri, 28 Aug 2020 06:36:33 -0300 Subject: [PATCH 2/2] 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 Co-authored-by: Antonin Raffin --- docs/guide/callbacks.rst | 30 ++++++++++++ docs/misc/changelog.rst | 6 ++- stable_baselines3/common/callbacks.py | 47 ++++++++++++++++++- .../common/off_policy_algorithm.py | 7 +-- .../common/on_policy_algorithm.py | 3 +- tests/test_callbacks.py | 33 ++++++++++++- 6 files changed, 118 insertions(+), 8 deletions(-) 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)