From 2a4350da5606535b2eec31ef09a43e4971b3037c Mon Sep 17 00:00:00 2001
From: Yoann Poupart <66315201+Xmaster6y@users.noreply.github.com>
Date: Wed, 29 Oct 2025 12:56:01 -0400
Subject: [PATCH 1/2] new methods

---
 benchmarl/experiment/callback.py   | 18 +++++++++++++++++-
 benchmarl/experiment/experiment.py |  2 ++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/benchmarl/experiment/callback.py b/benchmarl/experiment/callback.py
index e04900e5..3ace685f 100644
--- a/benchmarl/experiment/callback.py
+++ b/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
 
@@ -28,6 +28,10 @@ def on_setup(self):
         """A callback called atexperiment 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):
         """
         A callback called at the end of every collection step.
@@ -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)
diff --git a/benchmarl/experiment/experiment.py b/benchmarl/experiment/experiment.py
index 5dc78d55..df49737a 100644
--- a/benchmarl/experiment/experiment.py
+++ b/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"""

From 61616a7957deef4e048bb8d926f0a2e27c62eb95 Mon Sep 17 00:00:00 2001
From: Yoann Poupart <66315201+Xmaster6y@users.noreply.github.com>
Date: Wed, 29 Oct 2025 12:58:31 -0400
Subject: [PATCH 2/2] typo

---
 benchmarl/experiment/callback.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmarl/experiment/callback.py b/benchmarl/experiment/callback.py
index 3ace685f..1c24959f 100644
--- a/benchmarl/experiment/callback.py
+++ b/benchmarl/experiment/callback.py
@@ -25,7 +25,7 @@ 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]):