Package: evosax
evosax: JAX-Based Evolution Strategies 🦎Tired of having to handle asynchronous processes for neuroevolution? Do you want to leverage massive vectorization and high-throughput accelerators for evolution strategies (ES)?evosaxallows you to leverage JAX, XLA compilation and auto-vectorization/parallelization to scale ES to your favorite accelerators. The API is based on the classicalask,evaluate,tellcycle of ES. Bothaskandtellcalls are compatible withjit,vmap/pmapandlax.scan. It includes a vast set of both classic (e.g. CMA-ES, Differential Evolution, etc.) and modern neuroevolution (e.g. OpenAI-ES, Augmented RS, etc.) strategies. You can get started here 👉BasicevosaxAPI Usage 🍲importjaxfromevosaximportCMA_ES# Instantiate the search strategyrng=jax.random.PRNGKey(0)strategy=CMA_ES(popsize=20,num_dims=2,elite_ratio=0.5)es_params=strategy.default_paramsstate=strategy.initialize(rng,es_params)# Run ask-eval-tell loop - NOTE: By default minimization!fortinrange(num_generations):rng,rng_gen,rng_eval=jax.random.split(rng,3)x,state=strategy.ask(rng_gen,state,es_params)fitness=...# Your population evaluation fctstate=strategy.tell(x,fitness,state,es_params)# Get best overall population member & its fitnessstate.best_member,state.best_fitnessImplemented Evolution Strategies 🦎StrategyReferenceImportExampleOpenAI-ESSalimans et al. (2017)OpenESPGPESehnke et al. (2010)PGPEARSMania et al. (2018)ARSESMCMerchant et al. (2021)ESMCPersistent ESVicol et al. (2021)PersistentESNoise-Reuse ESLi et al. (2023)NoiseReuseESxNESWierstra et al. (2014)XNESSNESWierstra et al. (2014)SNESCR-FM-NESNomura & Ono (2022)CR_FM_NESGuided ESMaheswaranathan et al. (2018)GuidedESASEBOChoromanski et al. (2019)ASEBOCMA-ESHansen & Ostermeier (2001)CMA_ESSep-CMA-ESRos & Hansen (2008)Sep_CMA_ESBIPOP-CMA-ESHansen (2009)BIPOP_CMA_ESIPOP-CMA-ESAuer & Hansen (2005)IPOP_CMA_ESFull-iAMaLGaMBosman et al. (2013)Full_iAMaLGaMIndependent-iAMaLGaMBosman et al. (2013)Indep_iAMaLGaMMA-ESBayer & Sendhoff (2017)MA_ESLM-MA-ESLoshchilov et al. (2017)LM_MA_ESRmESLi & Zhang (2017)RmESSimple GeneticSuch et al. (2017)SimpleGASAMR-GAClune et al. (2008)SAMR_GAGESMR-GAKumar et al. (2022)GESMR_GAMR15-GARechenberg (1978)MR15_GALGALange et al. (2023b)LGASimple GaussianRechenberg (1978)SimpleESDESLange et al. (2023a)DESLESLange et al. (2023a)LESParticle Swarm OptimizationKennedy & Eberhart (1995)PSODifferential EvolutionStorn & Price (1997)DEGLDGolovin et al. (2019)GLDSimulated AnnealingRasdi Rere et al. (2015)SimAnnealPopulation-Based TrainingJaderberg et al. (2017)PBTRandom SearchBergstra & Bengio (2012)RandomSearchInstallation ⏳The latestevosaxrelease can directly be installed from PyPI:pip install evosaxIf you want to get the most recent commit, please install directly from the repository:pip install git+https://github.com/RobertTLange/evosax.git@mainIn order to use JAX on your accelerators, you can find more details in theJAX documentation.Examples 📖📓Classic ES Tasks: API introduction on Rosenbrock function (CMA-ES, Simple GA, etc.).📓CartPole-Control: OpenES & PEPG on theCartPole-v1gym task (MLP/LSTM controller).📓MNIST-Classifier: OpenES on MNIST with CNN network.📓LRateTune-PES: Persistent/Noise-Reuse ES on meta-learning problem as inVicol et al. (2021).📓Quadratic-PBT: PBT on toy quadratic problem as inJaderberg et al. (2017).📓Restart-Wrappers: Custom restart wrappers as e.g. used in (B)IPOP-CMA-ES.📓Brax Control: Evolve Tanh MLPs on Brax tasks using theEvoJAXwrapper.📓BBOB Visualizer: Visualize evolution rollouts on 2D fitness landscapes.Key Features 💵Strategy Diversity:evosaximplements more than 30 classical and modern neuroevolution strategies. All of them follow the same simpleask/evalAPI and come with tailored tools such as theClipUpoptimizer, parameter reshaping into PyTrees and fitness shaping (see below).Vectorization/Parallelization ofask/tellCalls: Bothaskandtellcalls can leveragejit,vmap/pmap. This enables vectorized/parallel rollouts of different evolution strategies.fromevosax.strategies.arsimportARS,EvoParams# E.g. vectorize over different initial perturbation stdsstrategy=ARS(popsize=100,num_dims=20)es_params=EvoParams(sigma_init=jnp.array([0.1,0.01,0.001]),sigma_decay=0.999,...)# Specify how to map over ES hyperparametersmap_dict=EvoParams(sigma_init=0,sigma_decay=None,...)# Vmap-composed batch initialize, ask and tell functionsbatch_init=jax.vmap(strategy.init,in_axes=(None,map_dict))batch_ask=jax.vmap(strategy.ask,in_axes=(None,0,map_dict))batch_tell=jax.vmap(strategy.tell,in_axes=(0,0,0,map_dict))Scan Through Evolution Rollouts: You can alsolax.scanthrough entireinit,ask,eval,tellloops for fast compilation of ES loops:@partial(jax.jit,static_argnums=(1,))defrun_es_loop(rng,num_steps):"""Run evolution ask-eval-tell loop."""es_params=strategy.default_paramsstate=strategy.initialize(rng,es_params)defes_step(state_input,tmp):"""Helper es step to lax.scan through."""rng,state=state_inputrng,rng_iter=jax.random.split(rng)x,state=strategy.ask(rng_iter,state,es_params)fitness=...state=strategy.tell(y,fitness,state,es_params)return[rng,state],fitness[jnp.argmin(fitness)]_,scan_out=jax.lax.scan(es_step,[rng,state],[jnp.zeros(num_steps)])returnjnp.min(scan_out)Population Parameter Reshaping: We provide aParamaterReshaperwrapper to reshape flat parameter vectors into PyTrees. The wrapper is compatible with JAX neural network libraries such as Flax/Haiku and makes it easier to afterwards evaluate network populations.fromflaximportlinenasnnfromevosaximportParameterReshaperclassMLP(nn.Module):num_hidden_units:int...@nn.compactdef__call__(self,obs):...return...network=MLP(64)net_params=network.init(rng,jnp.zeros(4,),rng)# Initialize reshaper based on placeholder network shapesparam_reshaper=ParameterReshaper(net_params)# Get population candidates & reshape into stacked pytreesx=strategy.ask(...)x_shaped=param_reshaper.reshape(x)Flexible Fitness Shaping: By defaultevosaxassumes that the fitness objective is to be minimized. If you would like to maximize instead, perform rank centering, z-scoring or add weight regularization you can use theFitnessShaper:fromevosaximportFitnessShaper# Instantiate jittable fitness shaper (e.g. for Open ES)fit_shaper=FitnessShaper(centered_rank=True,z_score=False,weight_decay=0.01,maximize=True)# Shape the evaluated fitness scoresfit_shaped=fit_shaper.apply(x,fitness)Additonal Work-In-Progress**Strategy Restart Wrappers**: *Work-in-progress*. You can also choose from a set of different restart mechanisms, which will relaunch a strategy (with e.g. new population size) based on termination criteria. Note: For all restart strategies which alter the population size the ask and tell methods will have to be re-compiled at the time of change. Note that all strategies can also be executed without explicitly providing `es_params`. In this case the default parameters will be used.```Python
from evosax import CMA_ES
from evosax.restarts import BIPOP_Restarter

# Define a termination criterion (kwargs - fitness, state, params)
def std_criterion(fitness, state, params):
    """Restart strategy if fitness std across population is small."""
    return fitness.std() < 0.001

# Instantiate Base CMA-ES & wrap with BIPOP restarts
# Pass strategy-specific kwargs separately (e.g. elite_ration or opt_name)
strategy = CMA_ES(num_dims, popsize, elite_ratio)
re_strategy = BIPOP_Restarter(
                strategy,
                stop_criteria=[std_criterion],
                strategy_kwargs={"elite_ratio": elite_ratio}
            )
state = re_strategy.initialize(rng)

# ask/tell loop - restarts are automatically handled 
rng, rng_gen, rng_eval = jax.random.split(rng, 3)
x, state = re_strategy.ask(rng_gen, state)
fitness = ...  # Your population evaluation fct 
state = re_strategy.tell(x, fitness, state)
```

- **Batch Strategy Rollouts**: *Work-in-progress*. We are currently also working on different ways of incorporating multiple subpopulations with different communication protocols.

```Python
from evosax.experimental.subpops import BatchStrategy

# Instantiates 5 CMA-ES subpops of 20 members
strategy = BatchStrategy(
        strategy_name="CMA_ES",
        num_dims=4096,
        popsize=100,
        num_subpops=5,
        strategy_kwargs={"elite_ratio": 0.5},
        communication="best_subpop",
    )

state = strategy.initialize(rng)
# Ask for evaluation candidates of different subpopulation ES
x, state = strategy.ask(rng_iter, state)
fitness = ...
state = strategy.tell(x, fitness, state)
```

- **Indirect Encodings**: *Work-in-progress*. ES can struggle with high-dimensional search spaces (e.g. due to harder estimation of covariances). One potential way to alleviate this challenge, is to use indirect parameter encodings in a lower dimensional space. So far we provide JAX-compatible encodings with random projections (Gaussian/Rademacher) and Hypernetworks for MLPs. They act as drop-in replacements for the `ParameterReshaper`:

```Python
from evosax.experimental.decodings import RandomDecoder, HyperDecoder

# For arbitrary network architectures / search spaces
num_encoding_dims = 6
param_reshaper = RandomDecoder(num_encoding_dims, net_params)
x_shaped = param_reshaper.reshape(x)

# For MLP-based models we also support a HyperNetwork en/decoding
reshaper = HyperDecoder(
        net_params,
        hypernet_config={
            "num_latent_units": 3,  # Latent units per module kernel/bias
            "num_hidden_units": 2,  # Hidden dimensionality of a_i^j embedding
        },
    )
x_shaped = param_reshaper.reshape(x)
```Resources & Other Great JAX-ES Tools 📝📺Rob's MLC Research Jam Talk: Small motivation talk at the ML Collective Research Jam.📝Rob's 02/2021 Blog: Tutorial on CMA-ES & leveraging JAX's primitives.💻Evojax: JAX-ES library by Google Brain with great rollout wrappers.💻QDax: Quality-Diversity algorithms in JAX.Acknowledgements & Citingevosax✏️If you useevosaxin 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 TRCand the Deutsche
Forschungsgemeinschaft (DFG, German Research Foundation) under Germany's Excellence Strategy - EXC 2002/1"Science of Intelligence"- project number 390523135.Development 👷You can run the test suite viapython -m pytest -vv --all. If you find a bug or are missing your favourite feature, feel free to create an issue and/or startcontributing🤗.Disclaimer ⚠️This repository contains an independent reimplementation of LES and DES based on the corresponding ICLR 2023 publication(Lange et al., 2023). It is unrelated to Google or DeepMind. The implementation has been tested to roughly reproduce the official results on a range of tasks.