Add a callback to log raw stats

utils/callbacks.py

@@ -10,6 +10,7 @@
 from sb3_contrib import TQC
 from stable_baselines3 import SAC
 from stable_baselines3.common.callbacks import BaseCallback, EvalCallback
+from stable_baselines3.common.logger import TensorBoardOutputFormat
 from stable_baselines3.common.vec_env import VecEnv
 
 
@@ -193,3 +194,36 @@ def _on_training_end(self) -> None:
             if self.verbose > 0:
                 print("Waiting for training thread to terminate")
             self.process.join()
+
+
+class RawStatisticsCallback(BaseCallback):
+    """
+    Callback used for logging raw episode data (return and episode length).
+    """
+
+    def __init__(self, verbose=0):
+        super(RawStatisticsCallback, self).__init__(verbose)
+        # Custom counter to reports stats
+        # (and avoid reporting multiple values for the same step)
+        self._timesteps_counter = 0
+        self._tensorboard_writer = None
+
+    def _init_callback(self) -> None:
+        # Retrieve tensorboard writer to not flood the logger output
+        for out_format in self.logger.output_formats:
+            if isinstance(out_format, TensorBoardOutputFormat):
+                self._tensorboard_writer = out_format
+        assert self._tensorboard_writer is not None, "You must activate tensorboard logging when using RawStatisticsCallback"
+
+    def _on_step(self) -> bool:
+        for info in self.locals["infos"]:
+            if "episode" in info:
+                logger_dict = {
+                    "raw/rollouts/episodic_return": info["episode"]["r"],
+                    "raw/rollouts/episodic_length": info["episode"]["l"],
+                }
+                exclude_dict = {key: None for key in logger_dict.keys()}
+                self._timesteps_counter += info["episode"]["l"]
+                self._tensorboard_writer.write(logger_dict, exclude_dict, self._timesteps_counter)
+
+        return True