Containers¶

class Archive(PyTreeNode):
    """Stores jnp.ndarray in a way that makes insertion jittable.

    An example of use of the archive is the algorithm QDPG: state
    descriptors are stored in this archive and a novelty scorer compares
    new state desciptors to the state descriptors stored in this archive.

    Note: notations suppose that the elements are called state desciptors.
    If we where to use this structure for another application, it would be
    better to change the variables name for another one. Does not seem
    necessary at the moment though.
    """

    data: jnp.ndarray  # initialised with nan everywhere
    current_position: int
    acceptance_threshold: float
    state_descriptor_size: int
    max_size: int

    @property
    def size(self) -> float:
        """Compute the number of state descriptors stored in the archive.

        Returns:
            Size of the archive.
        """
        # remove fake borders
        fake_data = jnp.isnan(self.data)

        # count number of real data
        return sum(~fake_data)

    @classmethod
    def create(
        cls,
        acceptance_threshold: float,
        state_descriptor_size: int,
        max_size: int,
    ) -> Archive:
        """Create an Archive instance.

        This class method provides a convenient way to create the archive while
        keeping the __init__ function for more general way to init an archive.

        Args:
            acceptance_threshold: the minimal distance to a stored descriptor to
                be respected for a new descriptor to be added.
            state_descriptor_size: the number of elements in a state descriptor.
            max_size: the maximal size of the archive. In case of overflow, previous
                elements are replaced by new ones. Defaults to 80000.

        Returns:
            A newly initialized archive.
        """
        init_data = jnp.ones((max_size, state_descriptor_size)) * jnp.nan
        return cls(  # type: ignore
            data=init_data,
            current_position=0,
            acceptance_threshold=acceptance_threshold,
            state_descriptor_size=state_descriptor_size,
            max_size=max_size,
        )

    @jax.jit
    def _single_insertion(self, state_descriptor: jnp.ndarray) -> Archive:
        """Insert a single element.

        If the archive is not full yet, the new element replaces a fake
        border, if it is full, it replaces the element that was inserted
        first in the archive.

        Args:
            state_descriptor: state descriptor to be added.

        Returns:
            Return the archive with the newly added element.
        """
        new_current_position = self.current_position + 1
        new_data = jax.lax.dynamic_update_slice_in_dim(
            self.data,
            state_descriptor.reshape(1, -1),
            start_index=self.current_position % self.max_size,
            axis=0,
        )

        return self.replace(  # type: ignore
            current_position=new_current_position, data=new_data
        )

    @jax.jit
    def _conditioned_single_insertion(
        self, condition: bool, state_descriptor: jnp.ndarray
    ) -> Tuple[Archive, jnp.ndarray]:
        """Inserts a single element under a condition.

        The function also retrieves the added elements.

        Args:
            condition: condition for being added in the archive.
            state_descriptor: state descriptor to be added under the
                given condition.

        Returns:
            The new archive and the added elements.
        """

        def true_fun(
            archive: Archive, state_descriptor: jnp.ndarray
        ) -> Tuple[Archive, jnp.ndarray]:
            return archive._single_insertion(state_descriptor), state_descriptor

        def false_fun(
            archive: Archive, state_descriptor: jnp.ndarray
        ) -> Tuple[Archive, jnp.ndarray]:
            return archive, jnp.ones_like(state_descriptor) * jnp.nan

        return jax.lax.cond(  # type: ignore
            condition, true_fun, false_fun, self, state_descriptor
        )

    @jax.jit
    def insert(self, state_descriptors: jnp.ndarray) -> Archive:
        """Tries to insert a batch of state descriptors in the archive.

        1. First, look at the distance of each new state descriptor with the
        already stored ones.
        2. Then, scan the state descriptors, check the distance with
        the descriptors inserted during the scan.
        3. If the state descriptor verified the first condition (not too close
        to a state descriptor in the old archive) and the second (not too close
        from a state descriptor that has just been added), then it is added
        to the archive.

        Note 1: the archive has a fixed size, hence, in case of overflow, the
        first elements added are removed first (FIFO style).
        Note 2: keep in mind that fake descriptors are used to help keep the size
        constant. Those correspond to a descriptor very far away from the typical
        values of the problem at hand.

        Args:
            state_descriptors: state descriptors to be added.

        Returns:
            New archive updated with the state descriptors.
        """
        state_descriptors = state_descriptors.reshape((-1, state_descriptors.shape[-1]))

        # get nearest neigbor for each new state descriptor
        values, _indices = knn(self.data, state_descriptors, 1)

        # get indices where distance bigger than threshold
        relevant_indices = jnp.where(
            values.squeeze() > self.acceptance_threshold, x=0, y=1
        )

        def iterate_fn(
            carry: Tuple[Archive, jnp.ndarray, int], condition_data: Dict
        ) -> Tuple[Tuple[Archive, jnp.ndarray, int], Any]:
            """Iterates over the archive to add elements one after the other.

            Args:
                carry: tuple containing the archive, the state descriptors and the
                    indices.

                condition_data: the first addition condition of the state descriptors
                    given, which corresponds to being sufficiently far away from already
                    stored descriptors.

            Returns:
                The update tuple.
            """
            archive, new_elements, index = carry

            first_condition = condition_data["condition"]
            state_descriptor = condition_data["state_descriptor"]

            # do the filtering among the added elements
            # get nearest neigbor for each new state descriptor
            values, _indices = knn(new_elements, state_descriptor.reshape(1, -1), 1)

            # get indices where distance bigger than threshold
            not_too_close = jnp.where(
                values.squeeze() > self.acceptance_threshold, x=0, y=1
            )
            second_condition = not_too_close.sum()
            condition = (first_condition + second_condition) == 0

            new_archive, added_element = archive._conditioned_single_insertion(
                condition, state_descriptor
            )
            new_elements = new_elements.at[index].set(added_element)
            index += 1

            return (
                (new_archive, new_elements, index),
                (),
            )

        new_elements = jnp.ones_like(state_descriptors) * jnp.nan

        # iterate over the indices
        (new_archive, _, _), _ = jax.lax.scan(
            iterate_fn,
            (self, new_elements, 0),
            {
                "condition": relevant_indices,
                "state_descriptor": state_descriptors,
            },
        )

        return new_archive  # type: ignore

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.
    """

    genotypes: Genotype
    fitnesses: Fitness

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

    def save(self, path: str = "./") -> None:
        """Saves the repertoire.

        Args:
            path: place to store the files. Defaults to "./".
        """

        def flatten_genotype(genotype: Genotype) -> jnp.ndarray:
            flatten_genotype, _ = ravel_pytree(genotype)
            return flatten_genotype

        # flatten all the genotypes
        flat_genotypes = jax.vmap(flatten_genotype)(self.genotypes)

        jnp.save(path + "genotypes.npy", flat_genotypes)
        jnp.save(path + "scores.npy", self.fitnesses)

    @classmethod
    def load(cls, reconstruction_fn: Callable, path: str = "./") -> GARepertoire:
        """Loads a GA Repertoire.

        Args:
            reconstruction_fn: Function to reconstruct a PyTree
                from a flat array.
            path: Path where the data is saved. Defaults to "./".

        Returns:
            A GA Repertoire.
        """

        flat_genotypes = jnp.load(path + "genotypes.npy")
        genotypes = jax.vmap(reconstruction_fn)(flat_genotypes)

        fitnesses = jnp.load(path + "fitnesses.npy")

        return cls(
            genotypes=genotypes,
            fitnesses=fitnesses,
        )

    @partial(jax.jit, static_argnames=("num_samples",))
    def sample(self, random_key: RNGKey, num_samples: int) -> Tuple[Genotype, RNGKey]:
        """Sample genotypes from the repertoire.

        Args:
            random_key: a random key to handle stochasticity.
            num_samples: the number of genotypes to sample.

        Returns:
            The sample of genotypes.
        """

        # prepare sampling probability
        mask = self.fitnesses != -jnp.inf
        p = jnp.any(mask, axis=-1) / jnp.sum(jnp.any(mask, axis=-1))

        # sample
        random_key, subkey = jax.random.split(random_key)
        samples = jax.tree_util.tree_map(
            lambda x: jax.random.choice(
                subkey, x, shape=(num_samples,), p=p, replace=False
            ),
            self.genotypes,
        )

        return samples, random_key

    @jax.jit
    def add(
        self, batch_of_genotypes: Genotype, batch_of_fitnesses: Fitness
    ) -> 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.
        """

        # gather individuals and fitnesses
        candidates = jax.tree_util.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_util.tree_map(
            lambda x: x[survivor_indices], candidates
        )

        new_repertoire = self.replace(
            genotypes=new_candidates, fitnesses=candidates_fitnesses[survivor_indices]
        )

        return new_repertoire  # type: ignore

    @classmethod
    def init(  # type: ignore
        cls,
        genotypes: Genotype,
        fitnesses: Fitness,
        population_size: int,
    ) -> 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

        Returns:
            An initial repertoire.
        """
        # create default fitnesses
        default_fitnesses = -jnp.inf * jnp.ones(
            shape=(population_size, fitnesses.shape[-1])
        )

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

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

        new_repertoire = repertoire.add(genotypes, fitnesses)

        return new_repertoire  # type: ignore

class MapElitesRepertoire(flax.struct.PyTreeNode):
    """Class for the repertoire in Map Elites.

    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).
    """

    genotypes: Genotype
    fitnesses: Fitness
    descriptors: Descriptor
    centroids: Centroid

    def save(self, path: str = "./") -> None:
        """Saves the repertoire on disk in the form of .npy files.

        Flattens the genotypes to store it with .npy format. Supposes that
        a user will have access to the reconstruction function when loading
        the genotypes.

        Args:
            path: Path where the data will be saved. Defaults to "./".
        """

        def flatten_genotype(genotype: Genotype) -> jnp.ndarray:
            flatten_genotype, _ = ravel_pytree(genotype)
            return flatten_genotype

        # flatten all the genotypes
        flat_genotypes = jax.vmap(flatten_genotype)(self.genotypes)

        # save data
        jnp.save(path + "genotypes.npy", flat_genotypes)
        jnp.save(path + "fitnesses.npy", self.fitnesses)
        jnp.save(path + "descriptors.npy", self.descriptors)
        jnp.save(path + "centroids.npy", self.centroids)

    @classmethod
    def load(cls, reconstruction_fn: Callable, path: str = "./") -> MapElitesRepertoire:
        """Loads a MAP Elites Repertoire.

        Args:
            reconstruction_fn: Function to reconstruct a PyTree
                from a flat array.
            path: Path where the data is saved. Defaults to "./".

        Returns:
            A MAP Elites Repertoire.
        """

        flat_genotypes = jnp.load(path + "genotypes.npy")
        genotypes = jax.vmap(reconstruction_fn)(flat_genotypes)

        fitnesses = jnp.load(path + "fitnesses.npy")
        descriptors = jnp.load(path + "descriptors.npy")
        centroids = jnp.load(path + "centroids.npy")

        return cls(
            genotypes=genotypes,
            fitnesses=fitnesses,
            descriptors=descriptors,
            centroids=centroids,
        )

    @partial(jax.jit, static_argnames=("num_samples",))
    def sample(self, random_key: RNGKey, num_samples: int) -> Tuple[Genotype, RNGKey]:
        """Sample elements in the repertoire.

        Args:
            random_key: a jax PRNG random key
            num_samples: the number of elements to be sampled

        Returns:
            samples: a batch of genotypes sampled in the repertoire
            random_key: an updated jax PRNG random key
        """

        repertoire_empty = self.fitnesses == -jnp.inf
        p = (1.0 - repertoire_empty) / jnp.sum(1.0 - repertoire_empty)

        random_key, subkey = jax.random.split(random_key)
        samples = jax.tree_util.tree_map(
            lambda x: jax.random.choice(subkey, x, shape=(num_samples,), p=p),
            self.genotypes,
        )

        return samples, random_key

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

        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. Its shape is (batch_size, num_descriptors)
            batch_of_fitnesses: an array that contains the fitnesses of the
                aforementioned genotypes. Its shape is (batch_size,)
            batch_of_extra_scores: unused tree that contains the extra_scores of
                aforementioned genotypes.

        Returns:
            The updated MAP-Elites repertoire.
        """

        batch_of_indices = get_cells_indices(batch_of_descriptors, self.centroids)
        batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)
        batch_of_fitnesses = jnp.expand_dims(batch_of_fitnesses, axis=-1)

        num_centroids = self.centroids.shape[0]

        # get fitness segment max
        best_fitnesses = jax.ops.segment_max(
            batch_of_fitnesses,
            batch_of_indices.astype(jnp.int32).squeeze(axis=-1),
            num_segments=num_centroids,
        )

        cond_values = jnp.take_along_axis(best_fitnesses, batch_of_indices, 0)

        # put dominated fitness to -jnp.inf
        batch_of_fitnesses = jnp.where(
            batch_of_fitnesses == cond_values, x=batch_of_fitnesses, y=-jnp.inf
        )

        # get addition condition
        repertoire_fitnesses = jnp.expand_dims(self.fitnesses, axis=-1)
        current_fitnesses = jnp.take_along_axis(
            repertoire_fitnesses, batch_of_indices, 0
        )
        addition_condition = batch_of_fitnesses > current_fitnesses

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

        # create new repertoire
        new_repertoire_genotypes = jax.tree_util.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.squeeze(axis=-1)
        )
        new_descriptors = self.descriptors.at[batch_of_indices.squeeze(axis=-1)].set(
            batch_of_descriptors
        )

        return MapElitesRepertoire(
            genotypes=new_repertoire_genotypes,
            fitnesses=new_fitnesses,
            descriptors=new_descriptors,
            centroids=self.centroids,
        )

    @classmethod
    def init(
        cls,
        genotypes: Genotype,
        fitnesses: Fitness,
        descriptors: Descriptor,
        centroids: Centroid,
        extra_scores: Optional[ExtraScores] = None,
    ) -> MapElitesRepertoire:
        """
        Initialize a 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,)
            descriptors: descriptors of the initial genotypes
                of shape (batch_size, num_descriptors)
            centroids: tesselation centroids of shape (batch_size, num_descriptors)
            extra_scores: unused 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,
        )

        # retrieve one genotype from the population
        first_genotype = jax.tree_util.tree_map(lambda x: x[0], genotypes)

        # create a repertoire with default values
        repertoire = cls.init_default(genotype=first_genotype, centroids=centroids)

        # add initial population to the repertoire
        new_repertoire = repertoire.add(genotypes, descriptors, fitnesses, extra_scores)

        return new_repertoire  # type: ignore

    @classmethod
    def init_default(
        cls,
        genotype: Genotype,
        centroids: Centroid,
    ) -> MapElitesRepertoire:
        """Initialize a 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:
            genotype: the typical genotype that will be stored.
            centroids: the centroids of the repertoire

        Returns:
            A repertoire filled with default values.
        """

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

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

        # default genotypes is all 0
        default_genotypes = jax.tree_util.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)

        return cls(
            genotypes=default_genotypes,
            fitnesses=default_fitnesses,
            descriptors=default_descriptors,
            centroids=centroids,
        )

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 save, load, 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 save(self, path: str = "./") -> None:
        """Saves the repertoire on disk in the form of .npy files.

        Flattens the genotypes to store it with .npy format. Supposes that
        a user will have access to the reconstruction function when loading
        the genotypes.

        Args:
            path: Path where the data will be saved. Defaults to "./".
        """

        def flatten_genotype(genotype: Genotype) -> jnp.ndarray:
            flatten_genotype, _ = ravel_pytree(genotype)
            return flatten_genotype

        # flatten all the genotypes
        flat_genotypes = jax.vmap(flatten_genotype)(self.genotypes)

        # save data
        jnp.save(path + "genotypes.npy", flat_genotypes)
        jnp.save(path + "fitnesses.npy", self.fitnesses)
        jnp.save(path + "descriptors.npy", self.descriptors)
        jnp.save(path + "centroids.npy", self.centroids)
        jnp.save(path + "spreads.npy", self.spreads)

    @classmethod
    def load(cls, reconstruction_fn: Callable, path: str = "./") -> MELSRepertoire:
        """Loads a MAP-Elites Low-Spread Repertoire.

        Args:
            reconstruction_fn: Function to reconstruct a PyTree
                from a flat array.
            path: Path where the data is saved. Defaults to "./".

        Returns:
            A MAP-Elites Low-Spread Repertoire.
        """

        flat_genotypes = jnp.load(path + "genotypes.npy")
        genotypes = jax.vmap(reconstruction_fn)(flat_genotypes)

        fitnesses = jnp.load(path + "fitnesses.npy")
        descriptors = jnp.load(path + "descriptors.npy")
        centroids = jnp.load(path + "centroids.npy")
        spreads = jnp.load(path + "spreads.npy")

        return cls(
            genotypes=genotypes,
            fitnesses=fitnesses,
            descriptors=descriptors,
            centroids=centroids,
            spreads=spreads,
        )

    @jax.jit
    def add(
        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.
        """
        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
        repertoire_fitnesses = jnp.expand_dims(self.fitnesses, axis=-1)
        current_fitnesses = jnp.take_along_axis(
            repertoire_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, x=batch_of_indices, y=num_centroids
        )

        # create new repertoire
        new_repertoire_genotypes = jax.tree_util.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.squeeze(axis=-1)
        )
        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)
        )

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

    @classmethod
    def init_default(
        cls,
        genotype: Genotype,
        centroids: Centroid,
    ) -> 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.

        Returns:
            A repertoire filled with default values.
        """

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

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

        # default genotypes is all 0
        default_genotypes = jax.tree_util.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)

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