From 59af0c1b01f5905026fa0074fa5339642418992d Mon Sep 17 00:00:00 2001 From: Anand Balakrishnan Date: Thu, 25 Aug 2022 01:57:51 -0700 Subject: [PATCH] `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)