Merge pull request #46 from instadeepai/feat/add-json-logger

Feat/add json logger

marl_eval/json_tools/__init__.py

@@ -14,3 +14,4 @@
 # limitations under the License.
 
 """JSON tools for data preprocessing."""
+from .json_logger import JsonLogger

marl_eval/json_tools/json_logger.py

@@ -0,0 +1,115 @@
+# python3
+# Copyright 2022 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import time
+from typing import Dict, Optional
+
+
+class JsonLogger:
+    """Logger to create JSON files for reporting experiment results.
+
+    This logger follows the suggested marl-eval protocol and was adapted
+    from the implementation found in BenchMARL which can be viewed at:
+    https://tinyurl.com/2t6fy548
+
+    Args:
+        path (str): folder path for saving the `metrics.json` file.
+        algorithm_name (str): algorithm name e.g PPO.
+        task_name (str): task name e.g 3s5z (for SMAC).
+        environment_name (str): environment name e.g SMAC.
+        seed (int): random seed of the experiment.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        algorithm_name: str,
+        task_name: str,
+        environment_name: str,
+        seed: int,
+    ):
+        """Initialises the JsonLogger and creates a metrics file if it doesn't exist."""
+        self.file_path = f"{path}/metrics.json"
+        self.run_data: Dict = {"absolute_metrics": {}}
+
+        # If the file already exists, load it
+        if os.path.isfile(self.file_path):
+            with open(self.file_path) as f:
+                data = json.load(f)
+        else:
+            # Create the logging directory if it doesn't exist
+            os.makedirs(path, exist_ok=True)
+            data = {}
+
+        # Merge the existing data with the new data
+        self.data = data
+        if environment_name not in self.data:
+            self.data[environment_name] = {}
+        if task_name not in self.data[environment_name]:
+            self.data[environment_name][task_name] = {}
+        if algorithm_name not in self.data[environment_name][task_name]:
+            self.data[environment_name][task_name][algorithm_name] = {}
+        self.data[environment_name][task_name][algorithm_name][
+            f"seed_{seed}"
+        ] = self.run_data
+
+        with open(self.file_path, "w") as f:
+            json.dump(self.data, f, indent=4)
+
+    def write(
+        self,
+        timestep: int,
+        key: str,
+        value: float,
+        evaluation_step: Optional[int] = None,
+        is_absolute_metric: bool = False,
+    ) -> None:
+        """Writes a step to the json reporting file.
+
+        Args:
+            timestep (int): the current environment timestep.
+            key (str): the name of the metric to be logged.
+            value (str): the value of the metric to be logged.
+            evaluation_step (int): the number of evaluations already run.
+            is_absolute_metric (bool): whether the metric being logged is
+                an absolute metric.
+        """
+
+        current_time = time.time()
+
+        # This will ensure the first logged time is 0, which avoids taking compilation
+        # into account for jax-based systems when plotting downstream.
+        if evaluation_step == 0:
+            self.start_time = current_time
+
+        metrics: Dict = {key: [value]}
+
+        if is_absolute_metric:
+            self.run_data["absolute_metrics"].update(metrics)
+        else:
+            step_metrics = {  # type: ignore
+                "step_count": timestep,
+                "elapsed_time": current_time - self.start_time,
+            } | metrics
+            step_str = f"step_{evaluation_step}"
+            if step_str in self.run_data:
+                self.run_data[step_str].update(step_metrics)
+            else:
+                self.run_data[step_str] = step_metrics
+
+        with open(self.file_path, "w") as f:
+            json.dump(self.data, f, indent=4)

requirements/requirements.txt

@@ -1,5 +1,5 @@
 colorcet==3.0.0
-matplotlib==3.5.3
+matplotlib>=3.5.3
 numpy
 rliable
 seaborn