From 01cc127d3277c9df19c662be040b3349dbb2019f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timoth=C3=A9?= <58693721+timothe-chaumont@users.noreply.github.com> Date: Mon, 22 Aug 2022 21:06:54 +0100 Subject: [PATCH 1/4] Support hparams logging to tensorboard (#984) * create Hparam class & support in all OutputFormats * add hparams documentation & example * add hparam tests * remove unnecessary test & fix name * format changes * support hyperparameters logging to tensorboard * fix HParams class docstring * use more explicit variable names * raise error instead of warning * Unpin protobuf * Add test for logging hparams Co-authored-by: Antonin RAFFIN --- docs/guide/tensorboard.rst | 49 ++++++++++++++++++++++++++++++ docs/misc/changelog.rst | 8 +++-- setup.py | 5 +-- stable_baselines3/common/logger.py | 32 +++++++++++++++++++ stable_baselines3/version.txt | 2 +- tests/test_logger.py | 14 +++++++++ tests/test_tensorboard.py | 39 ++++++++++++++++++++++-- 7 files changed, 139 insertions(+), 10 deletions(-) diff --git a/docs/guide/tensorboard.rst b/docs/guide/tensorboard.rst index 625c1be..89681d1 100644 --- a/docs/guide/tensorboard.rst +++ b/docs/guide/tensorboard.rst @@ -249,6 +249,55 @@ Here is an example of how to render an episode and log the resulting video to Te video_recorder = VideoRecorderCallback(gym.make("CartPole-v1"), render_freq=5000) model.learn(total_timesteps=int(5e4), callback=video_recorder) +Logging Hyperparameters +----------------------- + +TensorBoard supports logging of hyperparameters in its HPARAMS tab, which helps comparing agents trainings. + +.. warning:: + To display hyperparameters in the HPARAMS section, a ``metric_dict`` must be given (as well as a ``hparam_dict``). + + +Here is an example of how to save hyperparameters in TensorBoard: + +.. code-block:: python + + from stable_baselines3 import A2C + from stable_baselines3.common.callbacks import BaseCallback + from stable_baselines3.common.logger import HParam + + + class HParamCallback(BaseCallback): + def __init__(self): + """ + Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. + """ + super().__init__() + + def _on_training_start(self) -> None: + hparam_dict = { + "algorithm": self.model.__class__.__name__, + "learning rate": self.model.learning_rate, + "gamma": self.model.gamma, + } + # define the metrics that will appear in the `HPARAMS` Tensorboard tab by referencing their tag + # Tensorbaord will find & display metrics from the `SCALARS` tab + metric_dict = { + "rollout/ep_len_mean": 0, + "train/value_loss": 0, + } + self.logger.record( + "hparams", + HParam(hparam_dict, metric_dict), + exclude=("stdout", "log", "json", "csv"), + ) + + def _on_step(self) -> bool: + return True + + + model = A2C("MlpPolicy", "CartPole-v1", tensorboard_log="runs/", verbose=1) + model.learn(total_timesteps=int(5e4), callback=HParamCallback()) Directly Accessing The Summary Writer ------------------------------------- diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index eca1173..e30e28f 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,14 +3,16 @@ Changelog ========== -Release 1.6.1a0 (WIP) +Release 1.6.1a1 (WIP) --------------------------- Breaking Changes: ^^^^^^^^^^^^^^^^^ +- Switched minimum tensorboard version to 2.9.1 New Features: ^^^^^^^^^^^^^ +- Support logging hyperparameters to tensorboard (@timothe-chaumont) SB3-Contrib ^^^^^^^^^^^ @@ -33,12 +35,12 @@ Others: Documentation: ^^^^^^^^^^^^^^ +- Added an example of callback that logs hyperparameters to tensorboard. (@timothe-chaumont) - Fixed typo in docstring "nature" -> "Nature" (@Melanol) - Added info on split tensorboard logs into (@Melanol) - Fixed typo in ppo doc (@francescoluciano) - Fixed typo in install doc(@jlp-ue) - Release 1.6.0 (2022-07-11) --------------------------- @@ -1024,4 +1026,4 @@ And all the contributors: @eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP @simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485 @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede -@Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb +@Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont diff --git a/setup.py b/setup.py index 2816316..8c410a3 100644 --- a/setup.py +++ b/setup.py @@ -122,10 +122,7 @@ setup( "autorom[accept-rom-license]~=0.4.2", "pillow", # Tensorboard support - "tensorboard>=2.2.0", - # Protobuf >= 4 has breaking changes - # which does play well with tensorboard - "protobuf~=3.19.0", + "tensorboard>=2.9.1", # Checking memory taken by replay buffer "psutil", ], diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py index 1295e5b..c1e8433 100644 --- a/stable_baselines3/common/logger.py +++ b/stable_baselines3/common/logger.py @@ -14,6 +14,7 @@ from matplotlib import pyplot as plt try: from torch.utils.tensorboard import SummaryWriter + from torch.utils.tensorboard.summary import hparams except ImportError: SummaryWriter = None @@ -66,6 +67,22 @@ class Image: self.dataformats = dataformats +class HParam: + """ + Hyperparameter data class storing hyperparameters and metrics in dictionnaries + + :param hparam_dict: key-value pairs of hyperparameters to log + :param metric_dict: key-value pairs of metrics to log + A non-empty metrics dict is required to display hyperparameters in the corresponding Tensorboard section. + """ + + def __init__(self, hparam_dict: Dict[str, Union[bool, str, float, int, None]], metric_dict: Dict[str, Union[float, int]]): + self.hparam_dict = hparam_dict + if not metric_dict: + raise Exception("`metric_dict` must not be empty to display hyperparameters to the HPARAMS tensorboard tab.") + self.metric_dict = metric_dict + + class FormatUnsupportedError(NotImplementedError): """ Custom error to display informative message when @@ -165,6 +182,9 @@ class HumanOutputFormat(KVWriter, SeqWriter): elif isinstance(value, Image): raise FormatUnsupportedError(["stdout", "log"], "image") + elif isinstance(value, HParam): + raise FormatUnsupportedError(["stdout", "log"], "hparam") + elif isinstance(value, float): # Align left value_str = f"{value:<8.3g}" @@ -264,6 +284,8 @@ class JSONOutputFormat(KVWriter): raise FormatUnsupportedError(["json"], "figure") if isinstance(value, Image): raise FormatUnsupportedError(["json"], "image") + if isinstance(value, HParam): + raise FormatUnsupportedError(["json"], "hparam") if hasattr(value, "dtype"): if value.shape == () or len(value) == 1: # if value is a dimensionless numpy array or of length 1, serialize as a float @@ -333,6 +355,9 @@ class CSVOutputFormat(KVWriter): elif isinstance(value, Image): raise FormatUnsupportedError(["csv"], "image") + elif isinstance(value, HParam): + raise FormatUnsupportedError(["csv"], "hparam") + elif isinstance(value, str): # escape quotechars by prepending them with another quotechar value = value.replace(self.quotechar, self.quotechar + self.quotechar) @@ -389,6 +414,13 @@ class TensorBoardOutputFormat(KVWriter): if isinstance(value, Image): self.writer.add_image(key, value.image, step, dataformats=value.dataformats) + if isinstance(value, HParam): + # we don't use `self.writer.add_hparams` to have control over the log_dir + experiment, session_start_info, session_end_info = hparams(value.hparam_dict, metric_dict=value.metric_dict) + self.writer.file_writer.add_summary(experiment) + self.writer.file_writer.add_summary(session_start_info) + self.writer.file_writer.add_summary(session_end_info) + # Flush the output to the file self.writer.flush() diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 035e3b6..e36b727 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.6.1a0 +1.6.1a1 diff --git a/tests/test_logger.py b/tests/test_logger.py index a55f88a..516a622 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -17,6 +17,7 @@ from stable_baselines3.common.logger import ( CSVOutputFormat, Figure, FormatUnsupportedError, + HParam, HumanOutputFormat, Image, Logger, @@ -296,6 +297,19 @@ def test_report_figure_to_unsupported_format_raises_error(tmp_path, unsupported_ writer.close() +@pytest.mark.parametrize("unsupported_format", ["stdout", "log", "json", "csv"]) +def test_report_hparam_to_unsupported_format_raises_error(tmp_path, unsupported_format): + writer = make_output_format(unsupported_format, tmp_path) + + with pytest.raises(FormatUnsupportedError) as exec_info: + hparam_dict = {"learning rate": np.random.random()} + metric_dict = {"train/value_loss": 0} + hparam = HParam(hparam_dict=hparam_dict, metric_dict=metric_dict) + writer.write({"hparam": hparam}, key_excluded={"hparam": ()}) + assert unsupported_format in str(exec_info.value) + writer.close() + + def test_key_length(tmp_path): writer = make_output_format("stdout", tmp_path) assert writer.max_length == 36 diff --git a/tests/test_tensorboard.py b/tests/test_tensorboard.py index 6dccf41..8aa864d 100644 --- a/tests/test_tensorboard.py +++ b/tests/test_tensorboard.py @@ -3,6 +3,8 @@ import os import pytest from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3.common.callbacks import BaseCallback +from stable_baselines3.common.logger import HParam from stable_baselines3.common.utils import get_latest_run_id MODEL_DICT = { @@ -15,6 +17,34 @@ MODEL_DICT = { N_STEPS = 100 +class HParamCallback(BaseCallback): + def __init__(self): + """ + Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard. + """ + super().__init__() + + def _on_training_start(self) -> None: + hparam_dict = { + "algorithm": self.model.__class__.__name__, + "learning rate": self.model.learning_rate, + "gamma": self.model.gamma, + } + # define the metrics that will appear in the `HPARAMS` Tensorboard tab by referencing their tag + # Tensorbaord will find & display metrics from the `SCALARS` tab + metric_dict = { + "rollout/ep_len_mean": 0, + } + self.logger.record( + "hparams", + HParam(hparam_dict, metric_dict), + exclude=("stdout", "log", "json", "csv"), + ) + + def _on_step(self) -> bool: + return True + + @pytest.mark.parametrize("model_name", MODEL_DICT.keys()) def test_tensorboard(tmp_path, model_name): # Skip if no tensorboard installed @@ -22,8 +52,13 @@ def test_tensorboard(tmp_path, model_name): logname = model_name.upper() algo, env_id = MODEL_DICT[model_name] - model = algo("MlpPolicy", env_id, verbose=1, tensorboard_log=tmp_path) - model.learn(N_STEPS) + kwargs = {} + if model_name == "ppo": + kwargs["n_steps"] = 64 + elif model_name in {"sac", "td3"}: + kwargs["train_freq"] = 2 + model = algo("MlpPolicy", env_id, verbose=1, tensorboard_log=tmp_path, **kwargs) + model.learn(N_STEPS, callback=HParamCallback()) model.learn(N_STEPS, reset_num_timesteps=False) assert os.path.isdir(tmp_path / str(logname + "_1")) From 29a481a288b58008dbb4b0b4af5b360121aea120 Mon Sep 17 00:00:00 2001 From: Honglu Fan <64070721+honglu2875@users.noreply.github.com> Date: Tue, 23 Aug 2022 10:20:43 +0200 Subject: [PATCH 2/4] Include `running_mean` and `running_val` when updating target networks (#1004) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * include `running_mean` and `running_val` when updating target networks in DQN, SAC, TD3. * Update stable_baselines3/common/utils.py Co-authored-by: Antonin RAFFIN * Precompute batch norm parameters in `_setup_model` and directly copy them in the target update. * include `running_mean` and `running_val` when updating target networks in DQN, SAC, TD3. * Update stable_baselines3/common/utils.py Co-authored-by: Antonin RAFFIN * Precompute batch norm parameters in `_setup_model` and directly copy them in the target update. * Fix `DictReplayBuffer.next_observations` type (#1013) * Fix DictReplayBuffer.next_observations type * Update changelog Co-authored-by: Antonin RAFFIN * Fixed missing verbose parameter passing (#1011) Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com> * Support for `device=auto` buffers and set it as default value (#1009) * Default device is "auto" for buffer + auto device support in BufferBaseClass * Update docstring * Update tests * Unify tests * Update changelog * Fix tests on CUDA device Co-authored-by: Antonin RAFFIN Co-authored-by: Antonin Raffin * Precompute batch norm parameters in `_setup_model` and directly copy them in the target update. * Update test * Add comments and update tests * Bump version * Remove one extra space to conform code style. * Update docstrings Co-authored-by: Antonin RAFFIN Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com> Co-authored-by: Burak Demirbilek Co-authored-by: Antonin Raffin --- docs/misc/changelog.rst | 5 +++-- stable_baselines3/common/utils.py | 37 +++++++++++++++++++++---------- stable_baselines3/dqn/dqn.py | 7 +++++- stable_baselines3/sac/sac.py | 7 +++++- stable_baselines3/td3/td3.py | 10 ++++++++- stable_baselines3/version.txt | 2 +- tests/test_train_eval_mode.py | 24 +++++++++++++++----- tests/test_utils.py | 24 +++++++++++++++++++- 8 files changed, 91 insertions(+), 25 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index e30e28f..dd32743 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,7 +3,7 @@ Changelog ========== -Release 1.6.1a1 (WIP) +Release 1.6.1a2 (WIP) --------------------------- Breaking Changes: @@ -23,6 +23,7 @@ Bug Fixes: - Fixed division by zero error when computing FPS when a small number of time has elapsed in operating systems with low-precision timers. - Added multidimensional action space support (@qgallouedec) - Fixed missing verbose parameter passing in the ``EvalCallback`` constructor (@burakdmb) +- Fixed the issue that when updating the target network in DQN, SAC, TD3, the ``running_mean`` and ``running_var`` properties of batch norm layers are not updated (@honglu2875) Deprecations: ^^^^^^^^^^^^^ @@ -1026,4 +1027,4 @@ And all the contributors: @eleurent @ac-93 @cove9988 @theDebugger811 @hsuehch @Demetrio92 @thomasgubler @IperGiove @ScheiklP @simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485 @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede -@Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont +@Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 94cd658..a2126a2 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -4,7 +4,7 @@ import platform import random from collections import deque from itertools import zip_longest -from typing import Dict, Iterable, Optional, Tuple, Union +from typing import Dict, Iterable, List, Optional, Tuple, Union import gym import numpy as np @@ -67,8 +67,8 @@ def update_learning_rate(optimizer: th.optim.Optimizer, learning_rate: float) -> Update the learning rate for a given optimizer. Useful when doing linear schedule. - :param optimizer: - :param learning_rate: + :param optimizer: Pytorch optimizer + :param learning_rate: New learning rate value """ for param_group in optimizer.param_groups: param_group["lr"] = learning_rate @@ -79,8 +79,8 @@ def get_schedule_fn(value_schedule: Union[Schedule, float, int]) -> Schedule: Transform (if needed) learning rate and clip range (for PPO) to callable. - :param value_schedule: - :return: + :param value_schedule: Constant value of schedule function + :return: Schedule function (can return constant value) """ # If the passed schedule is a float # create a constant function @@ -104,7 +104,7 @@ def get_linear_fn(start: float, end: float, end_fraction: float) -> Schedule: :params end_fraction: fraction of ``progress_remaining`` where end is reached e.g 0.1 then end is reached after 10% of the complete training process. - :return: + :return: Linear schedule function. """ def func(progress_remaining: float) -> float: @@ -121,8 +121,8 @@ def constant_fn(val: float) -> Schedule: Create a function that returns a constant It is useful for learning rate schedule (to avoid code duplication) - :param val: - :return: + :param val: constant value + :return: Constant schedule function. """ def func(_): @@ -139,7 +139,7 @@ def get_device(device: Union[th.device, str] = "auto") -> th.device: By default, it tries to use the gpu. :param device: One for 'auto', 'cuda', 'cpu' - :return: + :return: Supported Pytorch device """ # Cuda by default if device == "auto": @@ -386,12 +386,25 @@ def safe_mean(arr: Union[np.ndarray, list, deque]) -> np.ndarray: Compute the mean of an array if there is at least one element. For empty array, return NaN. It is used for logging only. - :param arr: + :param arr: Numpy array or list of values :return: """ return np.nan if len(arr) == 0 else np.mean(arr) +def get_parameters_by_name(model: th.nn.Module, included_names: Iterable[str]) -> List[th.Tensor]: + """ + Extract parameters from the state dict of ``model`` + if the name contains one of the strings in ``included_names``. + + :param model: the model where the parameters come from. + :param included_names: substrings of names to include. + :return: List of parameters values (Pytorch tensors) + that matches the queried names. + """ + return [param for name, param in model.state_dict().items() if any([key in name for key in included_names])] + + def zip_strict(*iterables: Iterable) -> Iterable: r""" ``zip()`` function but enforces that iterables are of equal length. @@ -411,8 +424,8 @@ def zip_strict(*iterables: Iterable) -> Iterable: def polyak_update( - params: Iterable[th.nn.Parameter], - target_params: Iterable[th.nn.Parameter], + params: Iterable[th.Tensor], + target_params: Iterable[th.Tensor], tau: float, ) -> None: """ diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 0cd6dfb..839fe33 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -11,7 +11,7 @@ from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.policies import BasePolicy from stable_baselines3.common.preprocessing import maybe_transpose from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule -from stable_baselines3.common.utils import get_linear_fn, is_vectorized_observation, polyak_update +from stable_baselines3.common.utils import get_linear_fn, get_parameters_by_name, is_vectorized_observation, polyak_update from stable_baselines3.dqn.policies import CnnPolicy, DQNPolicy, MlpPolicy, MultiInputPolicy @@ -140,6 +140,9 @@ class DQN(OffPolicyAlgorithm): def _setup_model(self) -> None: super()._setup_model() self._create_aliases() + # Copy running stats, see GH issue #996 + self.batch_norm_stats = get_parameters_by_name(self.q_net, ["running_"]) + self.batch_norm_stats_target = get_parameters_by_name(self.q_net_target, ["running_"]) self.exploration_schedule = get_linear_fn( self.exploration_initial_eps, self.exploration_final_eps, @@ -170,6 +173,8 @@ class DQN(OffPolicyAlgorithm): self._n_calls += 1 if self._n_calls % self.target_update_interval == 0: polyak_update(self.q_net.parameters(), self.q_net_target.parameters(), self.tau) + # Copy running stats, see GH issue #996 + polyak_update(self.batch_norm_stats, self.batch_norm_stats_target, 1.0) self.exploration_rate = self.exploration_schedule(self._current_progress_remaining) self.logger.record("rollout/exploration_rate", self.exploration_rate) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 07f88d9..6969ef1 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -10,7 +10,7 @@ from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.policies import BasePolicy from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule -from stable_baselines3.common.utils import polyak_update +from stable_baselines3.common.utils import get_parameters_by_name, polyak_update from stable_baselines3.sac.policies import CnnPolicy, MlpPolicy, MultiInputPolicy, SACPolicy @@ -152,6 +152,9 @@ class SAC(OffPolicyAlgorithm): def _setup_model(self) -> None: super()._setup_model() self._create_aliases() + # Running mean and running var + self.batch_norm_stats = get_parameters_by_name(self.critic, ["running_"]) + self.batch_norm_stats_target = get_parameters_by_name(self.critic_target, ["running_"]) # Target entropy is used when learning the entropy coefficient if self.target_entropy == "auto": # automatically set target entropy if needed @@ -272,6 +275,8 @@ class SAC(OffPolicyAlgorithm): # Update target networks if gradient_step % self.target_update_interval == 0: polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau) + # Copy running stats, see GH issue #996 + polyak_update(self.batch_norm_stats, self.batch_norm_stats_target, 1.0) self._n_updates += gradient_steps diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 34a783d..f440b73 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -10,7 +10,7 @@ from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.policies import BasePolicy from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule -from stable_baselines3.common.utils import polyak_update +from stable_baselines3.common.utils import get_parameters_by_name, polyak_update from stable_baselines3.td3.policies import CnnPolicy, MlpPolicy, MultiInputPolicy, TD3Policy @@ -131,6 +131,11 @@ class TD3(OffPolicyAlgorithm): def _setup_model(self) -> None: super()._setup_model() self._create_aliases() + # Running mean and running var + self.actor_batch_norm_stats = get_parameters_by_name(self.actor, ["running_"]) + self.critic_batch_norm_stats = get_parameters_by_name(self.critic, ["running_"]) + self.actor_batch_norm_stats_target = get_parameters_by_name(self.actor_target, ["running_"]) + self.critic_batch_norm_stats_target = get_parameters_by_name(self.critic_target, ["running_"]) def _create_aliases(self) -> None: self.actor = self.policy.actor @@ -189,6 +194,9 @@ class TD3(OffPolicyAlgorithm): polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau) polyak_update(self.actor.parameters(), self.actor_target.parameters(), self.tau) + # Copy running stats, see GH issue #996 + polyak_update(self.critic_batch_norm_stats, self.critic_batch_norm_stats_target, 1.0) + polyak_update(self.actor_batch_norm_stats, self.actor_batch_norm_stats_target, 1.0) self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard") if len(actor_losses) > 0: diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index e36b727..51cf83a 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.6.1a1 +1.6.1a2 diff --git a/tests/test_train_eval_mode.py b/tests/test_train_eval_mode.py index 4f023e9..a1a63c0 100644 --- a/tests/test_train_eval_mode.py +++ b/tests/test_train_eval_mode.py @@ -143,7 +143,8 @@ def test_dqn_train_with_batch_norm(): policy_kwargs=dict(net_arch=[16, 16], features_extractor_class=FlattenBatchNormDropoutExtractor), learning_starts=0, seed=1, - tau=0, # do not clone the target + tau=0.0, # do not clone the target + target_update_interval=100, # Copy the stats to the target ) ( @@ -154,6 +155,9 @@ def test_dqn_train_with_batch_norm(): ) = clone_dqn_batch_norm_stats(model) model.learn(total_timesteps=200) + # Force stats copy + model.target_update_interval = 1 + model._on_step() ( q_net_bias_after, @@ -165,8 +169,12 @@ def test_dqn_train_with_batch_norm(): assert ~th.isclose(q_net_bias_before, q_net_bias_after).all() assert ~th.isclose(q_net_running_mean_before, q_net_running_mean_after).all() + # No weight update + assert th.isclose(q_net_bias_before, q_net_target_bias_after).all() assert th.isclose(q_net_target_bias_before, q_net_target_bias_after).all() - assert th.isclose(q_net_target_running_mean_before, q_net_target_running_mean_after).all() + # Running stat should be copied even when tau=0 + assert th.isclose(q_net_running_mean_before, q_net_target_running_mean_before).all() + assert th.isclose(q_net_running_mean_after, q_net_target_running_mean_after).all() def test_td3_train_with_batch_norm(): @@ -210,10 +218,12 @@ def test_td3_train_with_batch_norm(): assert ~th.isclose(critic_running_mean_before, critic_running_mean_after).all() assert th.isclose(actor_target_bias_before, actor_target_bias_after).all() - assert th.isclose(actor_target_running_mean_before, actor_target_running_mean_after).all() + # Running stat should be copied even when tau=0 + assert th.isclose(actor_running_mean_after, actor_target_running_mean_after).all() assert th.isclose(critic_target_bias_before, critic_target_bias_after).all() - assert th.isclose(critic_target_running_mean_before, critic_target_running_mean_after).all() + # Running stat should be copied even when tau=0 + assert th.isclose(critic_running_mean_after, critic_target_running_mean_after).all() def test_sac_train_with_batch_norm(): @@ -250,10 +260,12 @@ def test_sac_train_with_batch_norm(): assert ~th.isclose(actor_running_mean_before, actor_running_mean_after).all() assert ~th.isclose(critic_bias_before, critic_bias_after).all() - assert ~th.isclose(critic_running_mean_before, critic_running_mean_after).all() + # Running stat should be copied even when tau=0 + assert th.isclose(critic_running_mean_before, critic_target_running_mean_before).all() assert th.isclose(critic_target_bias_before, critic_target_bias_after).all() - assert th.isclose(critic_target_running_mean_before, critic_target_running_mean_after).all() + # Running stat should be copied even when tau=0 + assert th.isclose(critic_running_mean_after, critic_target_running_mean_after).all() @pytest.mark.parametrize("model_class", [A2C, PPO]) diff --git a/tests/test_utils.py b/tests/test_utils.py index 67f2ad1..57b4b39 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -14,7 +14,13 @@ from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_v from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise, VectorizedActionNoise -from stable_baselines3.common.utils import get_system_info, is_vectorized_observation, polyak_update, zip_strict +from stable_baselines3.common.utils import ( + get_parameters_by_name, + get_system_info, + is_vectorized_observation, + polyak_update, + zip_strict, +) from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv @@ -322,6 +328,22 @@ def test_vec_noise(): assert len(vec.noises) == num_envs +def test_get_parameters_by_name(): + model = th.nn.Sequential(th.nn.Linear(5, 5), th.nn.BatchNorm1d(5)) + # Initialize stats + model(th.ones(3, 5)) + included_names = ["weight", "bias", "running_"] + # 2 x weight, 2 x bias, 1 x running_mean, 1 x running_var; Ignore num_batches_tracked. + parameters = get_parameters_by_name(model, included_names) + assert len(parameters) == 6 + assert th.allclose(parameters[4], model[1].running_mean) + assert th.allclose(parameters[5], model[1].running_var) + parameters = get_parameters_by_name(model, ["running_"]) + assert len(parameters) == 2 + assert th.allclose(parameters[0], model[1].running_mean) + assert th.allclose(parameters[1], model[1].running_var) + + def test_polyak(): param1, param2 = th.nn.Parameter(th.ones((5, 5))), th.nn.Parameter(th.zeros((5, 5))) target1, target2 = th.nn.Parameter(th.ones((5, 5))), th.nn.Parameter(th.zeros((5, 5))) From 59af0c1b01f5905026fa0074fa5339642418992d Mon Sep 17 00:00:00 2001 From: Anand Balakrishnan Date: Thu, 25 Aug 2022 01:57:51 -0700 Subject: [PATCH 3/4] `CheckpointCallback` can now save replay buffer and `VecNormalize` (#1030) * CheckpointCallback now saves replay buffer (if present) * VecNormalize stats are saved at checkpoints * Make checkpointing replay buffer and VecNormalize opt-in * Edit changelog * Add documentation for new parameters * Update docs/misc/changelog.rst * Add documentation for new parameters * Implement suggested edits * Reformat code * Fix git conflict * Add .pkl suffix to VecNormalize checkpoints * Add tests for new CheckpointCallback params * Merge CheckpointCallback tests * Update test and add helper for checkpoint path Co-authored-by: Antonin RAFFIN --- docs/guide/callbacks.rst | 24 +++++++---- docs/misc/changelog.rst | 2 + stable_baselines3/common/callbacks.py | 57 +++++++++++++++++++++++---- tests/test_callbacks.py | 26 ++++++++++++ 4 files changed, 95 insertions(+), 14 deletions(-) diff --git a/docs/guide/callbacks.rst b/docs/guide/callbacks.rst index 6c7f4eb..7e22bbf 100644 --- a/docs/guide/callbacks.rst +++ b/docs/guide/callbacks.rst @@ -157,6 +157,10 @@ CheckpointCallback Callback for saving a model every ``save_freq`` calls to ``env.step()``, you must specify a log folder (``save_path``) and optionally a prefix for the checkpoints (``rl_model`` by default). +If you are using this callback to stop and resume training, you may want to optionally save the replay buffer if the +model has one (``save_replay_buffer``, ``False`` by default). +Additionally, if your environment uses a :ref:`VecNormalize ` wrapper, you can save the +corresponding statistics using ``save_vecnormalize`` (``False`` by default). .. warning:: @@ -168,14 +172,20 @@ and optionally a prefix for the checkpoints (``rl_model`` by default). .. code-block:: python - from stable_baselines3 import SAC - from stable_baselines3.common.callbacks import CheckpointCallback - # Save a checkpoint every 1000 steps - checkpoint_callback = CheckpointCallback(save_freq=1000, save_path='./logs/', - name_prefix='rl_model') + from stable_baselines3 import SAC + from stable_baselines3.common.callbacks import CheckpointCallback - model = SAC('MlpPolicy', 'Pendulum-v1') - model.learn(2000, callback=checkpoint_callback) + # Save a checkpoint every 1000 steps + checkpoint_callback = CheckpointCallback( + save_freq=1000, + save_path="./logs/", + name_prefix="rl_model", + save_replay_buffer=True, + save_vecnormalize=True, + ) + + model = SAC("MlpPolicy", "Pendulum-v1") + model.learn(2000, callback=checkpoint_callback) .. _EvalCallback: diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index dd32743..126de2e 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -13,6 +13,7 @@ Breaking Changes: New Features: ^^^^^^^^^^^^^ - Support logging hyperparameters to tensorboard (@timothe-chaumont) +- Added checkpoints for replay buffer and ``VecNormalize`` statistics (@anand-bala) SB3-Contrib ^^^^^^^^^^^ @@ -1028,3 +1029,4 @@ And all the contributors: @simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485 @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 +@anand-bala diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index e9f46fe..a2abd44 100644 --- a/stable_baselines3/common/callbacks.py +++ b/stable_baselines3/common/callbacks.py @@ -15,7 +15,8 @@ class BaseCallback(ABC): """ Base class for callback. - :param verbose: + :param verbose: Verbosity of the output (set to 1 for info messages, + 2 for debug) """ def __init__(self, verbose: int = 0): @@ -214,6 +215,10 @@ class CheckpointCallback(BaseCallback): """ Callback for saving a model every ``save_freq`` calls to ``env.step()``. + By default, it only saves model checkpoints, + you need to pass ``save_replay_buffer=True``, + and ``save_vecnormalize=True`` to also save replay buffer checkpoints + and normalization statistics checkpoints. .. warning:: @@ -221,29 +226,67 @@ class CheckpointCallback(BaseCallback): will effectively correspond to ``n_envs`` steps. To account for that, you can use ``save_freq = max(save_freq // n_envs, 1)`` - :param save_freq: + :param save_freq: Save checkpoints every ``save_freq`` call of the callback. :param save_path: Path to the folder where the model will be saved. :param name_prefix: Common prefix to the saved models - :param verbose: + :param save_replay_buffer: Save the model replay buffer + :param save_vecnormalize: Save the ``VecNormalize`` statistics + :param verbose: Verbosity of the output (set to 2 for debug messages) """ - def __init__(self, save_freq: int, save_path: str, name_prefix: str = "rl_model", verbose: int = 0): + def __init__( + self, + save_freq: int, + save_path: str, + name_prefix: str = "rl_model", + save_replay_buffer: bool = False, + save_vecnormalize: bool = False, + verbose: int = 0, + ): super().__init__(verbose) self.save_freq = save_freq self.save_path = save_path self.name_prefix = name_prefix + self.save_replay_buffer = save_replay_buffer + self.save_vecnormalize = save_vecnormalize 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 _checkpoint_path(self, checkpoint_type: str = "", extension: str = "") -> str: + """ + Helper to get checkpoint path for each type of checkpoint. + + :param checkpoint_type: empty for the model, "replay_buffer_" + or "vecnormalize_" for the other checkpoints. + :param extension: Checkpoint file extension (zip for model, pkl for others) + :return: Path to the checkpoint + """ + return os.path.join(self.save_path, f"{self.name_prefix}_{checkpoint_type}{self.num_timesteps}_steps.{extension}") + 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) + model_path = self._checkpoint_path(extension="zip") + self.model.save(model_path) if self.verbose > 1: - print(f"Saving model checkpoint to {path}") + print(f"Saving model checkpoint to {model_path}") + + if self.save_replay_buffer and hasattr(self.model, "replay_buffer") and self.model.replay_buffer is not None: + # If model has a replay buffer, save it too + replay_buffer_path = self._checkpoint_path("replay_buffer_", extension="pkl") + self.model.save_replay_buffer(replay_buffer_path) + if self.verbose > 1: + print(f"Saving model replay buffer checkpoint to {replay_buffer_path}") + + if self.save_vecnormalize and self.model.get_vec_normalize_env() is not None: + # Save the VecNormalize statistics + vec_normalize_path = self._checkpoint_path("vecnormalize_", extension="pkl") + self.model.get_vec_normalize_env().save(vec_normalize_path) + if self.verbose > 1: + print(f"Saving model VecNormalize to {vec_normalize_path}") + return True diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 6576f7d..2c7e0ba 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -203,3 +203,29 @@ def test_eval_friendly_error(): with pytest.warns(Warning): with pytest.raises(AssertionError): model.learn(100, callback=eval_callback) + + +def test_checkpoint_additional_info(tmp_path): + # tests if the replay buffer and the VecNormalize stats are saved with every checkpoint + dummy_vec_env = DummyVecEnv([lambda: gym.make("CartPole-v1")]) + env = VecNormalize(dummy_vec_env) + + checkpoint_dir = tmp_path / "checkpoints" + checkpoint_callback = CheckpointCallback( + save_freq=200, + save_path=checkpoint_dir, + save_replay_buffer=True, + save_vecnormalize=True, + verbose=2, + ) + + model = DQN("MlpPolicy", env, learning_starts=100, buffer_size=500, seed=0) + model.learn(200, callback=checkpoint_callback) + + assert os.path.exists(checkpoint_dir / "rl_model_200_steps.zip") + assert os.path.exists(checkpoint_dir / "rl_model_replay_buffer_200_steps.pkl") + assert os.path.exists(checkpoint_dir / "rl_model_vecnormalize_200_steps.pkl") + # Check that checkpoints can be properly loaded + model = DQN.load(checkpoint_dir / "rl_model_200_steps.zip") + model.load_replay_buffer(checkpoint_dir / "rl_model_replay_buffer_200_steps.pkl") + VecNormalize.load(checkpoint_dir / "rl_model_vecnormalize_200_steps.pkl", dummy_vec_env) From 2cc1477fa28e654caa4b7bf9b367febe726513e8 Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 25 Aug 2022 05:50:08 -0400 Subject: [PATCH 4/4] Fix advantage normalization with mini-batchsize of 1 (#1028) * fix nan in advnatages with batch size 1, for ppo * changelog * black * Simplify test * Bump version Co-authored-by: Antonin Raffin --- docs/misc/changelog.rst | 5 +++-- stable_baselines3/ppo/ppo.py | 7 ++++--- stable_baselines3/sac/sac.py | 2 +- stable_baselines3/version.txt | 2 +- tests/test_run.py | 26 ++++++++++++++++++++++++++ tests/test_utils.py | 15 +-------------- 6 files changed, 36 insertions(+), 21 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 126de2e..4643011 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,7 +3,7 @@ Changelog ========== -Release 1.6.1a2 (WIP) +Release 1.6.1a3 (WIP) --------------------------- Breaking Changes: @@ -20,6 +20,7 @@ SB3-Contrib Bug Fixes: ^^^^^^^^^^ +- Fixed issue where ``PPO`` gives NaN if rollout buffer provides a batch of size 1 (@hughperkins) - Fixed the issue that ``predict`` does not always return action as ``np.ndarray`` (@qgallouedec) - Fixed division by zero error when computing FPS when a small number of time has elapsed in operating systems with low-precision timers. - Added multidimensional action space support (@qgallouedec) @@ -1029,4 +1030,4 @@ And all the contributors: @simoninithomas @armandpl @manuel-delverme @Gautam-J @gianlucadecola @buoyancy99 @caburu @xy9485 @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 -@anand-bala +@anand-bala @hughperkins diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index 5b8d9e2..0f7f8e4 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -137,8 +137,8 @@ class PPO(OnPolicyAlgorithm): # Check that `n_steps * n_envs > 1` to avoid NaN # when doing advantage normalization buffer_size = self.env.num_envs * self.n_steps - assert ( - buffer_size > 1 + assert buffer_size > 1 or ( + not normalize_advantage ), f"`n_steps * n_envs` must be greater than 1. Currently n_steps={self.n_steps} and n_envs={self.env.num_envs}" # Check that the rollout buffer size is a multiple of the mini-batch size untruncated_batches = buffer_size // batch_size @@ -210,7 +210,8 @@ class PPO(OnPolicyAlgorithm): values = values.flatten() # Normalize advantage advantages = rollout_data.advantages - if self.normalize_advantage: + # Normalization does not make sense if mini batchsize == 1, see GH issue #325 + if self.normalize_advantage and len(advantages) > 1: advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) # ratio between old and new policy, should be one at the first iteration diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 6969ef1..ba27998 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -261,7 +261,7 @@ class SAC(OffPolicyAlgorithm): # Compute actor loss # Alternative: actor_loss = th.mean(log_prob - qf1_pi) - # Mean over all critic networks + # Min over all critic networks q_values_pi = th.cat(self.critic(replay_data.observations, actions_pi), dim=1) min_qf_pi, _ = th.min(q_values_pi, dim=1, keepdim=True) actor_loss = (ent_coef * log_prob - min_qf_pi).mean() diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 51cf83a..7a35b06 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.6.1a2 +1.6.1a3 diff --git a/tests/test_run.py b/tests/test_run.py index b0a9a11..655182d 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -213,3 +213,29 @@ def test_warn_dqn_multi_env(): buffer_size=100, target_update_interval=1, ) + + +def test_ppo_warnings(): + """Test that PPO warns and errors correctly on + problematic rollout buffer sizes""" + + # Only 1 step: advantage normalization will return NaN + with pytest.raises(AssertionError): + PPO("MlpPolicy", "Pendulum-v1", n_steps=1) + + # batch_size of 1 is allowed when normalize_advantage=False + model = PPO("MlpPolicy", "Pendulum-v1", n_steps=1, batch_size=1, normalize_advantage=False) + model.learn(4) + + # Truncated mini-batch + # Batch size 1 yields NaN with normalized advantage because + # torch.std(some_length_1_tensor) == NaN + # advantage normalization is automatically deactivated + # in that case + with pytest.warns(UserWarning, match="there will be a truncated mini-batch of size 1"): + model = PPO("MlpPolicy", "Pendulum-v1", n_steps=64, batch_size=63, verbose=1) + model.learn(64) + + loss = model.logger.name_to_value["train/loss"] + assert loss > 0 + assert not np.isnan(loss) # check not nan (since nan does not equal nan) diff --git a/tests/test_utils.py b/tests/test_utils.py index 57b4b39..2a9eade 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -8,7 +8,7 @@ import torch as th from gym import spaces import stable_baselines3 as sb3 -from stable_baselines3 import A2C, PPO +from stable_baselines3 import A2C from stable_baselines3.common.atari_wrappers import ClipRewardEnv, MaxAndSkipEnv from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_vec_env, unwrap_wrapper from stable_baselines3.common.evaluation import evaluate_policy @@ -388,19 +388,6 @@ def test_is_wrapped(): assert unwrap_wrapper(env, Monitor) == monitor_env -def test_ppo_warnings(): - """Test that PPO warns and errors correctly on - problematic rollour buffer sizes""" - - # Only 1 step: advantage normalization will return NaN - with pytest.raises(AssertionError): - PPO("MlpPolicy", "Pendulum-v1", n_steps=1) - - # Truncated mini-batch - with pytest.warns(UserWarning): - PPO("MlpPolicy", "Pendulum-v1", n_steps=6, batch_size=8) - - def test_get_system_info(): info, info_str = get_system_info(print_info=True) assert info["Stable-Baselines3"] == str(sb3.__version__)