[Feature] New Callback methods for state_dict

benchmarl/experiment/callback.py

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from typing import List
+from typing import Any, Dict, List
 
 from tensordict import TensorDictBase
 
@@ -25,7 +25,11 @@ def __init__(self):
         self.experiment = None
 
     def on_setup(self):
-        """A callback called atexperiment setup."""
+        """A callback called at experiment setup."""
+        pass
+
+    def on_load_state_dict(self, state_dict: Dict[str, Any]):
+        """A callback called at state_dict load."""
         pass
 
     def on_batch_collected(self, batch: TensorDictBase):
@@ -73,6 +77,10 @@ def on_evaluation_end(self, rollouts: List[TensorDictBase]):
         """
         pass
 
+    def on_state_dict(self, state_dict: Dict[str, Any]):
+        """A callback called at state_dict save."""
+        pass
+
 
 class CallbackNotifier:
     def __init__(self, experiment, callbacks: List[Callback]):
@@ -84,6 +92,10 @@ def _on_setup(self):
         for callback in self.callbacks:
             callback.on_setup()
 
+    def _on_load_state_dict(self, state_dict: Dict[str, Any]):
+        for callback in self.callbacks:
+            callback.on_load_state_dict(state_dict)
+
     def _on_batch_collected(self, batch: TensorDictBase):
         for callback in self.callbacks:
             callback.on_batch_collected(batch)
@@ -106,3 +118,7 @@ def _on_train_end(self, training_td: TensorDictBase, group: str):
     def _on_evaluation_end(self, rollouts: List[TensorDictBase]):
         for callback in self.callbacks:
             callback.on_evaluation_end(rollouts)
+
+    def _on_state_dict(self, state_dict: Dict[str, Any]):
+        for callback in self.callbacks:
+            callback.on_state_dict(state_dict)

benchmarl/experiment/experiment.py

@@ -972,6 +972,7 @@ def state_dict(self) -> OrderedDict:
         )
         if not self.config.collect_with_grad:
             state_dict.update({"collector": self.collector.state_dict()})
+        self._on_state_dict(state_dict)
         return state_dict
 
     def load_state_dict(self, state_dict: Dict) -> None:
@@ -993,6 +994,7 @@ def load_state_dict(self, state_dict: Dict) -> None:
         self.total_frames = state_dict["state"]["total_frames"]
         self.n_iters_performed = state_dict["state"]["n_iters_performed"]
         self.mean_return = state_dict["state"]["mean_return"]
+        self._on_load_state_dict(state_dict)
 
     def _save_experiment(self) -> None:
         """Checkpoint trainer"""