feat: add custom x_axis_name

instadeepai · Dec 14, 2023 · 707b7a1 · 707b7a1
1 parent 1821702
commit 707b7a1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/marl_eval/plotting_tools/plotting.py b/marl_eval/plotting_tools/plotting.py
@@ -359,6 +359,7 @@ def plot_single_task(
     legend_map: Optional[Dict[str, str]] = None,
     run_times: Optional[Dict[str, float]] = None,
     base_line_values: Optional[Dict[str, Dict[str, float]]] = None,
+    x_axis_name: Optional[str] = None,
 ) -> Figure:
     """Produces aggregated plot for a single task in an environment.
 
@@ -409,6 +410,8 @@ def plot_single_task(
     if run_times is not None:
         run_times = {algo.upper(): value for algo, value in run_times.items()}
         xlabel = "Time (Minutes)"
+    if x_axis_name is not None:
+        xlabel = x_axis_name
 
     if base_line_values is not None:
         task_baseline: Dict[str, float] = base_line_values[task_name]