evosax provides a high-performance, JAX-based library for implementing Evolution Strategies (ES). It targets researchers and practitioners in areas like neuroevolution and black-box optimization, enabling efficient scaling of evolutionary algorithms on modern hardware accelerators (CPUs, GPUs, TPUs) without complex distributed setups. The library offers a unified API for over 30 ES algorithms, including classics and recent advancements, fully leveraging JAX's compilation and transformation capabilities.

How It Works

evosax utilizes JAX's XLA compilation and transformation primitives (jit, vmap, lax.scan) to achieve high throughput and scalability. The core API follows the "ask-evaluate-tell" cycle common in ES. This design allows for efficient vectorization across populations and parallel execution on accelerators, abstracting away the complexities of distributed computing while enabling seamless integration with JAX's functional programming paradigm.

Quick Start & Requirements

• Install via pip: pip install evosax
• Requires Python 3.10+ and a working JAX installation.
• Upgrade to the latest version: pip install git+https://github.com/RobertTLange/evosax.git@main
• Examples and documentation are available: Basic evosax API Usage, Examples

Highlighted Details

• Implements over 30 classic and modern Evolution Strategies (e.g., CMA-ES, OpenAI-ES, Diffusion Evolution).
• Fully compatible with JAX transformations (jit, vmap, lax.scan) for acceleration and scalability.
• Optimized for high-throughput execution on CPUs, GPUs, and TPUs.
• Includes built-in optimizers (ClipUp), fitness shaping, and restart strategies.

Maintenance & Community

• Actively maintained by RobertTLange.
• Contributions are welcome via issues and pull requests.
• Related projects: Evojax, QDax.
• Citable via arXiv: evosax: JAX-based Evolution Strategies.

Licensing & Compatibility

• Licensed under the MIT License.
• Permissive license suitable for commercial use and integration into closed-source projects.

Limitations & Caveats

The library contains independent reimplementations of LES and DES and is stated to be unrelated to Google DeepMind. While tested to reproduce official results, users should be aware of potential discrepancies in highly specific research contexts.