Fix: issue 198 return metrics after initialisation of repertoires

examples/aurora.ipynb

@@ -405,7 +405,7 @@
     "\n",
     "# init step of the aurora algorithm\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state, aurora_extra_info = aurora.init(\n",
+    "repertoire, emitter_state, init_metrics, aurora_extra_info = aurora.init(\n",
     "    init_variables,\n",
     "    aurora_extra_info,\n",
     "    jnp.asarray(l_value_init),\n",

examples/cmame.ipynb

@@ -260,7 +260,7 @@
    "outputs": [],
    "source": [
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(initial_population, centroids, subkey)"
+    "repertoire, emitter_state, init_metrics = map_elites.init(initial_population, centroids, subkey)"
    ]
   },
   {

examples/cmamega.ipynb

@@ -235,7 +235,7 @@
    "outputs": [],
    "source": [
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(initial_population, centroids, subkey)"
+    "repertoire, emitter_state, init_metrics = map_elites.init(initial_population, centroids, subkey)"
    ]
   },
   {

examples/dcrlme.ipynb

@@ -352,7 +352,7 @@
     "\n",
     "# compute initial repertoire\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(\n",
+    "repertoire, emitter_state, init_metrics = map_elites.init(\n",
     "    init_params, centroids, subkey\n",
     ")"
    ]

examples/distributed_mapelites.ipynb

@@ -340,7 +340,7 @@
     ")\n",
     "\n",
     "# get initial elements\n",
-    "repertoire, emitter_state = map_elites.get_distributed_init_fn(\n",
+    "repertoire, emitter_state, init_metrics = map_elites.get_distributed_init_fn(\n",
     "    centroids=centroids,\n",
     "    devices=devices,\n",
     ")(genotypes=init_variables, key=keys)"

examples/jumanji_snake.ipynb

@@ -375,7 +375,7 @@
     "    )\n",
     "\n",
     "    key, subkey = jax.random.split(key)\n",
-    "    repertoire, emitter_state = algo_instance.init(\n",
+    "    repertoire, emitter_state, init_metrics = algo_instance.init(\n",
     "        init_variables, population_size, subkey\n",
     "    )\n",
     "\n",
@@ -402,7 +402,7 @@
     "\n",
     "    # Compute initial repertoire and emitter state\n",
     "    key, subkey = jax.random.split(key)\n",
-    "    repertoire, emitter_state = algo_instance.init(init_variables, centroids, subkey)"
+    "    repertoire, emitter_state, init_metrics = algo_instance.init(init_variables, centroids, subkey)"
    ]
   },
   {

examples/mapelites.ipynb

@@ -288,7 +288,7 @@
     "\n",
     "# Compute initial repertoire and emitter state\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(init_variables, centroids, subkey)"
+    "repertoire, emitter_state, init_metrics = map_elites.init(init_variables, centroids, subkey)"
    ]
   },
   {

examples/me_sac_pbt.ipynb

@@ -330,7 +330,7 @@
     ")(training_states)\n",
     "\n",
     "# initialize map-elites\n",
-    "repertoire, emitter_state = map_elites.get_distributed_init_fn(\n",
+    "repertoire, emitter_state, init_metrics = map_elites.get_distributed_init_fn(\n",
     "    devices=devices, centroids=centroids\n",
     ")(genotypes=training_states, key=keys)"
    ]

examples/me_td3_pbt.ipynb

@@ -332,7 +332,7 @@
     ")(training_states)\n",
     "\n",
     "# initialize map-elites\n",
-    "repertoire, emitter_state = map_elites.get_distributed_init_fn(\n",
+    "repertoire, emitter_state, init_metrics = map_elites.get_distributed_init_fn(\n",
     "    devices=devices, centroids=centroids\n",
     ")(genotypes=training_states, key=keys)"
    ]

examples/mees.ipynb

@@ -360,7 +360,7 @@
     "\n",
     "# compute initial repertoire\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(\n",
+    "repertoire, emitter_state, init_metrics = map_elites.init(\n",
     "    init_variables, centroids, subkey\n",
     ")"
    ]

examples/mels.ipynb

@@ -304,7 +304,7 @@
     "\n",
     "# Compute initial repertoire and emitter state\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = mels.init(init_variables, centroids, subkey)"
+    "repertoire, emitter_state, init_metrics = mels.init(init_variables, centroids, subkey)"
    ]
   },
   {

examples/mome.ipynb

@@ -299,7 +299,7 @@
    "outputs": [],
    "source": [
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = mome.init(\n",
+    "repertoire, emitter_state, init_metrics = mome.init(\n",
     "    genotypes,\n",
     "    centroids,\n",
     "    pareto_front_max_length,\n",

examples/nsga2_spea2.ipynb

@@ -232,7 +232,7 @@
     "\n",
     "# init nsga2\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = nsga2.init(\n",
+    "repertoire, emitter_state, init_metrics = nsga2.init(\n",
     "    genotypes,\n",
     "    population_size,\n",
     "    subkey\n",
@@ -296,7 +296,7 @@
     "\n",
     "# init spea2\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = spea2.init(\n",
+    "repertoire, emitter_state, init_metrics = spea2.init(\n",
     "    genotypes,\n",
     "    population_size,\n",
     "    num_neighbours,\n",

examples/omgmega.ipynb

@@ -228,7 +228,7 @@
    "outputs": [],
    "source": [
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(initial_population, centroids, subkey)"
+    "repertoire, emitter_state, init_metrics = map_elites.init(initial_population, centroids, subkey)"
    ]
   },
   {

examples/pga_aurora.ipynb

@@ -467,7 +467,7 @@
     "\n",
     "# init step of the aurora algorithm\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state, aurora_extra_info = aurora.init(\n",
+    "repertoire, emitter_state, init_metrics, aurora_extra_info = aurora.init(\n",
     "    init_variables,\n",
     "    aurora_extra_info,\n",
     "    jnp.asarray(l_value_init),\n",

examples/pgame.ipynb

@@ -325,7 +325,7 @@
     "\n",
     "# compute initial repertoire\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(\n",
+    "repertoire, emitter_state, init_metrics = map_elites.init(\n",
     "    init_variables, centroids, subkey\n",
     ")"
    ]

examples/qdpg.ipynb

@@ -373,7 +373,7 @@
     "\n",
     "# compute initial repertoire\n",
     "key, subkey = jax.random.split(key)\n",
-    "repertoire, emitter_state = map_elites.init(\n",
+    "repertoire, emitter_state, init_metrics = map_elites.init(\n",
     "    init_variables, centroids, subkey\n",
     ")"
    ]

qdax/baselines/genetic_algorithm.py

@@ -39,7 +39,7 @@ def __init__(
     @partial(jax.jit, static_argnames=("self", "population_size"))
     def init(
         self, genotypes: Genotype, population_size: int, key: RNGKey
-    ) -> Tuple[GARepertoire, Optional[EmitterState]]:
+    ) -> Tuple[GARepertoire, Optional[EmitterState], Metrics]:
         """Initialize a GARepertoire with an initial population of genotypes.
 
         Args:
@@ -73,7 +73,10 @@ def init(
             extra_scores=extra_scores,
         )
 
-        return repertoire, emitter_state
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics
 
     @partial(jax.jit, static_argnames=("self",))
     def update(

qdax/baselines/nsga2.py

@@ -13,7 +13,7 @@
 from qdax.baselines.genetic_algorithm import GeneticAlgorithm
 from qdax.core.containers.nsga2_repertoire import NSGA2Repertoire
 from qdax.core.emitters.emitter import EmitterState
-from qdax.custom_types import Genotype, RNGKey
+from qdax.custom_types import Genotype, Metrics, RNGKey
 
 
 class NSGA2(GeneticAlgorithm):
@@ -29,7 +29,7 @@ class NSGA2(GeneticAlgorithm):
     @partial(jax.jit, static_argnames=("self", "population_size"))
     def init(
         self, genotypes: Genotype, population_size: int, key: RNGKey
-    ) -> Tuple[NSGA2Repertoire, Optional[EmitterState]]:
+    ) -> Tuple[NSGA2Repertoire, Optional[EmitterState], Metrics]:
 
         # score initial genotypes
         key, subkey = jax.random.split(key)
@@ -62,4 +62,7 @@ def init(
             extra_scores=extra_scores,
         )
 
-        return repertoire, emitter_state
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics

qdax/baselines/spea2.py

@@ -15,7 +15,7 @@
 from qdax.baselines.genetic_algorithm import GeneticAlgorithm
 from qdax.core.containers.spea2_repertoire import SPEA2Repertoire
 from qdax.core.emitters.emitter import EmitterState
-from qdax.custom_types import Genotype, RNGKey
+from qdax.custom_types import Genotype, Metrics, RNGKey
 
 
 class SPEA2(GeneticAlgorithm):
@@ -37,7 +37,7 @@ def init(
         population_size: int,
         num_neighbours: int,
         key: RNGKey,
-    ) -> Tuple[SPEA2Repertoire, Optional[EmitterState]]:
+    ) -> Tuple[SPEA2Repertoire, Optional[EmitterState], Metrics]:
 
         # score initial genotypes
         key, subkey = jax.random.split(key)
@@ -71,4 +71,7 @@ def init(
             extra_scores=extra_scores,
         )
 
-        return repertoire, emitter_state
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics

qdax/core/aurora.py

@@ -123,7 +123,9 @@ def init(
         l_value: jnp.ndarray,
         max_size: int,
         key: RNGKey,
-    ) -> Tuple[UnstructuredRepertoire, Optional[EmitterState], AuroraExtraInfo]:
+    ) -> Tuple[
+        UnstructuredRepertoire, Optional[EmitterState], Metrics, AuroraExtraInfo
+    ]:
         """Initialize an unstructured repertoire with an initial population of
         genotypes. Also performs the first training of the AURORA encoder.
 
@@ -174,7 +176,10 @@ def init(
             repertoire, aurora_extra_info.model_params, iteration=0, key=key
         )
 
-        return repertoire, emitter_state, updated_aurora_extra_info
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics, updated_aurora_extra_info
 
     @partial(jax.jit, static_argnames=("self",))
     def update(

qdax/core/distributed_map_elites.py

@@ -21,7 +21,7 @@ def init(
         genotypes: Genotype,
         centroids: Centroid,
         key: RNGKey,
-    ) -> Tuple[MapElitesRepertoire, Optional[EmitterState]]:
+    ) -> Tuple[MapElitesRepertoire, Optional[EmitterState], Metrics]:
         """
         Initialize a Map-Elites repertoire with an initial population of genotypes.
         Requires the definition of centroids that can be computed with any method
@@ -80,8 +80,10 @@ def init(
             descriptors=descriptors,
             extra_scores=extra_scores,
         )
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
 
-        return repertoire, emitter_state
+        return repertoire, emitter_state, metrics
 
     @partial(jax.jit, static_argnames=("self",))
     def update(

qdax/core/map_elites.py

@@ -56,7 +56,7 @@ def init(
         genotypes: Genotype,
         centroids: Centroid,
         key: RNGKey,
-    ) -> Tuple[MapElitesRepertoire, Optional[EmitterState]]:
+    ) -> Tuple[MapElitesRepertoire, Optional[EmitterState], Metrics]:
         """
         Initialize a Map-Elites repertoire with an initial population of genotypes.
         Requires the definition of centroids that can be computed with any method
@@ -96,7 +96,10 @@ def init(
             extra_scores=extra_scores,
         )
 
-        return repertoire, emitter_state
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics
 
     @partial(jax.jit, static_argnames=("self",))
     def update(

qdax/core/mels.py

@@ -59,7 +59,7 @@ def init(
         genotypes: Genotype,
         centroids: Centroid,
         key: RNGKey,
-    ) -> Tuple[MELSRepertoire, Optional[EmitterState]]:
+    ) -> Tuple[MELSRepertoire, Optional[EmitterState], Metrics]:
         """Initialize a MAP-Elites Low-Spread repertoire with an initial
         population of genotypes. Requires the definition of centroids that can
         be computed with any method such as CVT or Euclidean mapping.
@@ -107,4 +107,8 @@ def init(
             descriptors=descriptors,
             extra_scores=extra_scores,
         )
-        return repertoire, emitter_state
+
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics

qdax/core/mome.py

@@ -9,7 +9,7 @@
 from qdax.core.containers.mome_repertoire import MOMERepertoire
 from qdax.core.emitters.emitter import EmitterState
 from qdax.core.map_elites import MAPElites
-from qdax.custom_types import Centroid, RNGKey
+from qdax.custom_types import Centroid, Metrics, RNGKey
 
 
 class MOME(MAPElites):
@@ -27,7 +27,7 @@ def init(
         centroids: Centroid,
         pareto_front_max_length: int,
         key: RNGKey,
-    ) -> Tuple[MOMERepertoire, Optional[EmitterState]]:
+    ) -> Tuple[MOMERepertoire, Optional[EmitterState], Metrics]:
         """Initialize a MOME grid with an initial population of genotypes. Requires
         the definition of centroids that can be computed with any method such as
         CVT or Euclidean mapping.
@@ -68,4 +68,7 @@ def init(
             extra_scores=extra_scores,
         )
 
-        return repertoire, emitter_state
+        # calculate the initial metrics
+        metrics = self._metrics_function(repertoire)
+
+        return repertoire, emitter_state, metrics

tests/baselines_test/cmame_test.py

@@ -112,7 +112,9 @@ def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:
     )
 
     key, subkey = jax.random.split(key)
-    repertoire, emitter_state = map_elites.init(initial_population, centroids, subkey)
+    repertoire, emitter_state, init_metrics = map_elites.init(
+        initial_population, centroids, subkey
+    )
 
     (
         repertoire,

tests/baselines_test/cmamega_test.py

@@ -124,7 +124,9 @@ def metrics_fn(repertoire: MapElitesRepertoire) -> Dict[str, jnp.ndarray]:
     )
 
     key, subkey = jax.random.split(key)
-    repertoire, emitter_state = map_elites.init(initial_population, centroids, subkey)
+    repertoire, emitter_state, init_metrics = map_elites.init(
+        initial_population, centroids, subkey
+    )
 
     (
         repertoire,

tests/baselines_test/dcrlme_test.py

@@ -197,7 +197,9 @@ def play_step_fn(
 
     # compute initial repertoire
     key, subkey = jax.random.split(key)
-    repertoire, emitter_state = map_elites.init(init_params, centroids, subkey)
+    repertoire, emitter_state, init_metrics = map_elites.init(
+        init_params, centroids, subkey
+    )
 
     @jax.jit
     def update_scan_fn(carry: Any, key: RNGKey) -> Any: