Add extra_scores to mome repertoire as well

qdax/core/containers/mome_repertoire.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Any, Tuple
+from typing import Any, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -17,6 +17,7 @@
 from qdax.types import (
     Centroid,
     Descriptor,
+    ExtraScores,
     Fitness,
     Genotype,
     Mask,
@@ -259,6 +260,7 @@ def add(
         batch_of_genotypes: Genotype,
         batch_of_descriptors: Descriptor,
         batch_of_fitnesses: Fitness,
+        batch_of_extra_scores: Optional[ExtraScores] = None,
     ) -> MOMERepertoire:
         """Insert a batch of elements in the repertoire.
 
@@ -274,6 +276,8 @@ def add(
                 trying to add to the repertoire.
             batch_of_fitnesses: the fitnesses of the genotypes we are trying
                 to add to the repertoire.
+            batch_of_extra_scores: unused tree that contains the extra_scores of
+                aforementioned genotypes.
 
         Returns:
             The updated repertoire with potential new individuals.
@@ -355,6 +359,7 @@ def init(  # type: ignore
         descriptors: Descriptor,
         centroids: Centroid,
         pareto_front_max_length: int,
+        extra_scores: Optional[ExtraScores] = None,
     ) -> MOMERepertoire:
         """
         Initialize a Multi Objective Map-Elites repertoire with an initial population
@@ -373,6 +378,7 @@ def init(  # type: ignore
                 of shape (batch_size, num_descriptors)
             centroids: tesselation centroids of shape (batch_size, num_descriptors)
             pareto_front_max_length: maximum size of the pareto fronts
+            extra_scores: unused extra_scores of the initial genotypes
 
         Returns:
             An initialized MAP-Elite repertoire
@@ -410,7 +416,7 @@ def init(  # type: ignore
         )
 
         # add first batch of individuals in the repertoire
-        new_repertoire = repertoire.add(genotypes, descriptors, fitnesses)
+        new_repertoire = repertoire.add(genotypes, descriptors, fitnesses, extra_scores)
 
         return new_repertoire  # type: ignore
 

qdax/core/mome.py

@@ -55,6 +55,7 @@ def init(
             descriptors=descriptors,
             centroids=centroids,
             pareto_front_max_length=pareto_front_max_length,
+            extra_scores=extra_scores,
         )
 
         # get initial state of the emitter