evosax0.2.0

Published 2 months ago

Evolution Strategies in JAX

pip install evosax

Package Downloads

Authors

Robert Tjarko Lange

Project URLs

Homepage

Requires Python

>=3.10

Dependencies

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