[Feature Request] CheckpointCallback should also save replay buffer

[anand-bala]

Important Note: We do not do technical support, nor consulting and don't answer personal questions per email.
Please post your question on the RL Discord, Reddit or Stack Overflow in that case.

🚀 Feature

CheckpointCallback should also call model.save_replay_buffer(...) when it is applicable to do so.

Motivation

The primary motivation for using the CheckpointCallback is to save the trained model periodically to be able to resume training if something goes wrong, or to continue training for better convergence. But, in the context of most off-policy RL algorithms, if the replay buffer isn't also saved, there is a significant lack of continuity in training performance (see #326).

Pitch

The CheckpointCallback._on_step method needs to add something like the following after the linked line (modulo some changes to __init__ to get a suffix for the replay buffer):

if hasattr(self.model, "replay_buffer") and self.model.replay_buffer is not None:
    self.model.save_replay_buffer(replay_buffer_path)

Alternatives

The alternative is to have the users define their own CheckpointCallback that does this (as I am doing now) but I thought it made sense to have this feature default in the library.

### Checklist

• I have checked that there is no similar issue in the repo (required)

[araffin]

Hello,
this sounds like a reasonable feature, but it should be deactivated by default (saving buffer takes time and space).
Feel free to submit a PR ;) (and don't forget to read the contributing guide)

Could you also add an option to save VecNormalize statistics? (see DLR-RM/rl-baselines3-zoo#278)

[anand-bala]

Sorry for the delay in getting back.

Could you also add an option to save VecNormalize statistics? (see DLR-RM/rl-baselines3-zoo#278)

Would this work the same as SaveVecNormalizeCallback? That is to say, do you want me to merge CheckpointCallback with SaveVecNormalizeCallback while adding the replay buffer functionality?

[araffin]

That is to say, do you want me to merge CheckpointCallback with SaveVecNormalizeCallback while adding the replay buffer functionality?

exactly =)