evosax

evosax: JAX 中的进化策略 你是否厌倦了需要处理异步过程来进行神经进化?你是否希望利用大规模向量化和高吞吐量的加速器来实现进化策略?evosax 提供了一个全面、高性能的库,在 JAX 中实现了进化策略(ES)。通过利用 XLA 编译和 JAX 的转换原语,evosax 使研究人员和从业者能够高效地将进化算法扩展到现代硬件加速器,而无需传统的分布式实现带来的开销。

API 遵循 ES 的经典 ask - eval - tell 循环,完全支持 JAX 的转换(jit,vmap,lax.scan)。库中包含 30 多种进化策略,从经典的 CMA-ES 和差分进化到现代的 OpenAI-ES 和扩散进化。

fitness_init 参数会为第一次调用 ask(...) 所使用的存储的种群排名种子。初始化后,ga.ask(...) 返回下一个需要评估的种群,而 ga.tell(...) 会使用已评估的种群和适应度来更新算法状态。