Feat: Add MAP-Annealing Repertoire (to run CMA-MAE)

qdax/core/containers/mae_repertoire.py

@@ -0,0 +1,305 @@
+from __future__ import annotations
+
+import warnings
+from typing import Callable, Optional
+
+import flax
+import jax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+from typing_extensions import TypeAlias
+
+from qdax.core.containers.mapelites_repertoire import (
+ MapElitesRepertoire,
+ get_cells_indices,
+)
+from qdax.types import Centroid, Descriptor, ExtraScores, Fitness, Genotype
+
+Threshold: TypeAlias = Fitness
+
+
+class MAERepertoire(MapElitesRepertoire):
+ """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).
+ """
+
+ thresholds: Threshold
+ archive_learning_rate: float = flax.struct.field(pytree_node=False)
+
+ 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 + "thresholds.npy", self.thresholds)
+
+ @classmethod
+ def load(cls, reconstruction_fn: Callable, path: str = "./") -> MAERepertoire:
+ """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")
+ thresholds = jnp.load(path + "thresholds.npy")
+
+ return cls(
+ genotypes=genotypes,
+ fitnesses=fitnesses,
+ descriptors=descriptors,
+ centroids=centroids,
+ thresholds=thresholds,
+ )
+
+ @jax.jit
+ def add(
+ self,
+ batch_of_genotypes: Genotype,
+ batch_of_descriptors: Descriptor,
+ batch_of_fitnesses: Fitness,
+ batch_of_extra_scores: Optional[ExtraScores] = None,
+ ) -> MAERepertoire:
+ """
+ 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_thresholds = jnp.expand_dims(self.thresholds, axis=-1)
+ current_thresholds_with_inf = jnp.take_along_axis(
+ repertoire_thresholds, batch_of_indices, 0
+ )
+ current_thresholds = jnp.where(
+ jnp.isinf(current_thresholds_with_inf),
+ x=-jnp.inf,
+ y=current_thresholds_with_inf,
+ )
+ addition_condition = batch_of_fitnesses > current_thresholds
+
+ # 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)
+ )
+
+ previous_thresholds = self.thresholds.at[
+ batch_of_indices.squeeze(axis=-1)
+ ].get()
+ updated_thresholds = jnp.where(
+ jnp.isnan(previous_thresholds),
+ x=batch_of_fitnesses.squeeze(axis=-1),
+ y=previous_thresholds,
+ )
+ updated_thresholds = (
+ updated_thresholds * (1.0 - self.archive_learning_rate)
+ + batch_of_fitnesses.squeeze(axis=-1) * self.archive_learning_rate
+ )
+
+ new_thresholds = self.thresholds.at[batch_of_indices.squeeze(axis=-1)].set(
+ updated_thresholds
+ )
+ new_descriptors = self.descriptors.at[batch_of_indices.squeeze(axis=-1)].set(
+ batch_of_descriptors
+ )
+
+ return self.replace( # type: ignore
+ genotypes=new_repertoire_genotypes,
+ fitnesses=new_fitnesses,
+ descriptors=new_descriptors,
+ centroids=self.centroids,
+ thresholds=new_thresholds,
+ archive_learning_rate=self.archive_learning_rate,
+ )
+
+ @classmethod
+ def init(
+ cls,
+ genotypes: Genotype,
+ fitnesses: Fitness,
+ descriptors: Descriptor,
+ centroids: Centroid,
+ archive_learning_rate: float,
+ extra_scores: Optional[ExtraScores] = None,
+ min_threshold: Optional[float] = None,
+ ) -> MAERepertoire:
+ """
+ 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,
+ archive_learning_rate=archive_learning_rate,
+ min_threshold=min_threshold,
+ )
+
+ # 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,
+ archive_learning_rate: float,
+ min_threshold: Optional[float] = None,
+ ) -> MAERepertoire:
+ """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)
+ if min_threshold is not None:
+ default_thresholds = jnp.full_like(default_fitnesses, min_threshold)
+ else:
+ default_thresholds = jnp.full_like(default_fitnesses, -jnp.inf)
+
+ # 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,
+ thresholds=default_thresholds,
+ archive_learning_rate=archive_learning_rate,
+ )