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 <antonin.raffin@ensta.org>

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 <vec_env>` 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:

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

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
 
 

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)