chore: Custom x label for plotting (#233)

Author: TemplierPaul

qdax/utils/plotting.py

@@ -222,6 +222,7 @@ def plot_map_elites_results(
     repertoire: MapElitesRepertoire,
     min_descriptor: jnp.ndarray,
     max_descriptor: jnp.ndarray,
+    x_label: str = "Environment steps",
 ) -> Tuple[Optional[Figure], Axes]:
     """Plots three usual QD metrics, namely the coverage, the maximum fitness
     and the QD-score, along the number of environment steps. This function also
@@ -235,6 +236,7 @@ def plot_map_elites_results(
         repertoire: the final repertoire obtained.
         min_descriptor: the minimal possible values for the descriptor.
         max_descriptor: the maximal possible values for the descriptor.
+        x_label: label for the x axis, defaults to environment steps
 
     Returns:
         A figure and axes with the plots of the metrics and visualisation of the grid.
@@ -259,19 +261,19 @@ def plot_map_elites_results(
     # env_steps = jnp.arange(num_iterations) * episode_length * batch_size
 
     axes[0].plot(env_steps, metrics["coverage"])
-    axes[0].set_xlabel("Environment steps")
+    axes[0].set_xlabel(x_label)
     axes[0].set_ylabel("Coverage in %")
     axes[0].set_title("Coverage evolution during training")
     axes[0].set_aspect(0.95 / axes[0].get_data_ratio(), adjustable="box")
 
     axes[1].plot(env_steps, metrics["max_fitness"])
-    axes[1].set_xlabel("Environment steps")
+    axes[1].set_xlabel(x_label)
     axes[1].set_ylabel("Maximum fitness")
     axes[1].set_title("Maximum fitness evolution during training")
     axes[1].set_aspect(0.95 / axes[1].get_data_ratio(), adjustable="box")
 
     axes[2].plot(env_steps, metrics["qd_score"])
-    axes[2].set_xlabel("Environment steps")
+    axes[2].set_xlabel(x_label)
     axes[2].set_ylabel("QD Score")
     axes[2].set_title("QD Score evolution during training")
     axes[2].set_aspect(0.95 / axes[2].get_data_ratio(), adjustable="box")