Containers¶

class DominatedNoveltyRepertoire(GARepertoire):
    """Repertoire that keeps the individuals with highest dominated novelty.

    Args:
        genotypes: population genotypes with shape (population_size, ...)
        fitnesses: population fitnesses with shape (population_size, fitness_dim)
        descriptors: population descriptors with shape (population_size, D)
        k: number of neighbors for novelty and dominated novelty
        extra_scores: extra scores resulting from the evaluation of the genotypes
        keys_extra_scores: keys of the extra scores to store in the repertoire
    """

    descriptors: Descriptor
    k: int = flax.struct.field(pytree_node=False)

    @jax.jit
    def add(  # type: ignore
        self,
        batch_of_genotypes: Genotype,
        batch_of_descriptors: Descriptor,
        batch_of_fitnesses: Fitness,
        batch_of_extra_scores: Optional[ExtraScores] = None,
    ) -> DominatedNoveltyRepertoire:
        """Add a batch and keep the top individuals by dominated novelty.

        Parents and offsprings are gathered and only the population_size
        best according to dominated novelty are kept.
        """

        if batch_of_extra_scores is None:
            batch_of_extra_scores = {}

        filtered_batch_of_extra_scores = self.filter_extra_scores(batch_of_extra_scores)

        batch_of_fitnesses = jnp.reshape(
            batch_of_fitnesses, (batch_of_fitnesses.shape[0], 1)
        )

        # Gather candidates
        candidates_genotypes = jax.tree.map(
            lambda x, y: jnp.concatenate((x, y), axis=0),
            self.genotypes,
            batch_of_genotypes,
        )
        candidates_fitnesses = jnp.concatenate(
            (self.fitnesses, batch_of_fitnesses), axis=0
        )
        candidates_descriptors = jnp.concatenate(
            (self.descriptors, batch_of_descriptors), axis=0
        )
        candidates_extra_scores = jax.tree.map(
            lambda x, y: jnp.concatenate((x, y), axis=0),
            self.extra_scores,
            filtered_batch_of_extra_scores,
        )

        # Compute dominated novelty
        _, dominated_novelty = _novelty_and_dominated_novelty(
            fitness=candidates_fitnesses[:, 0],
            descriptor=candidates_descriptors,
            novelty_k=self.k,
            dominated_novelty_k=self.k,
        )

        # Use dominated novelty as meta-fitness, invalid individuals get -inf
        valid = candidates_fitnesses[:, 0] != -jnp.inf
        meta_fitness = jnp.where(valid, dominated_novelty, -jnp.inf)

        # Select survivors
        indices = jnp.argsort(meta_fitness)[::-1]
        survivor_indices = indices[: self.size]

        new_genotypes = jax.tree.map(
            lambda x: x[survivor_indices], candidates_genotypes
        )
        new_fitnesses = candidates_fitnesses[survivor_indices]
        new_descriptors = candidates_descriptors[survivor_indices]
        new_extra_scores = jax.tree.map(
            lambda x: x[survivor_indices], candidates_extra_scores
        )

        return self.replace(  # type: ignore
            genotypes=new_genotypes,
            fitnesses=new_fitnesses,
            descriptors=new_descriptors,
            extra_scores=new_extra_scores,
        )

    @classmethod
    def init(  # type: ignore
        cls,
        genotypes: Genotype,
        fitnesses: Fitness,
        descriptors: Descriptor,
        population_size: int,
        k: int,
        *args,
        extra_scores: Optional[ExtraScores] = None,
        keys_extra_scores: Tuple[str, ...] = (),
        **kwargs,
    ) -> DominatedNoveltyRepertoire:
        """Initialize the repertoire and add the first batch.

        Args:
            genotypes: first batch of genotypes (batch_size, ...)
            fitnesses: fitnesses of shape (batch_size, fitness_dim)
            descriptors: descriptors of shape (batch_size, num_descriptors)
            population_size: maximum number of individuals kept
            k: number of neighbors for novelty metrics
            extra_scores: extra scores of the first batch
            keys_extra_scores: keys of extra scores to store
        """

        if extra_scores is None:
            extra_scores = {}

        # retrieve one genotype and one extra score prototype
        first_genotype = jax.tree.map(lambda x: x[0], genotypes)
        first_extra_scores = jax.tree.map(lambda x: x[0], extra_scores)

        # create a repertoire with default values
        repertoire = cls.init_default(
            genotype=first_genotype,
            descriptor_dim=descriptors.shape[-1],
            population_size=population_size,
            one_extra_score=first_extra_scores,
            keys_extra_scores=keys_extra_scores,
            k=k,
        )

        # add initial population to the repertoire
        return repertoire.add(  # type: ignore
            genotypes, descriptors, fitnesses, extra_scores
        )

    @classmethod
    def init_default(
        cls,
        genotype: Genotype,
        descriptor_dim: int,
        population_size: int,
        one_extra_score: Optional[ExtraScores] = None,
        keys_extra_scores: Tuple[str, ...] = (),
        k: int = 15,
    ) -> DominatedNoveltyRepertoire:
        """Create a DNS repertoire with default values.

        Args:
            genotype: a representative genotype PyTree (leaf shapes define storage).
            descriptor_dim: number of descriptor dimensions.
            population_size: maximum number of individuals kept.
            one_extra_score: a representative extra score PyTree to size buffers.
            keys_extra_scores: keys of extra scores to store in the repertoire.
            k: number of neighbors for novelty metrics.

        Returns:
            A repertoire filled with default values.
        """
        if one_extra_score is None:
            one_extra_score = {}

        one_extra_score = {
            key: value
            for key, value in one_extra_score.items()
            if key in keys_extra_scores
        }

        # default fitness is -inf
        default_fitnesses = -jnp.inf * jnp.ones(shape=(population_size, 1))

        # default genotypes is all zeros
        default_genotypes = jax.tree.map(
            lambda x: jnp.zeros(shape=(population_size,) + x.shape, dtype=x.dtype),
            genotype,
        )

        # default descriptors is NaN (uninitialized)
        default_descriptors = jnp.full(
            shape=(population_size, descriptor_dim), fill_value=jnp.nan
        )

        # default extra scores buffers
        default_extra_scores = jax.tree.map(
            lambda x: jnp.zeros(shape=(population_size,) + x.shape, dtype=x.dtype),
            one_extra_score,
        )

        return cls(
            genotypes=default_genotypes,
            fitnesses=default_fitnesses,
            descriptors=default_descriptors,
            extra_scores=default_extra_scores,
            keys_extra_scores=keys_extra_scores,
            k=k,
        )

class GARepertoire(Repertoire):
    """Class for a simple repertoire for a simple genetic
    algorithm.

    Args:
        genotypes: a PyTree containing the genotypes of the
            individuals in the population. Each leaf has the
            shape (population_size, num_features).
        fitnesses: an array containing the fitness of the individuals
            in the population. With shape (population_size, fitness_dim).
            The implementation of GARepertoire was thought for the case
            where fitness_dim equals 1 but the class can be herited and
            rules adapted for cases where fitness_dim is greater than 1.
        extra_scores: extra scores resulting from the evaluation of the genotypes
        keys_extra_scores: keys of the extra scores to store in the repertoire
    """

    genotypes: Genotype
    fitnesses: Fitness
    extra_scores: ExtraScores
    keys_extra_scores: Tuple[str, ...] = flax.struct.field(
        pytree_node=False,
    )

    @property
    def size(self) -> int:
        """Gives the size of the population."""
        first_leaf = jax.tree.leaves(self.genotypes)[0]
        return int(first_leaf.shape[0])

    def select(
        self,
        key: RNGKey,
        num_samples: int,
        selector: Optional[Selector[GARepertoireT]] = None,
    ) -> GARepertoireT:
        if selector is None:
            selector = UniformSelector(select_with_replacement=True)
        repertoire = selector.select(self, key, num_samples)
        return repertoire

    def filter_extra_scores(self, extra_scores: ExtraScores) -> ExtraScores:
        filtered_extra_scores = {
            key: value
            for key, value in extra_scores.items()
            if key in self.keys_extra_scores
        }
        return filtered_extra_scores

    def add(  # type: ignore
        self,
        batch_of_genotypes: Genotype,
        batch_of_fitnesses: Fitness,
        batch_of_extra_scores: Optional[ExtraScores] = None,
    ) -> GARepertoire:
        """Implements the repertoire addition rules.

        Parents and offsprings are gathered and only the population_size
        bests are kept. The others are killed.

        Args:
            batch_of_genotypes: new genotypes that we try to add.
            batch_of_fitnesses: fitness of those new genotypes.

        Returns:
            The updated repertoire.
        """
        if batch_of_extra_scores is None:
            batch_of_extra_scores = {}

        filtered_batch_of_extra_scores = self.filter_extra_scores(batch_of_extra_scores)

        # gather individuals and fitnesses
        candidates = jax.tree.map(
            lambda x, y: jnp.concatenate((x, y), axis=0),
            self.genotypes,
            batch_of_genotypes,
        )
        candidates_fitnesses = jnp.concatenate(
            (self.fitnesses, batch_of_fitnesses), axis=0
        )

        # sort by fitnesses
        indices = jnp.argsort(jnp.sum(candidates_fitnesses, axis=1))[::-1]

        # keep only the best ones
        survivor_indices = indices[: self.size]

        # keep only the best ones
        new_candidates = jax.tree.map(lambda x: x[survivor_indices], candidates)
        new_extra_scores = jax.tree.map(
            lambda x: x[survivor_indices], filtered_batch_of_extra_scores
        )
        new_repertoire = self.replace(
            genotypes=new_candidates,
            fitnesses=candidates_fitnesses[survivor_indices],
            extra_scores=new_extra_scores,
        )

        return new_repertoire  # type: ignore

    @classmethod
    def init(  # type: ignore
        cls,
        genotypes: Genotype,
        fitnesses: Fitness,
        population_size: int,
        *args,
        extra_scores: Optional[ExtraScores] = None,
        keys_extra_scores: Tuple[str, ...] = (),
        **kwargs,
    ) -> GARepertoire:
        """Initializes the repertoire.

        Start with default values and adds a first batch of genotypes
        to the repertoire.

        Args:
            genotypes: first batch of genotypes
            fitnesses: corresponding fitnesses
            population_size: size of the population we want to evolve
            extra_scores: extra scores resulting from the evaluation of the genotypes
            keys_extra_scores: keys of the extra scores to store in the repertoire

        Returns:
            An initial repertoire.
        """

        if extra_scores is None:
            extra_scores = {}

        # create default fitnesses
        default_fitnesses = -jnp.inf * jnp.ones(
            shape=(population_size, fitnesses.shape[-1])
        )

        # create default genotypes
        default_genotypes = jax.tree.map(
            lambda x: jnp.zeros(shape=(population_size,) + x.shape[1:]), genotypes
        )

        # create default extra scores
        filtered_extra_scores = {
            key: value
            for key, value in extra_scores.items()
            if key in keys_extra_scores
        }

        default_extra_scores = jax.tree.map(
            lambda x: jnp.zeros(shape=(population_size,) + x.shape[1:]),
            filtered_extra_scores,
        )

        # create an initial repertoire with those default values
        repertoire = cls(
            genotypes=default_genotypes,
            fitnesses=default_fitnesses,
            extra_scores=default_extra_scores,
            keys_extra_scores=keys_extra_scores,
        )

        new_repertoire = repertoire.add(genotypes, fitnesses, extra_scores)

        return new_repertoire  # type: ignore

class MELSRepertoire(MapElitesRepertoire):
    """Class for the repertoire in MAP-Elites Low-Spread.

    This class inherits from MapElitesRepertoire. In addition to the stored data in
    MapElitesRepertoire (genotypes, fitnesses, descriptors, centroids), this repertoire
    also maintains an array of spreads. We overload the add, and
    init_default methods of MapElitesRepertoire.

    Refer to Mace 2023 for more info on MAP-Elites Low-Spread:
    https://dl.acm.org/doi/abs/10.1145/3583131.3590433

    Args:
        genotypes: a PyTree containing all the genotypes in the repertoire ordered
            by the centroids. Each leaf has a shape (num_centroids, num_features). The
            PyTree can be a simple JAX array or a more complex nested structure such
            as to represent parameters of neural network in Flax.
        fitnesses: an array that contains the fitness of solutions in each cell of the
            repertoire, ordered by centroids. The array shape is (num_centroids,).
        descriptors: an array that contains the descriptors of solutions in each cell
            of the repertoire, ordered by centroids. The array shape
            is (num_centroids, num_descriptors).
        centroids: an array that contains the centroids of the tessellation. The array
            shape is (num_centroids, num_descriptors).
        spreads: an array that contains the spread of solutions in each cell of the
            repertoire, ordered by centroids. The array shape is (num_centroids,).
    """

    spreads: Spread

    def add(  # type: ignore
        self,
        batch_of_genotypes: Genotype,
        batch_of_descriptors: Descriptor,
        batch_of_fitnesses: Fitness,
        batch_of_extra_scores: Optional[ExtraScores] = None,
    ) -> MELSRepertoire:
        """
        Add a batch of elements to the repertoire.

        The key difference between this method and the default add() in
        MapElitesRepertoire is that it expects each individual to be evaluated
        `num_samples` times, resulting in `num_samples` fitnesses and
        `num_samples` descriptors per individual.

        If multiple individuals may be added to a single cell, this method will
        arbitrarily pick one -- the exact choice depends on the implementation of
        jax.at[].set(), which can be non-deterministic:
        https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html
        We do not currently check if one of the multiple individuals dominates the
        others (dominate means that the individual has both highest fitness and lowest
        spread among the individuals for that cell).

        If `num_samples` is only 1, the spreads will default to 0.

        Args:
            batch_of_genotypes: a batch of genotypes to be added to the repertoire.
                Similarly to the self.genotypes argument, this is a PyTree in which
                the leaves have a shape (batch_size, num_features)
            batch_of_descriptors: an array that contains the descriptors of the
                aforementioned genotypes over all evals. Its shape is
                (batch_size, num_samples, num_descriptors). Note that we "aggregate"
                descriptors by finding the most frequent cell of each individual. Thus,
                the actual descriptors stored in the repertoire are just the coordinates
                of the centroid of the most frequent cell.
            batch_of_fitnesses: an array that contains the fitnesses of the
                aforementioned genotypes over all evals. Its shape is (batch_size,
                num_samples)
            batch_of_extra_scores: unused tree that contains the extra_scores of
                aforementioned genotypes.

        Returns:
            The updated repertoire.
        """

        if batch_of_extra_scores is None:
            batch_of_extra_scores = {}

        filtered_batch_of_extra_scores = self.filter_extra_scores(batch_of_extra_scores)

        batch_size, num_samples = batch_of_fitnesses.shape

        # Compute indices/cells of all descriptors.
        batch_of_all_indices = get_cells_indices(
            batch_of_descriptors.reshape(batch_size * num_samples, -1), self.centroids
        ).reshape((batch_size, num_samples))

        # Compute most frequent cell of each solution.
        batch_of_indices = jax.vmap(_mode)(batch_of_all_indices)[:, None]

        # Compute dispersion / spread. The dispersion is set to zero if
        # num_samples is 1.
        batch_of_spreads = jax.lax.cond(
            num_samples == 1,
            lambda desc: jnp.zeros(batch_size),
            lambda desc: jax.vmap(_dispersion)(
                desc.reshape((batch_size, num_samples, -1))
            ),
            batch_of_descriptors,
        )
        batch_of_spreads = jnp.expand_dims(batch_of_spreads, axis=-1)

        # Compute canonical descriptors as the descriptor of the centroid of the most
        # frequent cell. Note that this line redefines the earlier batch_of_descriptors.
        batch_of_descriptors = jnp.take_along_axis(
            self.centroids, batch_of_indices, axis=0
        )

        # Compute canonical fitnesses as the average fitness.
        #
        # Shape: (batch_size, 1)
        batch_of_fitnesses = batch_of_fitnesses.mean(axis=-1, keepdims=True)

        num_centroids = self.centroids.shape[0]

        # get current repertoire fitnesses and spreads
        current_fitnesses = jnp.take_along_axis(self.fitnesses, batch_of_indices, 0)

        repertoire_spreads = jnp.expand_dims(self.spreads, axis=-1)
        current_spreads = jnp.take_along_axis(repertoire_spreads, batch_of_indices, 0)

        # get addition condition
        addition_condition_fitness = batch_of_fitnesses > current_fitnesses
        addition_condition_spread = batch_of_spreads <= current_spreads
        addition_condition = jnp.logical_and(
            addition_condition_fitness, addition_condition_spread
        )

        # assign fake position when relevant : num_centroids is out of
        # bound
        batch_of_indices = jnp.where(
            addition_condition, batch_of_indices, num_centroids
        )

        # create new repertoire
        new_repertoire_genotypes = jax.tree.map(
            lambda repertoire_genotypes, new_genotypes: repertoire_genotypes.at[
                batch_of_indices.squeeze(axis=-1)
            ].set(new_genotypes),
            self.genotypes,
            batch_of_genotypes,
        )

        # compute new fitness and descriptors
        new_fitnesses = self.fitnesses.at[batch_of_indices.squeeze(axis=-1)].set(
            batch_of_fitnesses,
        )
        new_descriptors = self.descriptors.at[batch_of_indices.squeeze(axis=-1)].set(
            batch_of_descriptors
        )
        new_spreads = self.spreads.at[batch_of_indices.squeeze(axis=-1)].set(
            batch_of_spreads.squeeze(axis=-1)
        )

        # update extra scores
        new_extra_scores = jax.tree.map(
            lambda repertoire_scores, new_scores: repertoire_scores.at[
                batch_of_indices.squeeze(axis=-1)
            ].set(new_scores),
            self.extra_scores,
            filtered_batch_of_extra_scores,
        )

        return MELSRepertoire(
            genotypes=new_repertoire_genotypes,
            fitnesses=new_fitnesses,
            extra_scores=new_extra_scores,
            keys_extra_scores=self.keys_extra_scores,
            descriptors=new_descriptors,
            centroids=self.centroids,
            spreads=new_spreads,
        )

    @classmethod
    def init_default(
        cls,
        genotype: Genotype,
        centroids: Centroid,
        one_extra_score: Optional[ExtraScores] = None,
        keys_extra_scores: Tuple[str, ...] = (),
    ) -> MELSRepertoire:
        """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.

        Note: this function has been kept outside of the object MELS, so
        it can be called easily called from other modules.

        Args:
            genotype: the typical genotype that will be stored.
            centroids: the centroids of the repertoire.
            extra_scores: extra scores to store in the repertoire
            keys_extra_scores: keys of the extra scores to store in the repertoire

        Returns:
            A repertoire filled with default values.
        """
        if one_extra_score is None:
            one_extra_score = {}

        one_extra_score = {
            key: value
            for key, value in one_extra_score.items()
            if key in keys_extra_scores
        }

        # get number of centroids
        num_centroids = centroids.shape[0]

        # default fitness is -inf
        default_fitnesses = -jnp.inf * jnp.ones(shape=(num_centroids, 1))

        # default genotypes is all 0
        default_genotypes = jax.tree.map(
            lambda x: jnp.zeros(shape=(num_centroids,) + x.shape, dtype=x.dtype),
            genotype,
        )

        # default descriptor is all zeros
        default_descriptors = jnp.zeros_like(centroids)

        # default spread is inf so that any spread will be less
        default_spreads = jnp.full(shape=num_centroids, fill_value=jnp.inf)

        # default extra scores is empty dict
        default_extra_scores = jax.tree.map(
            lambda x: jnp.zeros(shape=(num_centroids,) + x.shape, dtype=x.dtype),
            one_extra_score,
        )

        return cls(
            genotypes=default_genotypes,
            fitnesses=default_fitnesses,
            descriptors=default_descriptors,
            centroids=centroids,
            spreads=default_spreads,
            extra_scores=default_extra_scores,
            keys_extra_scores=keys_extra_scores,
        )

class MOMERepertoire(MapElitesRepertoire):
    """Class for the repertoire in Multi Objective Map Elites

    This class inherits from MAPElitesRepertoire. The stored data
    is the same: genotypes, fitnesses, descriptors, centroids.

    The shape of genotypes is (in the case where it's an array):
    (num_centroids, pareto_front_length, genotype_dim).
    When the genotypes is a PyTree, the two first dimensions are the same
    but the third will depend on the leafs.

    The shape of fitnesses is: (num_centroids, pareto_front_length, num_criteria)

    The shape of descriptors and centroids are:
    (num_centroids, num_descriptors, pareto_front_length).
    """

    @property
    def repertoire_capacity(self) -> int:
        """Returns the maximum number of solutions the repertoire can
        contain which corresponds to the number of cells times the
        maximum pareto front length.

        Returns:
            The repertoire capacity.
        """
        first_leaf = jax.tree.leaves(self.genotypes)[0]
        return int(first_leaf.shape[0] * first_leaf.shape[1])

    def select(
        self,
        key: RNGKey,
        num_samples: int,
        selector: Optional[Selector[MOMERepertoireT]] = None,
    ) -> MOMERepertoireT:
        if selector is None:
            selector = MOMEUniformSelector()
        repertoire = selector.select(self, key, num_samples)
        return repertoire

    def _update_masked_pareto_front(
        self,
        pareto_front_fitnesses: ParetoFront[Fitness],
        pareto_front_genotypes: ParetoFront[Genotype],
        pareto_front_descriptors: ParetoFront[Descriptor],
        pareto_front_extra_scores: ParetoFront[ExtraScores],
        mask: Mask,
        new_batch_of_fitnesses: Fitness,
        new_batch_of_genotypes: Genotype,
        new_batch_of_descriptors: Descriptor,
        new_batch_of_extra_scores: ExtraScores,
        new_mask: Mask,
    ) -> Tuple[
        ParetoFront[Fitness],
        ParetoFront[Genotype],
        ParetoFront[Descriptor],
        ParetoFront[ExtraScores],
        Mask,
    ]:
        """Takes a fixed size pareto front, its mask and new points to add.
        Returns updated front and mask.

        Args:
            pareto_front_fitnesses: fitness of the pareto front
            pareto_front_genotypes: corresponding genotypes
            pareto_front_descriptors: corresponding descriptors
            pareto_front_extra_scores: corresponding extra scores
            mask: mask of the front, to hide void parts
            new_batch_of_fitnesses: new batch of fitness that is considered
                to be added to the pareto front
            new_batch_of_genotypes: corresponding genotypes
            new_batch_of_descriptors: corresponding descriptors
            new_batch_of_extra_scores: corresponding extra scores
            new_mask: corresponding mask (no one is masked)

        Returns:
            The updated pareto front.
        """
        # get dimensions
        batch_size = new_batch_of_fitnesses.shape[0]
        num_criteria = new_batch_of_fitnesses.shape[1]

        pareto_front_len = pareto_front_fitnesses.shape[0]  # type: ignore

        descriptors_dim = new_batch_of_descriptors.shape[1]

        # gather all data
        cat_mask = jnp.concatenate([mask, new_mask], axis=-1)
        cat_fitnesses = jnp.concatenate(
            [pareto_front_fitnesses, new_batch_of_fitnesses], axis=0
        )
        cat_genotypes = jax.tree.map(
            lambda x, y: jnp.concatenate([x, y], axis=0),
            pareto_front_genotypes,
            new_batch_of_genotypes,
        )
        cat_descriptors = jnp.concatenate(
            [pareto_front_descriptors, new_batch_of_descriptors], axis=0
        )
        cat_extra_scores = jax.tree.map(
            lambda x, y: jnp.concatenate([x, y], axis=0),
            pareto_front_extra_scores,
            new_batch_of_extra_scores,
        )

        # get new front
        cat_bool_front = compute_masked_pareto_front(
            batch_of_criteria=cat_fitnesses, mask=cat_mask
        )

        # get corresponding indices
        indices = (
            jnp.arange(start=0, stop=pareto_front_len + batch_size) * cat_bool_front
        )
        indices = indices + ~cat_bool_front * (batch_size + pareto_front_len - 1)
        indices = jnp.sort(indices)

        # get new fitness, genotypes and descriptors
        new_front_fitness = jnp.take(cat_fitnesses, indices, axis=0)
        new_front_genotypes = jax.tree.map(
            lambda x: jnp.take(x, indices, axis=0), cat_genotypes
        )
        new_front_descriptors = jnp.take(cat_descriptors, indices, axis=0)
        new_front_extra_scores = jax.tree.map(
            lambda x: jnp.take(x, indices, axis=0), cat_extra_scores
        )

        # compute new mask
        num_front_elements = jnp.sum(cat_bool_front)
        new_mask_indices = jnp.arange(start=0, stop=batch_size + pareto_front_len)
        new_mask_indices = (num_front_elements - new_mask_indices) > 0

        new_mask = jnp.where(
            new_mask_indices,
            jnp.ones(shape=batch_size + pareto_front_len, dtype=bool),
            jnp.zeros(shape=batch_size + pareto_front_len, dtype=bool),
        )

        fitness_mask = jnp.repeat(
            jnp.expand_dims(new_mask, axis=-1), num_criteria, axis=-1
        )
        new_front_fitness = new_front_fitness * fitness_mask

        front_size = len(pareto_front_fitnesses)  # type: ignore
        new_front_fitness = new_front_fitness[:front_size, :]

        new_front_genotypes = jax.tree.map(
            lambda x: x * new_mask_indices[0], new_front_genotypes
        )
        new_front_genotypes = jax.tree.map(
            lambda x: x[:front_size], new_front_genotypes
        )

        descriptors_mask = jnp.repeat(
            jnp.expand_dims(new_mask, axis=-1), descriptors_dim, axis=-1
        )
        new_front_descriptors = new_front_descriptors * descriptors_mask
        new_front_descriptors = new_front_descriptors[:front_size, :]

        new_front_extra_scores = jax.tree.map(
            lambda x: x * new_mask_indices[0], new_front_extra_scores
        )
        new_front_extra_scores = jax.tree.map(
            lambda x: x[:front_size], new_front_extra_scores
        )

        new_mask = ~new_mask[:front_size]

        return (
            new_front_fitness,
            new_front_genotypes,
            new_front_descriptors,
            new_front_extra_scores,
            new_mask,
        )

    def add(  # type: ignore
        self,
        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.

        Shape of the batch_of_genotypes (if an array):
        (batch_size, genotypes_dim)
        Shape of the batch_of_descriptors: (batch_size, num_descriptors)
        Shape of the batch_of_fitnesses: (batch_size, num_criteria)

        Args:
            batch_of_genotypes: a batch of genotypes that we are trying to
                insert into the repertoire.
            batch_of_descriptors: the descriptors of the genotypes we are
                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.
        """
        if batch_of_extra_scores is None:
            batch_of_extra_scores = {}

        filtered_batch_of_extra_scores = self.filter_extra_scores(batch_of_extra_scores)

        # get the indices that corresponds to the descriptors in the repertoire
        batch_of_indices = get_cells_indices(batch_of_descriptors, self.centroids)
        batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)

        def _add_one(
            carry: MOMERepertoire,
            data: Tuple[Genotype, Descriptor, Fitness, ExtraScores, jax.Array],
        ) -> Tuple[MOMERepertoire, Any]:
            # unwrap data
            genotype, descriptors, fitness, extra_scores, index = data

            index = index.astype(jnp.int32)

            # get current repertoire cell data
            cell_genotype = jax.tree.map(lambda x: x[index][0], carry.genotypes)
            cell_fitness = carry.fitnesses[index][0]
            cell_descriptor = carry.descriptors[index][0]
            cell_extra_scores = jax.tree.map(lambda x: x[index][0], carry.extra_scores)
            cell_mask = jnp.any(cell_fitness == -jnp.inf, axis=-1)

            new_genotypes = jax.tree.map(lambda x: jnp.expand_dims(x, axis=0), genotype)

            # update pareto front
            (
                cell_fitness,
                cell_genotype,  # new pf for cell
                cell_descriptor,
                cell_extra_scores,
                cell_mask,
            ) = self._update_masked_pareto_front(
                pareto_front_fitnesses=cell_fitness,
                pareto_front_genotypes=cell_genotype,
                pareto_front_descriptors=cell_descriptor,
                pareto_front_extra_scores=cell_extra_scores,
                mask=cell_mask,
                new_batch_of_fitnesses=jnp.expand_dims(fitness, axis=0),
                new_batch_of_genotypes=new_genotypes,
                new_batch_of_descriptors=jnp.expand_dims(descriptors, axis=0),
                new_batch_of_extra_scores=jax.tree.map(
                    lambda x: jnp.expand_dims(x, axis=0), extra_scores
                ),
                new_mask=jnp.zeros(shape=(1,), dtype=bool),
            )

            # update cell fitness
            cell_fitness = cell_fitness - jnp.inf * jnp.expand_dims(cell_mask, axis=-1)

            # update grid
            new_genotypes = jax.tree.map(
                lambda x, y: x.at[index].set(y), carry.genotypes, cell_genotype
            )
            new_fitnesses = carry.fitnesses.at[index].set(cell_fitness)
            new_descriptors = carry.descriptors.at[index].set(cell_descriptor)
            new_extra_scores = jax.tree.map(
                lambda x, y: x.at[index].set(y), carry.extra_scores, cell_extra_scores
            )
            carry = carry.replace(  # type: ignore
                genotypes=new_genotypes,
                descriptors=new_descriptors,
                fitnesses=new_fitnesses,
                extra_scores=new_extra_scores,
            )

            # return new grid
            return carry, ()

        # scan the addition operation for all the data
        self, _ = jax.lax.scan(
            _add_one,
            self,
            (
                batch_of_genotypes,
                batch_of_descriptors,
                batch_of_fitnesses,
                filtered_batch_of_extra_scores,
                batch_of_indices,
            ),
        )

        return self

    @classmethod
    def init(  # type: ignore
        cls,
        genotypes: Genotype,
        fitnesses: Fitness,
        descriptors: Descriptor,
        centroids: Centroid,
        pareto_front_max_length: int,
        *args,
        extra_scores: Optional[ExtraScores] = None,
        keys_extra_scores: Tuple[str, ...] = (),
        **kwargs,
    ) -> MOMERepertoire:
        """
        Initialize a Multi Objective Map-Elites 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.

        Note: this function has been kept outside of the object MapElites, so it can
        be called easily called from other modules.

        Args:
            genotypes: initial genotypes, pytree in which leaves
                have shape (batch_size, num_features)
            fitnesses: fitness of the initial genotypes of shape:
                (batch_size, num_criteria)
            descriptors: descriptors of the initial genotypes
                of shape (batch_size, num_descriptors)
            centroids: tessellation 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
            keys_extra_scores: keys of the extra_scores of the initial genotypes

        Returns:
            An initialized MAP-Elite repertoire
        """

        warnings.warn(
            (
                "This type of repertoire does not store the extra scores "
                "computed by the scoring function"
            ),
            stacklevel=2,
        )

        if extra_scores is None:
            extra_scores = {}

        filtered_extra_scores = {key: extra_scores[key] for key in keys_extra_scores}

        # get dimensions
        num_criteria = fitnesses.shape[1]
        num_descriptors = descriptors.shape[1]
        num_centroids = centroids.shape[0]

        # create default values
        default_fitnesses = -jnp.inf * jnp.ones(
            shape=(num_centroids, pareto_front_max_length, num_criteria)
        )
        default_genotypes = jax.tree.map(
            lambda x: jnp.zeros(
                shape=(
                    num_centroids,
                    pareto_front_max_length,
                )
                + x.shape[1:]
            ),
            genotypes,
        )
        default_descriptors = jnp.zeros(
            shape=(num_centroids, pareto_front_max_length, num_descriptors)
        )

        default_extra_scores = jax.tree.map(
            lambda x: jnp.zeros(
                shape=(num_centroids, pareto_front_max_length, *x.shape[1:])
            ),
            filtered_extra_scores,
        )

        # create repertoire with default values
        repertoire = MOMERepertoire(  # type: ignore
            genotypes=default_genotypes,
            fitnesses=default_fitnesses,
            descriptors=default_descriptors,
            centroids=centroids,
            extra_scores=default_extra_scores,
            keys_extra_scores=keys_extra_scores,
        )

        # add first batch of individuals in the repertoire
        new_repertoire = repertoire.add(genotypes, descriptors, fitnesses, extra_scores)

        return new_repertoire  # type: ignore

    @jax.jit
    def compute_global_pareto_front(
        self,
    ) -> Tuple[ParetoFront[Fitness], Mask]:
        """Merge all the pareto fronts of the MOME repertoire into a single one
        called global pareto front.

        Returns:
            The pareto front and its mask.
        """
        fitnesses = jnp.concatenate(self.fitnesses, axis=0)
        mask = jnp.any(fitnesses == -jnp.inf, axis=-1)
        pareto_mask = compute_masked_pareto_front(fitnesses, mask)
        pareto_front = fitnesses - jnp.inf * (~jnp.array([pareto_mask, pareto_mask]).T)

        return pareto_front, pareto_mask