evosax: Evolution Strategies in JAX 🦎

April 2, 2026 · View on GitHub

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
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 - 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
📓 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 strategies

📺 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},
}