ES with Parallelization

README.md

@@ -1,6 +1,6 @@
 # `evosax`: Evolution Strategies in JAX 🦎
 
-[![Pyversions](https://img.shields.io/pypi/pyversions/evosax.svg?style=flat-square)](https://pypi.python.org/pypi/evosax) [![PyPI version](https://badge.fury.io/py/evosax.svg)](https://badge.fury.io/py/evosax)
+[![Pyversions](https://img.shields.io/pypi/pyversions/evosax.svg?style=flat)](https://pypi.python.org/pypi/evosax) [![PyPI version](https://badge.fury.io/py/evosax.svg)](https://badge.fury.io/py/evosax)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![codecov](https://codecov.io/gh/RobertTLange/evosax/branch/main/graph/badge.svg?token=5FUSX35KWO)](https://codecov.io/gh/RobertTLange/evosax)
 [![Paper](http://img.shields.io/badge/paper-arxiv.2212.04180-B31B1B.svg)](http://arxiv.org/abs/2212.04180)
@@ -108,9 +108,10 @@ pip install git+https://github.com/RobertTLange/evosax.git@main
 * 📓 [Reinforcement Learning](https://github.com/RobertTLange/evosax/blob/main/examples/02_rl.ipynb) - Learning MLP control policies
 * 📓 [Vision](https://github.com/RobertTLange/evosax/blob/main/examples/03_vision.ipynb) - Training CNNs for classification
 * 📓 [Restart ES](https://github.com/RobertTLange/evosax/blob/main/examples/04_restart_es.ipynb) - Implementing restart strategies
-* 📓 [Diffusion Evolution](https://github.com/RobertTLange/evosax/blob/main/examples/05_diffusion_evolution.ipynb) - Optimization with diffusion evolution.
+* 📓 [Diffusion Evolution](https://github.com/RobertTLange/evosax/blob/main/examples/05_diffusion_evolution.ipynb) - Optimization with diffusion evolution
 * 📓 [Stein Variational ES](https://github.com/RobertTLange/evosax/blob/main/examples/06_sv_es.ipynb) - Using SV-ES on BBOB problems
 * 📓 [Persistent/Noise-Reuse ES](https://github.com/RobertTLange/evosax/blob/main/examples/07_persistent_es.ipynb) - ES for meta-learning problems
+* 📓 [Parallelization](https://github.com/RobertTLange/evosax/blob/main/examples/08_parallelization.ipynb) - ES with parallelization on multiple devices
 
 ## Key Features 💎
 

examples/02_rl.ipynb

@@ -149,8 +149,11 @@
    "source": [
     "from evosax.algorithms import SNES as ES\n",
     "\n",
+    "num_generations = 64\n",
+    "population_size = 128\n",
+    "\n",
     "es = ES(\n",
-    "    population_size=128,\n",
+    "    population_size=population_size,\n",
     "    solution=solution,\n",
     "    optimizer=optax.adam(learning_rate=0.01),\n",
     ")\n",
@@ -192,8 +195,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "num_generations = 64\n",
-    "\n",
     "key, subkey = jax.random.split(key)\n",
     "state = es.init(subkey, solution, params)\n",
     "\n",
@@ -312,6 +313,8 @@
     "from evosax.algorithms import Open_ES as ES\n",
     "\n",
     "num_generations = 512\n",
+    "population_size = 256\n",
+    "\n",
     "lr_schedule = optax.exponential_decay(\n",
     "    init_value=0.01,\n",
     "    transition_steps=num_generations,\n",
@@ -323,7 +326,7 @@
     "    decay_rate=0.2,\n",
     ")\n",
     "es = ES(\n",
-    "    population_size=256,\n",
+    "    population_size=population_size,\n",
     "    solution=solution,\n",
     "    optimizer=optax.adam(learning_rate=lr_schedule),\n",
     "    std_schedule=std_schedule,\n",
@@ -395,6 +398,7 @@
     "fitness_log = []\n",
     "log_period = 32\n",
     "for i in range(num_generations // log_period):\n",
+    "    # Train\n",
     "    key, subkey = jax.random.split(key)\n",
     "    keys = jax.random.split(subkey, log_period)\n",
     "    (state, params, problem_state), metrics = jax.lax.scan(\n",
@@ -403,8 +407,8 @@
     "        keys,\n",
     "    )\n",
     "\n",
+    "    # Eval\n",
     "    mean = es.get_mean(state)\n",
-    "\n",
     "    key, subkey = jax.random.split(key)\n",
     "    fitness, problem_state, info = problem.eval(\n",
     "        key, jax.tree.map(lambda x: x[None], mean), problem_state\n",
@@ -579,7 +583,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.8"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,