Repository files navigation

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 classicalask-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 👉

importjaxfromevosax.algorithmsimportCMA_ES# Instantiate the search strategyes=CMA_ES(population_size=32,solution=dummy_solution)params=es.default_params# Initialize statekey=jax.random.key(0)state=es.init(key,params)# Ask-Eval-Tell loopforiinrange(num_generations):key,key_ask,key_eval=jax.random.split(key,3)# Generate a set of candidate solutions to evaluatepopulation,state=es.ask(key_ask,state,params)# Evaluate the fitness of the populationfitness= ...# Update the evolution strategystate=es.tell(population,fitness,state,params)# Get best solutionstate.best_solution,state.best_fitness

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

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

Then, installevosax from PyPi:

pip install evosax

To upgrade to the latest version ofevosax, you can use:

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

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

• 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.

If you useevosax in your research, please cite the followingpaper:

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

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

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

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.