Feature/delete old ckpts

howto/logs_and_checkpoints.md

@@ -169,13 +169,15 @@ By default the checkpointing is enabled with the following settings:
 every: 100
 resume_from: null
 save_last: True
+keep_last: 5
 ```
 
 meaning that:
 
 * `every` is the number of policy steps (number of steps played in the environment, e.g. if one has 2 processes with 4 environments per process then the policy steps are 2*4=8) between two consecutive checkpointing operations. For more info about the policy steps, check the [Work with Steps Tutorial](./work_with_steps.md).
 * `resume_from` is the path of the checkpoint to resume from. If `null`, then the checkpointing is not resumed.
 * `save_last` is a boolean flag that enables/disables the saving of the last checkpoint.
+* `keep_last` is the number of checkpoints you want to keep during the experiment. If `null`, all the checkpoints are kept.
 
 > **Note**
 >

sheeprl/configs/checkpoint/default.yaml

@@ -1,3 +1,4 @@
 every: 100
 resume_from: null
 save_last: True
+keep_last: 5

sheeprl/configs/fabric/default.yaml

@@ -6,3 +6,4 @@ accelerator: "cpu"
 precision: "32-true"
 callbacks:
   - _target_: sheeprl.utils.callback.CheckpointCallback
+    keep_last: "${checkpoint.keep_last}"

sheeprl/configs/model_manager/dreamer_v2.yaml

@@ -18,4 +18,4 @@ models:
   target_critic:
     model_name: "${exp_name}_target_critic"
     description: "DreamerV2 Target Critic used in ${env.id} Environment"
-    tags: {}
+    tags: {}

sheeprl/utils/callback.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import os
+import pathlib
 from typing import Any, Dict, Optional, Sequence, Union
 
 from lightning.fabric import Fabric
@@ -22,6 +24,9 @@ class CheckpointCallback:
     When the buffer is added to the state of the checkpoint, it is assumed that the episode is truncated.
     """
 
+    def __init__(self, keep_last: int | None = None) -> None:
+        self.keep_last = keep_last
+
     def on_checkpoint_coupled(
         self,
         fabric: Fabric,
@@ -47,6 +52,8 @@ def on_checkpoint_coupled(
         fabric.save(ckpt_path, state)
         if replay_buffer is not None:
             self._experiment_consistent_rb(replay_buffer, rb_state)
+        if self.keep_last:
+            self._delete_old_checkpoints(pathlib.Path(ckpt_path).parent)
 
     def on_checkpoint_player(
         self,
@@ -64,6 +71,8 @@ def on_checkpoint_player(
         fabric.save(ckpt_path, state)
         if replay_buffer is not None:
             self._experiment_consistent_rb(replay_buffer, rb_state)
+        if self.keep_last:
+            self._delete_old_checkpoints(pathlib.Path(ckpt_path).parent)
 
     def on_checkpoint_trainer(
         self, fabric: Fabric, player_trainer_collective: TorchCollective, state: Dict[str, Any], ckpt_path: str
@@ -128,3 +137,9 @@ def _experiment_consistent_rb(
         elif isinstance(rb, EpisodeBuffer):
             # reinsert the open episodes to continue the training
             rb._open_episodes = state
+
+    def _delete_old_checkpoints(self, ckpt_folder: str | pathlib.Path):
+        ckpts = list(sorted(ckpt_folder.glob("*.ckpt"), key=os.path.getmtime))
+        if len(ckpts) > self.keep_last:
+            to_delete = ckpts[: -self.keep_last]
+            [f.unlink() for f in to_delete]