evosax: Thirty Evolution Strategies That Run Inside JAX

The problem it removes is the asynchronous plumbing

evosax exists because of a specific irritation: handling asynchronous processes for neuroevolution. The pitch is that you want to use massive vectorization and high-throughput accelerators for Evolution Strategies, and the library provides a comprehensive, high-performance implementation of Evolution Strategies in JAX. The mechanism is XLA compilation and JAX's transformation primitives, which allow evolutionary algorithms to scale to modern hardware accelerators without the overhead of distributed implementations. In practice that means the population of candidates is not a list of processes to communicate with but an array to evaluate in one vectorized call, and the update to the strategy state is a function you can compile. That is the whole proposition, and it is why the library is written in JAX rather than as a set of numpy routines with an accelerator backend.

Ask, evaluate, tell, and nothing else in the interface

The API follows the classical ask-eval-tell cycle of Evolution Strategies, with full support for JAX's transformations including jit, vmap and lax.scan. The shape of a run is short: instantiate the strategy with a population size and a solution template, take its default parameters, create a random key, initialise the state, and then loop. Inside the loop the key is split, the strategy asks for a population, you evaluate the fitness of that population however you like, and you tell the strategy the results.

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)

The last line of the snippet is the payoff: the best solution and its fitness are read straight off the state object after the loop ends.

Population-based strategies initialise differently

Not every strategy initialises the same way, and the library is explicit about the difference. SimpleGA and other population-based algorithms use the same ask-tell loop, but their state is initialised from an explicit initial population and its fitness rather than from a template alone.

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 that, the pattern is the same as before: ask returns the next population to evaluate and tell updates the state with the evaluated population and its fitness. The `jax.tree.map` over the solution template is the idiomatic JAX way to build a batch of identical structures, and it works for any pytree, which is how the library supports parameter dictionaries as well as flat arrays.

Thirty strategies from Rechenberg to diffusion evolution

The library includes more than 30 evolution strategies, and the reference table names them with the original paper for each. The classics are there: Simple Evolution Strategy from Rechenberg in 1978, OpenAI-ES from Salimans and colleagues in 2017, CMA-ES from Hansen and Ostermeier in 2001, and the separable variant Sep-CMA-ES from Ros and Hansen in 2008. There are the natural evolution strategy family from Wierstra and colleagues in 2014, with xNES and SNES among them, and the table continues into modern approaches including Diffusion Evolution. Every row links to the source paper, the import path inside the package and a Colab notebook that runs it, so you can read the method and then execute it without leaving the browser. That combination of a citation and a runnable notebook per strategy is unusually useful for a library whose point is comparison.

A checkpoint mechanism and the newer strategies

The package data configuration shows that several algorithms ship pretrained checkpoint files inside the distribution, with directories for les, lga, evotf and cma_anchored_les holding pickle files, plus a JSON file for the CMA-anchored variant. That is a hint at what the more recent releases added. The version history is informative: v0.3.0 was titled restarts, parallelization and optimizer fixes, v0.3.1 modern JAX and Gymnax compatibility fixes, and v0.3.2 CMA-anchored learned evolution, published on 2026-10-11. The presence of a CMA-anchored learned evolution strategy, with checkpoints, suggests the library now includes methods that learn an anchoring distribution rather than hand-tuning one. The example directory backs that up with notebooks for restart ES, diffusion evolution, separable ES, persistent ES and parallelization, plus Python scripts for a CMA-anchored approach on Brax and a learned sphere problem.

What you need installed to run the examples

The core install is small. The project metadata lists numpy, jax, flax, dotmap and matplotlib as dependencies, with jax at 0.5.0 or newer and flax at 0.10.0 or newer, and it requires Python 3.10 or newer, with classifiers for 3.10 through 3.13. The build backend is setuptools. The optional extras split by purpose: a dev group with pytest, pytest-cov, ruff and mypy, and an examples group with torch, torchvision, brax, gymnax, mediapy, ipywidgets and ipykernel. That tells you the RL and vision examples are built on Brax and Gymnax, which is where the library's benchmark and environment work comes from. For linting, ruff is configured with a line length of 88 and a targeted rule selection, and mypy runs with ignore_missing_imports set, which is the pragmatic setting for a library whose dependencies ship their own type information inconsistently.

Where the library fits against other approaches

The alternative approaches to a large-scale evolutionary search are a plain numpy implementation you write yourself, a distributed framework such as Ray or a dedicated evolutionary framework, and a gradient-based optimiser. The difference evosax offers is that one implementation of the ask-eval-tell interface is shared by more than thirty strategies, and that interface is the same one you would write against a numpy implementation, so swapping a strategy in or out does not change your loop. Compared with a distributed framework, the win is that there is no process management at all, since the population is an array and the accelerator does the parallel work, and the cost is that you are bound by what JAX's transformations can express, which is why the library leans on vmap and lax.scan rather than on custom kernels. Compared with gradient descent, the difference is fundamental rather than incremental: evolution strategies need only a scalar fitness and no gradients, which is exactly what makes them usable on problems where the objective is not differentiable.

Licence, versioning and how to start

The licence is Apache-2.0, and the repository is not archived with the last push on 2026-10-11, the same day as the v0.3.2 release. Versions are numbered at 0.3.2, so the API is still moving, and the CHANGELOG.md is present for tracking that. The fastest way to start is the getting started notebook, hosted on Colab and linked from the readme, which needs no local setup and no GPU to begin with. Each strategy in the table also has its own Colab notebook, covering the black-box optimisation tasks, the reinforcement learning example and the vision example, so you can compare methods on the same task family rather than on tasks each paper chose. The examples directory in the repository mirrors those notebooks with numbered files, so the offline path is the same set of notebooks with a local kernel. Documentation lives in a docs/ directory, and the paper behind the library is on arXiv.

Editorial conclusion

evosax fits a researcher who wants to try several evolution strategies on the same problem without rewriting each one for accelerators, and who already works in JAX. It is the wrong choice if your optimiser needs gradients, since evolution strategies are gradient-free, and the wrong choice if you need a strategy that is not in the list of more than thirty implementations. Before you commit, run the getting started notebook in Colab on a problem you know, and compare against your current implementation rather than a benchmark table.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. README
4. Releases
5. RobertTLange/evosax on GitHub