evosax

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

…

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