RobertTLange/evosax

解决的问题

evosax 的设计旨在消除在实现神经进化时处理复杂异步过程的需求。它使研究人员能够利用大规模向量化和高吞吐量硬件加速器（如 GPU 和 TPU）进行进化策略（ES），而无需承担通常与分布式实现相关的开销。

工作原理

该库基于 JAX 构建，利用 XLA 编译和 JAX 的变换原语（如 jit、vmap 和 lax.scan）实现高性能。它遵循标准化的 ask-eval-tell API 循环：

1. Ask：策略生成一组待评估的候选解（种群）。
2. Eval：用户评估该种群的适应度。
3. Tell：用户将适应度结果反馈给策略，以更新其内部状态并为下一代进化种群。

适用人群

适用于希望使用高性能、JAX 原生库将进化算法和神经进化实验扩展到现代硬件加速器的研究人员和实践者。

主要亮点

• 丰富的算法库：包含 30 多种已实现的进化策略，从 CMA-ES 和差分进化等经典方法，到 OpenAI-ES 和扩散进化等现代方法。
• JAX 原生支持：完全支持 JAX 变换，实现高效扩展和 XLA 编译。
• 标准化 API：在不同策略之间（包括基于种群的算法如 SimpleGA）采用一致的 ask-eval-tell 循环。
• 高吞吐量：针对硬件加速器进行优化，避免传统分布式系统的开销。