RobertTLange/evosax

evosax – JAX 中的进化策略

它是什么

• 一个纯 Python 库，在 JAX 之上实现了超过 30 种现代和经典的进化策略 (ES) 以及基于种群的优化器。通过利用 JAX 的 XLA 编译和转换原语 (jit, vmap, lax.scan)，这些算法可以在 CPU、GPU 和 TPU 上高效运行，而无需编写自定义的分布式代码。

为什么它很重要

• ES 是一类无梯度优化方法，广泛应用于神经进化、强化学习策略搜索、超参数调优和黑盒优化。evosax 使这些方法实现了向量化和硬件加速，让研究人员能够扩展那些原本受限于 Python 循环或手动并行化的实验。

主要特性

特性	详情
广泛的算法目录	30 多种策略，包括 CMA‑ES, Sep‑CMA‑ES, xNES, SNES, OpenAI‑ES, ARS, Guided‑ES, ASEBO, Persistent‑ES, Noise‑Reuse‑ES, SV‑CMA‑ES, SimpleGA, MR15‑GA 等。
JAX 原生 API	使用经典的 ask → evaluate → tell 循环；所有组件均与 JAX 兼容，并可通过 jit, vmap, scan 进行转换以实现大规模并行。
无状态设计	算法状态是一个显式的不可变对象，易于进行检查点保存、重放或集成到函数式流水线中。
开箱即用的 Notebook	每个算法都提供了 Colab 链接（例如 BBOB 基准测试、RL 策略搜索、持久化 ES），无需任何配置即可尝试该库。
经过测试与 CI 保护	徽章显示了持续集成覆盖率、使用 Ruff 进行的代码检查以及已发布的 PyPI 包。
参考支持	每个实现都链接到原始论文，以便用户可以追溯确切的算法细节。

典型用例

• 神经进化：为 RL 智能体或监督任务进化神经网络权重或架构。
• 黑盒优化：在无法获取梯度的情况下优化昂贵的函数（例如模拟参数、超参数）。
• 基准测试：运行标准的 BBOB 或 RL 套件来比较不同的 ES 变体。
• 研究：通过继承基础 Algorithm 类并重用 JAX 工具来原型化新的 ES 想法。

安装

pip install evosax

该包需要 Python 3.8+ 和 JAX（根据需要选择 CPU、GPU 或 TPU 版本）。

快速开始

import jax
from evosax.algorithms import CMA_ES

# 虚拟目标函数 – 请替换为您自己的函数
def fitness_fn(params):
    return -jnp.sum(jnp.square(params))

# 创建一个 CMA‑ES 优化器
es = CMA_ES(population_size=32, solution=jax.random.normal(jax.random.key(0), (10,)))
params = es.default_params
key = jax.random.key(0)
state = es.init(key, es.solution, params)

for gen in range(100):
    key, ask_key, eval_key = jax.random.split(key, 3)
    pop, state = es.ask(ask_key, state, params)          # 提出候选解
    fitness = fitness_fn(pop)                           # 评估它们
    state, metrics = es.tell(eval_key, pop, fitness, state, params)

print('最佳解:', state.best_solution)
print('最佳适应度 :', state.best_fitness)

该循环反映了教科书式的 ES 工作流程，但得益于 JAX，整个种群的评估可以在一次向量化调用中完成。

文档与示例

• 该仓库在 docs/ 文件夹中提供了徽标和徽章，每个算法条目都包含一个指向 Colab notebook (examples/) 的直接链接。
• 随附的 arXiv 论文 (2212.04180) 描述了其设计理念和性能基准。

谁应该尝试它

• 已经在深度学习或科学计算中使用 JAX，并且需要快速、可扩展的黑盒优化器的研究人员和工程师。
• 任何有兴趣在同一硬件栈上比较 ES 与基于梯度的方法的人。

---

以上所有信息均直接取自该项目的 README；未添加任何外部声明。