Merge branch 'develop' into feat/store-ppo-experience

BDEvan5 · Dec 12, 2023 · 5dc8ff1 · 5dc8ff1
2 parents ded3660 + f88e1f5
commit 5dc8ff1
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/mava/utils/logger_tools.py b/mava/utils/logger_tools.py
@@ -15,6 +15,7 @@
 import json
 import logging
 import os
+import time
 from datetime import datetime
 from typing import Dict, Optional
 
@@ -216,13 +217,20 @@ def write(
             evaluation_step (int): the evaluation step
         """
 
+        current_time = time.time()
+
+        # This will ensure the first logged time is 0, which avoids taking compilation into account
+        # when plotting downstream.
+        if evaluation_step == 0:
+            self.start_time = current_time
+
         logging_prefix, *metric_key = key.split("/")
         metric_key = "/".join(metric_key)
 
         metrics = {metric_key: [value]}
 
         if logging_prefix == "evaluator":
-            step_metrics = {"step_count": timestep}
+            step_metrics = {"step_count": timestep, "elapsed_time": current_time - self.start_time}
             step_metrics.update(metrics)  # type: ignore
             step_str = f"step_{evaluation_step}"
             if step_str in self.run_data: