Using Individuals PyTrees to increase the flexibility of Information stored in MAPElites Repertoires

[Lookatator]

Here is one of the structures I thought of to increase the flexibility of the content we save in the repertoires, while not breaking too much the current API.

• all the information related with individuals (genotypes, fitnesses, descriptors...) are stored in a pytree IndividualQD
• If more information is needed (such as observations in AURORA), we can inherit from IndividualQD (see IndividualAURORA below for example)
• the repertoire now stores individuals + the repertoire properties (e.g. the l_value for the Unstructured Archive, the grid_shape for the Grid Repertoire, the centroids for the CVT repertoire...).
• For backward compatibility: some properties genotypes, fitnesses and descriptors can be added to the repertoire (see below).

With this new structure, we could create any type of QD individual with plenty of additional information, such as: curiosity score, novelty estimations, observations...

class IndividualQD(flax.struct.PyTreeNode):
    genotypes: Genotype
    fitnesses: Fitness
    descriptors: Descriptor

    @classmethod
    def create(cls, genotypes, fitnesses, descriptors, extra_scores):
        return cls(genotypes, fitnesses, descriptors)

#  For example with AURORA:
class IndividualAURORA(IndividualQD):
    genotypes: Genotype
    fitnesses: Fitness
    descriptors: Descriptor
    observations: Observations

    @classmethod
    def create(cls, genotypes, fitnesses, descriptors, extra_scores):
        return cls(genotypes, fitnesses, descriptors, extra_scores['observations'])

class MapElitesRepertoire:
    individuals: IndividualQD
    centroids: Centroid
    _create_individual_fn: Callable[[Genotype, Fitness, Descriptor], IndividualQD] = flax.struct.field(pytree_node=False)

    @classmethod
    def create(cls, genotypes, fitnesses, descriptors, centroids, extra_scores, create_individual_fn=None):
        if create_individual_fn is None:
            #  Default function for creating the most standard individuals
            create_individual_fn = IndividualQD.create

        individuals = create_individual_fn(genotypes, fitnesses, descriptors, extra_scores)
        return MapElitesRepertoire(individuals, centroids, create_individual_fn)


    ### for backward compatibility ###
    @property
    def genotypes(self) -> Genotype:
        return self.individuals.genotypes

    @property
    def fitnesses(self) -> Fitness:
        return self.individuals.fitnesses

    @property
    def descriptors(self) -> Descriptor:
        return self.individuals.descriptors

def example_usage():
    # For MAP-Elites 
    MapElitesRepertoire.create(genotypes, descriptors, centroids, extra_scores)

    # For AURORA
    MapElitesRepertoire.create(genotypes, descriptors, centroids, extra_scores, create_individual_fn=IndividualAURORA.create)