RobertTLange/evosax

evosax – JAX 기반 진화 전략(Evolution Strategies)

개요

• JAX를 기반으로 30개 이상의 현대적이고 고전적인 진화 전략(ES) 및 모집단 기반 최적화 알고리즘을 구현한 순수 Python 라이브러리입니다. JAX의 XLA 컴파일 및 변환 프리미티브(jit, vmap, lax.scan)를 사용하여 별도의 분산 코드 작성 없이도 CPU, GPU, TPU에서 알고리즘을 효율적으로 실행할 수 있습니다.

중요성

• ES는 신경 진화, 강화 학습 정책 탐색, 하이퍼파라미터 튜닝, 블랙박스 최적화 등에 널리 사용되는 그래디언트가 필요 없는 최적화 방법입니다. evosax는 이러한 방법들을 벡터화하고 하드웨어 가속을 지원하여, Python 수준의 루프나 수동 병렬화로 인해 제한되었던 실험을 연구자들이 대규모로 확장할 수 있게 해줍니다.

주요 기능

기능	세부 정보
광범위한 알고리즘 카탈로그	CMA‑ES, Sep‑CMA‑ES, xNES, SNES, OpenAI‑ES, ARS, Guided‑ES, ASEBO, Persistent‑ES, Noise‑Reuse‑ES, SV‑CMA‑ES, SimpleGA, MR15‑GA 등을 포함한 30개 이상의 전략
JAX 네이티브 API	고전적인 ask → evaluate → tell 루프를 사용하며, 모든 구성 요소는 JAX와 호환되어 jit, vmap, scan을 통해 대규모 병렬 처리가 가능합니다.
상태 비저장(Stateless) 설계	알고리즘 상태가 명시적인 불변 객체로 관리되어 체크포인트 생성, 재생 또는 함수형 파이프라인과의 통합이 용이합니다.
즉시 실행 가능한 노트북	모든 알고리즘에 대한 Colab 링크(예: BBOB 벤치마크, RL 정책 탐색, 지속적 ES)를 제공하여 별도의 설정 없이 라이브러리를 체험할 수 있습니다.
테스트 및 CI 보호	지속적 통합(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('Best solution:', state.best_solution)
print('Best fitness :', state.best_fitness)

이 루프는 교과서적인 ES 워크플로우를 반영하지만, JAX 덕분에 전체 세대를 단일 벡터화 호출로 평가할 수 있습니다.

문서 및 예제

• 저장소에는 로고와 배지가 포함된 docs/ 폴더가 있으며, 모든 알고리즘 항목에는 Colab 노트북(examples/)으로 연결되는 직접 링크가 포함되어 있습니다.
• 함께 제공되는 arXiv 논문(2212.04180)은 설계 철학과 성능 벤치마크를 설명합니다.

추천 대상

• 딥러닝이나 과학적 컴퓨팅을 위해 이미 JAX를 사용 중이며, 빠르고 확장 가능한 블랙박스 최적화 도구가 필요한 연구자 및 엔지니어.
• 동일한 하드웨어 스택에서 ES와 그래디언트 기반 방법을 비교하는 데 관심이 있는 모든 사람.

---

위의 모든 정보는 프로젝트의 README에서 직접 가져온 것이며, 외부 주장은 추가되지 않았습니다.