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, 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)	CR_FM_NES
Guided ES	Maheswaranathan et al. (2018)	GuidedES
ASEBO	Choromanski et al. (2019)	ASEBO
Discovered ES	Lange et al. (2023a)	DES
Learned ES	Lange et al. (2023a)	LES
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)	GLD
Simulated Annealing	Rasdi Rere et al. (2015)	SimAnneal
Hill Climbing	Rasdi Rere et al. (2015)	SimAnneal
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)	LGA
Diffusion Evolution	Zhang et al. (2024)	DiffusionEvolution
Differential Evolution	Storn & Price (1997)	DE
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 - Implementing restart strategies
• 📓 Diffusion Evolution - Optimization with diffusion evolution.
• 📓 Stein Variational ES - Using SV-ES on BBOB problems
• 📓 Persistent/Noise-Reuse ES - ES for meta-learning problems

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 strategies

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.