evosax: Evolution Strategies in JAX 🦎

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, dummy_solution, 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, metrics = es.tell(key, population, fitness, state, params)

# Get best solution
state.best_solution, state.best_fitness

Population-based algorithms such as SimpleGA use the same ask-tell loop, but their state is initialized from an explicit initial population and its fitness:

import jax
import jax.numpy as jnp
from evosax.algorithms import SimpleGA


key = jax.random.key(0)
solution = dummy_solution
ga = SimpleGA(population_size=32, solution=solution)
params = ga.default_params

# Replicate a template solution (or provide your own evaluated population).
population_init = jax.tree.map(
    lambda x: jnp.repeat(x[None, ...], ga.population_size, axis=0),
    solution,
)
fitness_init = ...

# Initialize state from the initial population and its fitness.
state = ga.init(key, population_init, fitness_init, params)

The fitness_init argument seeds the stored population ranking used by the first ask(...) call. After initialization, ga.ask(...) returns the next population to evaluate and ga.tell(...) updates the algorithm state with the evaluated population and 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)	CR_FM_NES
Guided ES	Maheswaranathan et al. (2018)	GuidedES
ASEBO	Choromanski et al. (2019)	ASEBO
Discovered ES	Lange et al. (2023a)	DiscoveredES
LES	Lange et al. (2023a)	LearnedES
EvoTF	Lange et al. (2024)	EvoTF_ES
iAMaLGaM-Full	Bosman et al. (2013)	iAMaLGaM_Full
iAMaLGaM-Univariate	Bosman et al. (2013)	iAMaLGaM_Univariate
Gradientless Descent	Golovin et al. (2019)	GradientlessDescent
Simulated Annealing	Rasdi Rere et al. (2015)	SimAnneal
Hill Climbing	Rasdi Rere et al. (2015)	HillClimbing
Random Search	Bergstra & Bengio (2012)	RandomSearch
SV-CMA-ES	Braun et al. (2024)	SV_CMA_ES
SV-OpenAI-ES	Liu et al. (2017)	SV_OpenES
Simple Genetic Algorithm	Such et al. (2017)	SimpleGA
MR15-GA	Rechenberg (1978)	MR15_GA
SAMR-GA	Clune et al. (2008)	SAMR_GA
GESMR-GA	Kumar et al. (2022)	GESMR_GA
LGA	Lange et al. (2023b)	LearnedGA
Diffusion Evolution	Zhang et al. (2024)	DiffusionEvolution
Differential Evolution	Storn & Price (1997)	DifferentialEvolution
Particle Swarm Optimization	Kennedy & Eberhart (1995)	PSO

Installation ⏳

You will need Python 3.10 or later, and a working JAX installation.

Then, install evosax from PyPi:

pip install evosax

To upgrade to the latest version of evosax, you can use:

pip install git+https://github.com/RobertTLange/evosax.git@main

Examples 📖

• 📓 Getting Started - Introduction to the library
• 📓 Black Box Optimization Benchmark - Optimization of common test functions
• 📓 Reinforcement Learning - Learning MLP control policies
• 📓 Vision - Training CNNs for classification
• 📓 Restart ES - Using simple, IPOP, and BIPOP restart controllers
• 📓 Diffusion Evolution - Optimization with diffusion evolution
• 📓 Stein Variational ES - Using SV-ES on BBOB problems
• 📓 Persistent/Noise-Reuse ES - ES for meta-learning problems
• 📓 Parallelization - ES with parallelization on multiple devices

Key Features 💎

• Comprehensive Algorithm Collection: 30+ classic and modern evolution strategies with a unified API
• JAX Acceleration: Fully compatible with JAX transformations for speed and scalability
• Vectorization & Parallelization: Fast execution on CPUs, GPUs, and TPUs
• Production Ready: Well-tested, documented, and used in research environments
• Batteries Included: Comes with optimizers like ClipUp, fitness shaping, and restart controllers

Related Resources 📚

• 📺 Rob's MLC Research Jam Talk - Overview at the ML Collective Research Jam
• 📝 Rob's 02/2021 Blog - Blog post on implementing CMA-ES in JAX
• 💻 Evojax - Hardware-Accelerated Neuroevolution with great rollout wrappers.
• 💻 QDax: Quality-Diversity algorithms in JAX.

Citing evosax ✏️

If you use evosax in your research, please cite the following paper:

@article{evosax2022github,
    author  = {Robert Tjarko Lange},
    title   = {evosax: JAX-based Evolution Strategies},
    journal = {arXiv preprint arXiv:2212.04180},
    year    = {2022},
}

Acknowledgements 🙏

We acknowledge financial support by the Google TRC and the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany's Excellence Strategy - EXC 2002/1 "Science of Intelligence" - project number 390523135.

Contributing 👷

Contributions are welcome! If you find a bug or are missing your favorite feature, please open an issue or submit a pull request following our contribution guidelines 🤗.

Disclaimer ⚠️

This repository contains independent reimplementations of LES and DES based and is unrelated to Google DeepMind. The implementation has been tested to reproduce the official results on a range of tasks.