# evosax: JAX-based Evolution Strategies

arXiv: https://arxiv.org/abs/2212.04180 · HTML: https://mlsanity.com/paper/2212.04180/

## Key facts

- Cited by 41 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 49 references that are in the map.
- First posted on arXiv 8 Dec 2022.

## Builds on (most cited first, 18 of 18)

- [Evolution Strategies as a Scalable Alternative to Reinforcement Learning](https://mlsanity.com/paper/1703.03864/) (2017)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Population Based Training of Neural Networks](https://mlsanity.com/paper/1711.09846/) (2017)
- [The Hardware Lottery](https://mlsanity.com/paper/2009.06489/) (2020)
- [Deep Neuroevolution: Genetic Algorithms Are a Competitive Alternative for Training Deep Neural Networks for Reinforcement Learning](https://mlsanity.com/paper/1712.06567/) (2017)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies](https://mlsanity.com/paper/2112.13835/) (2021)
- [Discovering Evolution Strategies via Meta-Black-Box Optimization](https://mlsanity.com/paper/2211.11260/) (2022)
- [Gradientless Descent: High-Dimensional Zeroth-Order Optimization](https://mlsanity.com/paper/1911.06317/) (2019)
- [Practical tradeoffs between memory, compute, and performance in learned optimizers](https://mlsanity.com/paper/2203.11860/) (2022)
- [EvoJAX: Hardware-Accelerated Neuroevolution](https://mlsanity.com/paper/2202.05008/) (2022)
- [Accelerated Quality-Diversity through Massive Parallelism](https://mlsanity.com/paper/2202.01258/) (2022)
- [Learn2Hop: Learned Optimization on Rough Landscapes](https://mlsanity.com/paper/2107.09661/) (2021)
- [ClipUp: A Simple and Powerful Optimizer for Distribution-based Policy Evolution](https://mlsanity.com/paper/2008.02387/) (2020)
- [Effective Mutation Rate Adaptation through Group Elite Selection](https://mlsanity.com/paper/2204.04817/) (2022)
- [Limited-Memory Matrix Adaptation for Large Scale Black-box Optimization](https://mlsanity.com/paper/1705.06693/) (2017)
- [Fast Moving Natural Evolution Strategy for High-Dimensional Problems](https://mlsanity.com/paper/2201.11422/) (2022)

## Built on by (most cited first, 30 of 41)

- [LLaMoCo: Instruction Tuning of Large Language Models for Optimization Code Generation](https://mlsanity.com/paper/2403.01131/) (2024)
- [Discovering Attention-Based Genetic Algorithms via Meta-Black-Box Optimization](https://mlsanity.com/paper/2304.03995/) (2023)
- [Evolutionary Reinforcement Learning: A Survey](https://mlsanity.com/paper/2303.04150/) (2023)
- [Discovering Temporally-Aware Reinforcement Learning Algorithms](https://mlsanity.com/paper/2402.05828/) (2024)
- [Can Learned Optimization Make Reinforcement Learning Less Difficult?](https://mlsanity.com/paper/2407.07082/) (2024)
- [Scaling Opponent Shaping to High Dimensional Games](https://mlsanity.com/paper/2312.12568/) (2023)
- [Flow-Lenia: Towards open-ended evolution in cellular automata through mass conservation and parameter localization](https://mlsanity.com/paper/2212.07906/) (2022)
- [minimax: Efficient Baselines for Autocurricula in JAX](https://mlsanity.com/paper/2311.12716/) (2023)
- [EvoTorch: Scalable Evolutionary Computation in Python](https://mlsanity.com/paper/2302.12600/) (2023)
- [EvIL: Evolution Strategies for Generalisable Imitation Learning](https://mlsanity.com/paper/2406.11905/) (2024)
- [CAX: Cellular Automata Accelerated in JAX](https://mlsanity.com/paper/2410.02651/) (2024)
- [Towards Universal Offline Black-Box Optimization via Learning Language Model Embeddings](https://mlsanity.com/paper/2506.07109/) (2025)
- [EvoFed: Leveraging Evolutionary Strategies for Communication-Efficient Federated Learning](https://mlsanity.com/paper/2311.07485/) (2023)
- [Behaviour Distillation](https://mlsanity.com/paper/2406.15042/) (2024)
- [How Should We Meta-Learn Reinforcement Learning Algorithms?](https://mlsanity.com/paper/2507.17668/) (2025)
- [Creative Text-to-Audio Generation via Synthesizer Programming](https://mlsanity.com/paper/2406.00294/) (2024)
- [Evolving Self-Assembling Neural Networks: From Spontaneous Activity to Experience-Dependent Learning](https://mlsanity.com/paper/2406.09787/) (2024)
- [Enabling Population-Level Parallelism in Tree-Based Genetic Programming for GPU Acceleration](https://mlsanity.com/paper/2501.17168/) (2025)
- [Kozax: Flexible and Scalable Genetic Programming in JAX](https://mlsanity.com/paper/2502.03047/) (2025)
- [Going faster to see further: GPU-accelerated value iteration and simulation for perishable inventory control using JAX](https://mlsanity.com/paper/2303.10672/) (2023)
- [Solving Deep Reinforcement Learning Tasks with Evolution Strategies and Linear Policy Networks](https://mlsanity.com/paper/2402.06912/) (2024)
- [Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks](https://mlsanity.com/paper/2402.18994/) (2024)
- [Gradient-free neural topology optimization: Towards effective fracture-resistant designs](https://mlsanity.com/paper/2403.04937/) (2024)
- [Stein Variational Evolution Strategies](https://mlsanity.com/paper/2410.10390/) (2024)
- [Meta-Learning Objectives for Preference Optimization](https://mlsanity.com/paper/2411.06568/) (2024)
- [Flow-Lenia: Emergent evolutionary dynamics in mass conservative continuous cellular automata](https://mlsanity.com/paper/2506.08569/) (2025)
- [Hypernetworks That Evolve Themselves](https://mlsanity.com/paper/2512.16406/) (2025)
- [Learning mirror maps in policy mirror descent](https://mlsanity.com/paper/2402.05187/) (2024)
- [Learning Evolution via Optimization Knowledge Adaptation](https://mlsanity.com/paper/2501.02200/) (2025)
- [A Learning-Based Cooperative Coevolution Framework for Heterogeneous Large-Scale Global Optimization](https://mlsanity.com/paper/2604.01241/) (2026)

## Abstract (arXiv)

The deep learning revolution has greatly been accelerated by the 'hardware lottery': Recent advances in modern hardware accelerators and compilers paved the way for large-scale batch gradient optimization. Evolutionary optimization, on the other hand, has mainly relied on CPU-parallelism, e.g. using Dask scheduling and distributed multi-host infrastructure. Here we argue that also modern evolutionary computation can significantly benefit from the massive computational throughput provided by GPUs and TPUs. In order to better harness these resources and to enable the next generation of black-box optimization algorithms, we release evosax: A JAX-based library of evolution strategies which allows researchers to leverage powerful function transformations such as just-in-time compilation, automatic vectorization and hardware parallelization. evosax implements 30 evolutionary optimization algorithms including finite-difference-based, estimation-of-distribution evolution strategies and various genetic algorithms. Every single algorithm can directly be executed on hardware accelerators and automatically vectorized or parallelized across devices using a single line of code. It is designed in a modular fashion and allows for flexible usage via a simple ask-evaluate-tell API. We thereby hope to facilitate a new wave of scalable evolutionary optimization algorithms.