RobertTLange/evosax

evosax – JAXにおける進化戦略

概要

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`で変換して大規模な並列処理が可能です。
ステートレス設計	アルゴリズムの状態は明示的な不変オブジェクトであり、チェックポイントの作成、リプレイ、関数型パイプラインへの統合が容易です。
すぐに実行可能なノートブック	すべてのアルゴリズムに対してColabリンク（例：BBOBベンチマーク、RLポリシー探索、Persistent 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('Best solution:', state.best_solution)
print('Best fitness :', state.best_fitness)

このループは教科書的なESのワークフローを反映していますが、JAXのおかげで世代全体を単一のベクトル化された呼び出しで評価できます。

ドキュメントと例

リポジトリにはロゴやバッジを含む docs/ フォルダがあり、すべてのアルゴリズムのエントリにはColabノートブック（examples/）への直接リンクが含まれています。
付属のarXiv論文（2212.04180）では、設計思想とパフォーマンスのベンチマークについて説明しています。

推奨ユーザー

すでにディープラーニングや科学計算でJAXを使用しており、高速でスケーラブルなブラックボックスオプティマイザを必要としている研究者やエンジニア。
同じハードウェアスタック上でESと勾配ベースの手法を比較したいと考えている方。

上記の情報はすべてプロジェクトのREADMEから直接引用したものであり、外部からの主張は追加されていません。

RobertTLange/evosax

evosax – JAXにおける進化戦略

関連