evosax: Evolution Strategies in JAX 🦎

<a href="https://github.com/RobertTLange/evosax/blob/main/docs/logo.png?raw=true"><img src="https://github.com/RobertTLange/evosax/blob/main/docs/logo.png?raw=true" width="170" align="right" /></a>

Tired of having to handle asynchronous processes for neuroevolution? Do you want to leverage massive vectorization and high-throughput accelerators for Evolution Strategies? evosax provides a comprehensive, high-performance library that implements Evolution Strategies (ES) in JAX. By leveraging XLA compilation and JAX's transformation primitives, evosax enables researchers and practitioners to efficiently scale evolutionary algorithms to modern hardware accelerators without the traditional overhead of distributed implementations.

The API follows the classical ask-eval-tell cycle of ES, with full support for JAX's transformations (jit, vmap, lax.scan). The library includes 30+ evolution strategies, from classics like CMA-ES and Differential Evolution to modern approaches like OpenAI-ES and Diffusion Evolution.

Get started here 👉

Basic evosax API Usage 🍲

import jax
from evosax.algorithms import CMA_ES


# Instantiate the search strategy
es = CMA_ES(population_size=32, solution=dummy_solution)
params = es.default_params

# Initialize state
key = jax.random.key(0)
state = es.init(key, params)

# Ask-Eval-Tell loop
for i in range(num_generations):
    key, key_ask, key_eval = jax.random.split(key, 3)

    # Generate a set of candidate solutions to evaluate
    population, state = es.ask(key_ask, state, params)

    # Evaluate the fitness of the population
    fitness = ...

    # Update the evolution strategy
    state = es.tell(population, fitness, state, params)

# Get best solution
state.best_solution, state.best_fitness

Implemented Evolution Strategies 🦎

| Strategy | Reference | Import | Example | |-----------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------| --- | | Simple Evolution Strategy | Rechenberg (1978) | SimpleES | | OpenAI-ES | Salimans et al. (2017) | Open_ES | | CMA-ES | Hansen & Ostermeier (2001) | CMA_ES | | Sep-CMA-ES | Ros & Hansen (2008) | Sep_CMA_ES | | xNES | Wierstra et al. (2014) | xNES | | SNES | Wierstra et al. (2014) | SNES | | MA-ES | Bayer & Sendhoff (2017) | MA_ES | | LM-MA-ES | Loshchilov et al. (2017) | LM_MA_ES | | Rm_ES | Li & Zhang (2017) | Rm_ES | | PGPE | Sehnke et al. (2010) | PGPE | | ARS | Mania et al. (2018) | ARS | | ESMC | Merchant et al. (2021) | ESMC | | Persistent ES | Vicol et al. (2021) | PersistentES | | Noise-Reuse ES | Li et al. (2023) | NoiseReuseES | | CR-FM-NES | Nomura & Ono (2022)