Optimizing multiple objectives with NSGA2 & SPEA2 in JAX¶

%pip install -U "jax[cuda]"

%pip install -U "qdax[examples]"

import jax.numpy as jnp
import jax

from typing import Tuple

import matplotlib.pyplot as plt

from functools import partial

from qdax.baselines.nsga2 import (
    NSGA2
)
from qdax.baselines.spea2 import (
    SPEA2
)

from qdax.core.emitters.mutation_operators import (
    polynomial_crossover,
    polynomial_mutation
)
from qdax.core.emitters.standard_emitters import MixingEmitter
from qdax.utils.pareto_front import compute_pareto_front
from qdax.utils.plotting import plot_global_pareto_front

from qdax.utils.pareto_front import compute_pareto_front
from qdax.utils.plotting import plot_global_pareto_front
from qdax.utils.metrics import default_ga_metrics

from qdax.custom_types import Genotype, Fitness, Descriptor

#@markdown ---
population_size = 1000 #@param {type:"integer"}
num_iterations = 1000 #@param {type:"integer"}
proportion_mutation = 0.80 #@param {type:"number"}
minval = -5.12 #@param {type:"number"}
maxval = 5.12 #@param {type:"number"}
genotype_dim = 6 #@param {type:"integer"}
lag = 2.2 #@param {type:"number"}
base_lag = 0 #@param {type:"number"}
# for spea2
num_neighbours=1 #@param {type:"integer"}
#@markdown ---

def rastrigin_scorer(
    genotypes: jax.Arraybase_lag: float, lag: float
) -> Tuple[Fitness, Descriptor]:
    """
    Rastrigin Scorer with first two dimensions as descriptors
    """
    descriptors = genotypes[:, :2]
    f1 = -(
        10 * genotypes.shape[1]
        + jnp.sum(
            (genotypes - base_lag) ** 2
            - 10 * jnp.cos(2 * jnp.pi * (genotypes - base_lag)),
            axis=1,
        )
    )

    f2 = -(
        10 * genotypes.shape[1]
        + jnp.sum(
            (genotypes - lag) ** 2 - 10 * jnp.cos(2 * jnp.pi * (genotypes - lag)),
            axis=1,
        )
    )
    scores = jnp.stack([f1, f2], axis=-1)

    return scores, descriptors

# Scoring function
scoring_function = partial(
    rastrigin_scorer,
    lag=lag,
    base_lag=base_lag
)

def scoring_fn(x, key):
    return scoring_function(x)[0], {}

# Initial population
key = jax.random.key(0)
key, subkey = jax.random.split(key)
genotypes = jax.random.uniform(
    subkey, (population_size, genotype_dim), minval=minval, maxval=maxval, dtype=jnp.float32
)

# Mutation & Crossover
crossover_function = partial(
    polynomial_crossover, 
    proportion_var_to_change=0.5,
)

mutation_function = partial(
    polynomial_mutation, 
    proportion_to_mutate=0.5, 
    eta=0.05, 
    minval=minval, 
    maxval=maxval
)

# Define the emitter
# NSGA-II and SPEA2 use batch size = population size
mixing_emitter = MixingEmitter(
    mutation_fn=mutation_function, 
    variation_fn=crossover_function, 
    variation_percentage=1-proportion_mutation, 
    batch_size=population_size, 
)

# instantitiate nsga2
nsga2 = NSGA2(
    scoring_function=scoring_fn,
    emitter=mixing_emitter,
    metrics_function=default_ga_metrics
)

# init nsga2
key, subkey = jax.random.split(key)
repertoire, emitter_state, init_metrics = nsga2.init(
    genotypes,
    population_size,
    subkey
)

# Run optimization loop
(repertoire, emitter_state, key), _ = jax.lax.scan(
    nsga2.scan_update, (repertoire, emitter_state, key), (), length=num_iterations
)

fig, ax = plt.subplots(figsize=(9, 6))
pareto_bool = compute_pareto_front(repertoire.fitnesses)
plot_global_pareto_front(repertoire.fitnesses[pareto_bool], ax=ax)
ax.set_title("Pareto front obtained by NSGA2", fontsize=16)
ax.set_xlabel("Fitness Dimension 1", fontsize=14)
ax.set_ylabel("Fitness Dimension 2", fontsize=14)
plt.grid()
plt.show()

# instantitiate spea2
spea2 = SPEA2(
    scoring_function=scoring_fn,
    emitter=mixing_emitter,
    metrics_function=default_ga_metrics
)

# init spea2
key, subkey = jax.random.split(key)
repertoire, emitter_state, init_metrics = spea2.init(
    genotypes,
    population_size,
    num_neighbours,
    subkey,
)

# run optimization loop
(repertoire, emitter_state, key), _ = jax.lax.scan(
    spea2.scan_update, (repertoire, emitter_state, key), (), length=num_iterations
)

fig, ax = plt.subplots(figsize=(9, 6))
pareto_bool = compute_pareto_front(repertoire.fitnesses)
plot_global_pareto_front(repertoire.fitnesses[pareto_bool], ax=ax)
ax.set_title("Pareto front obtained by SPEA2", fontsize=16)
ax.set_xlabel("Fitness Dimension 1", fontsize=14)
ax.set_ylabel("Fitness Dimension 2", fontsize=14)
plt.grid()
plt.show()