Training an SNN with Neuroevolution¶

This tutorial trains a spiking controller for the gymnax CartPole-v1 environment using a black-box evolutionary strategy from evosax. The controller is a small spiking network defined with the new Spyx Flax-NNX API.

Required extras:

pip install gymnax evosax

In [1]:

Copied!





import os

os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", ".70")

import jax
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
from flax import nnx
from tqdm import trange

# evosax >= 0.2 moved strategies under `evosax.algorithms.*`. Older releases
# expose them at the package root; pick whichever import works.
try:
    from evosax.algorithms.distribution_based.cma_es import CMA_ES
except ImportError:
    from evosax import CMA_ES

import gymnax

import spyx.nn as snn
import os

os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", ".70")

import jax
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
from flax import nnx
from tqdm import trange

# evosax >= 0.2 moved strategies under `evosax.algorithms.*`. Older releases
# expose them at the package root; pick whichever import works.
try:
    from evosax.algorithms.distribution_based.cma_es import CMA_ES
except ImportError:
    from evosax import CMA_ES

import gymnax

import spyx.nn as snn

/home/kade/Code/spyx/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Environment¶

Spin up CartPole-v1 and take a single step to inspect what gymnax returns.

In [2]:

Copied!





rng = jax.random.PRNGKey(0)
rng, key_reset, key_act, key_step = jax.random.split(rng, 4)

env, env_params = gymnax.make("CartPole-v1")

obs, env_state = env.reset(key_reset, env_params)
action = env.action_space(env_params).sample(key_act)
n_obs, n_state, reward, done, _ = env.step(key_step, env_state, action, env_params)
print("obs:", obs.shape, "reward:", float(reward), "done:", bool(done))
rng = jax.random.PRNGKey(0)
rng, key_reset, key_act, key_step = jax.random.split(rng, 4)

env, env_params = gymnax.make("CartPole-v1")

obs, env_state = env.reset(key_reset, env_params)
action = env.action_space(env_params).sample(key_act)
n_obs, n_state, reward, done, _ = env.step(key_step, env_state, action, env_params)
print("obs:", obs.shape, "reward:", float(reward), "done:", bool(done))

obs: (4,) reward: 1.0 done: False

/home/kade/Code/spyx/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:125: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return lax_numpy.astype(self, dtype, copy=copy, device=device)

Spike encoder¶

CartPole produces continuous observations. We bin each component (cart position, cart velocity, pole angle, pole angular velocity) into one-hot bins so the controller receives a sparse spike vector at every timestep.

In [3]:

Copied!





class Binarize:
    def __init__(self, bins, lo, hi):
        self.bins = bins
        self.edges = jnp.linspace(lo, hi, bins)

    def __call__(self, value):
        idx = jnp.digitize(value, self.edges)
        return jax.nn.one_hot(idx, self.bins + 1)


class NeuromorphicCartpole:
    """Adapter mapping a 4-D continuous CartPole obs to a binary spike vector."""

    def __init__(self, bins=16):
        self.cart_v = Binarize(bins, -3.5, 3.5)
        self.pole_a = Binarize(bins, -0.21, 0.21)
        self.pole_w = Binarize(bins, -3.5, 3.5)

    def __call__(self, obs):
        return jnp.concatenate([
            self.cart_v(obs[1]),
            self.pole_a(obs[2]),
            self.pole_w(obs[3]),
        ])


adapter = NeuromorphicCartpole(bins=16)
input_dim = adapter(obs).shape[0]
print("encoded input dim:", input_dim)
class Binarize:
    def __init__(self, bins, lo, hi):
        self.bins = bins
        self.edges = jnp.linspace(lo, hi, bins)

    def __call__(self, value):
        idx = jnp.digitize(value, self.edges)
        return jax.nn.one_hot(idx, self.bins + 1)


class NeuromorphicCartpole:
    """Adapter mapping a 4-D continuous CartPole obs to a binary spike vector."""

    def __init__(self, bins=16):
        self.cart_v = Binarize(bins, -3.5, 3.5)
        self.pole_a = Binarize(bins, -0.21, 0.21)
        self.pole_w = Binarize(bins, -3.5, 3.5)

    def __call__(self, obs):
        return jnp.concatenate([
            self.cart_v(obs[1]),
            self.pole_a(obs[2]),
            self.pole_w(obs[3]),
        ])


adapter = NeuromorphicCartpole(bins=16)
input_dim = adapter(obs).shape[0]
print("encoded input dim:", input_dim)

encoded input dim: 51

Controller network¶

A two-layer SNN: a hidden LIF layer followed by two non-spiking leaky-integrate (LI) readout neurons. The action is the index of the LI neuron with the higher membrane potential.

Under the new Flax-NNX API the network is an nnx.Module; spyx.nn.Sequential chains stateful spiking layers and threads each layer's hidden state through __call__.

In [4]:

Copied!





class Controller(nnx.Module):
    def __init__(self, in_dim, hidden=64, n_actions=2, *, rngs):
        self.core = snn.Sequential(
            nnx.Linear(in_dim, hidden, use_bias=False, rngs=rngs),
            snn.LIF((hidden,), beta=0.8, rngs=rngs),
            nnx.Linear(hidden, n_actions, use_bias=False, rngs=rngs),
            snn.LI((n_actions,), rngs=rngs),
        )

    def __call__(self, x, state):
        return self.core(x, state)

    def initial_state(self, batch_size=1):
        return self.core.initial_state(batch_size)


policy = Controller(input_dim, rngs=nnx.Rngs(0))
graphdef, params_state = nnx.split(policy, nnx.Param)
flat_params, unravel = ravel_pytree(params_state)
print("trainable parameters:", flat_params.size)
class Controller(nnx.Module):
    def __init__(self, in_dim, hidden=64, n_actions=2, *, rngs):
        self.core = snn.Sequential(
            nnx.Linear(in_dim, hidden, use_bias=False, rngs=rngs),
            snn.LIF((hidden,), beta=0.8, rngs=rngs),
            nnx.Linear(hidden, n_actions, use_bias=False, rngs=rngs),
            snn.LI((n_actions,), rngs=rngs),
        )

    def __call__(self, x, state):
        return self.core(x, state)

    def initial_state(self, batch_size=1):
        return self.core.initial_state(batch_size)


policy = Controller(input_dim, rngs=nnx.Rngs(0))
graphdef, params_state = nnx.split(policy, nnx.Param)
flat_params, unravel = ravel_pytree(params_state)
print("trainable parameters:", flat_params.size)

trainable parameters: 3395

Episode rollout¶

evaluate_one runs a single episode under a flat parameter vector. We use nnx.merge(graphdef, unravel(p)) to reconstitute the model inside JIT, then drive it step-by-step against gymnax. evaluate_population vmaps over a (P, D) matrix of candidate parameters.

In [5]:

Copied!





def action_selection(li_voltage):
    # li_voltage shape: (1, n_actions); pick the higher voltage neuron.
    return jnp.argmax(li_voltage[0])


def evaluate_one(flat_p, key, max_steps=500):
    model = nnx.merge(graphdef, unravel(flat_p))
    state = model.initial_state(batch_size=1)

    key_reset, key_step = jax.random.split(key)
    obs0, env_state0 = env.reset(key_reset, env_params)

    def step(carry, key_t):
        obs_t, env_state_t, model_state_t, total_reward, alive = carry
        spikes = adapter(obs_t)[None, :]
        li_v, model_state_next = model(spikes, model_state_t)
        action = action_selection(li_v)
        next_obs, next_env_state, reward, done, _ = env.step(
            key_t, env_state_t, action, env_params
        )
        new_total = total_reward + reward * alive
        new_alive = alive * (1.0 - done.astype(jnp.float32))
        return (next_obs, next_env_state, model_state_next, new_total, new_alive), None

    keys = jax.random.split(key_step, max_steps)
    init = (obs0, env_state0, state, jnp.float32(0.0), jnp.float32(1.0))
    (_, _, _, total_reward, _), _ = jax.lax.scan(step, init, keys)
    return total_reward


def evaluate_population(pop_params, keys):
    return jax.vmap(evaluate_one)(pop_params, keys)


evaluate_population_jit = jax.jit(evaluate_population)
def action_selection(li_voltage):
    # li_voltage shape: (1, n_actions); pick the higher voltage neuron.
    return jnp.argmax(li_voltage[0])


def evaluate_one(flat_p, key, max_steps=500):
    model = nnx.merge(graphdef, unravel(flat_p))
    state = model.initial_state(batch_size=1)

    key_reset, key_step = jax.random.split(key)
    obs0, env_state0 = env.reset(key_reset, env_params)

    def step(carry, key_t):
        obs_t, env_state_t, model_state_t, total_reward, alive = carry
        spikes = adapter(obs_t)[None, :]
        li_v, model_state_next = model(spikes, model_state_t)
        action = action_selection(li_v)
        next_obs, next_env_state, reward, done, _ = env.step(
            key_t, env_state_t, action, env_params
        )
        new_total = total_reward + reward * alive
        new_alive = alive * (1.0 - done.astype(jnp.float32))
        return (next_obs, next_env_state, model_state_next, new_total, new_alive), None

    keys = jax.random.split(key_step, max_steps)
    init = (obs0, env_state0, state, jnp.float32(0.0), jnp.float32(1.0))
    (_, _, _, total_reward, _), _ = jax.lax.scan(step, init, keys)
    return total_reward


def evaluate_population(pop_params, keys):
    return jax.vmap(evaluate_one)(pop_params, keys)


evaluate_population_jit = jax.jit(evaluate_population)

Evolution¶

We use evosax's CMA-ES strategy on the flat parameter vector. The cell below targets the modern (evosax >= 0.2) CMA_ES(population_size=, solution=) constructor and the init / ask / tell methods that all take a key argument. The import block above falls back to the legacy top-level evosax.CMA_ES if the new path is unavailable.

Expected result: with POP_SIZE=64 and GENERATIONS=25, CMA-ES reliably drives the spiking controller to the CartPole-v1 return cap of 500.0 (a fully solved episode), usually within the first ~10–20 generations. If your run plateaus well below 500, try raising GENERATIONS, widening the encoder bins, or adding a hidden layer.

In [6]:

Copied!





POP_SIZE = 64
GENERATIONS = 25

key = jax.random.PRNGKey(42)
key, init_key = jax.random.split(key)

strategy = CMA_ES(population_size=POP_SIZE, solution=jnp.zeros_like(flat_params))
es_params = strategy.default_params
es_state = strategy.init(init_key, mean=flat_params, params=es_params)

best_so_far = -jnp.inf
best_params = flat_params

for gen in trange(GENERATIONS):
    key, ask_key, tell_key, eval_key = jax.random.split(key, 4)
    candidates, es_state = strategy.ask(ask_key, es_state, es_params)
    eval_keys = jax.random.split(eval_key, POP_SIZE)
    rewards = evaluate_population_jit(candidates, eval_keys)
    # evosax minimises by default; flip sign so larger reward wins.
    es_state, _ = strategy.tell(tell_key, candidates, -rewards, es_state, es_params)
    gen_best = jnp.max(rewards)
    if gen_best > best_so_far:
        best_so_far = gen_best
        best_params = candidates[jnp.argmax(rewards)]

print("best return seen:", float(best_so_far))
POP_SIZE = 64
GENERATIONS = 25

key = jax.random.PRNGKey(42)
key, init_key = jax.random.split(key)

strategy = CMA_ES(population_size=POP_SIZE, solution=jnp.zeros_like(flat_params))
es_params = strategy.default_params
es_state = strategy.init(init_key, mean=flat_params, params=es_params)

best_so_far = -jnp.inf
best_params = flat_params

for gen in trange(GENERATIONS):
    key, ask_key, tell_key, eval_key = jax.random.split(key, 4)
    candidates, es_state = strategy.ask(ask_key, es_state, es_params)
    eval_keys = jax.random.split(eval_key, POP_SIZE)
    rewards = evaluate_population_jit(candidates, eval_keys)
    # evosax minimises by default; flip sign so larger reward wins.
    es_state, _ = strategy.tell(tell_key, candidates, -rewards, es_state, es_params)
    gen_best = jnp.max(rewards)
    if gen_best > best_so_far:
        best_so_far = gen_best
        best_params = candidates[jnp.argmax(rewards)]

print("best return seen:", float(best_so_far))

  0%|          | 0/25 [00:00<?, ?it/s]

  4%|▍         | 1/25 [00:01<00:42,  1.77s/it]

  8%|▊         | 2/25 [00:02<00:27,  1.19s/it]

 12%|█▏        | 3/25 [00:03<00:18,  1.16it/s]

 16%|█▌        | 4/25 [00:03<00:14,  1.43it/s]

 20%|██        | 5/25 [00:03<00:12,  1.58it/s]

 24%|██▍       | 6/25 [00:04<00:10,  1.74it/s]

 28%|██▊       | 7/25 [00:04<00:10,  1.78it/s]

 32%|███▏      | 8/25 [00:05<00:09,  1.82it/s]

 36%|███▌      | 9/25 [00:05<00:08,  1.90it/s]

 40%|████      | 10/25 [00:06<00:07,  2.00it/s]

 44%|████▍     | 11/25 [00:06<00:06,  2.10it/s]

 48%|████▊     | 12/25 [00:07<00:05,  2.18it/s]

 52%|█████▏    | 13/25 [00:07<00:05,  2.24it/s]

 56%|█████▌    | 14/25 [00:08<00:04,  2.24it/s]

 60%|██████    | 15/25 [00:08<00:04,  2.31it/s]

 64%|██████▍   | 16/25 [00:08<00:03,  2.28it/s]

 68%|██████▊   | 17/25 [00:09<00:03,  2.27it/s]

 72%|███████▏  | 18/25 [00:09<00:03,  2.16it/s]

 76%|███████▌  | 19/25 [00:10<00:02,  2.20it/s]

 80%|████████  | 20/25 [00:10<00:02,  2.11it/s]

 84%|████████▍ | 21/25 [00:11<00:01,  2.07it/s]

 88%|████████▊ | 22/25 [00:11<00:01,  2.12it/s]

 92%|█████████▏| 23/25 [00:12<00:00,  2.15it/s]

 96%|█████████▌| 24/25 [00:12<00:00,  2.17it/s]

100%|██████████| 25/25 [00:13<00:00,  2.18it/s]

100%|██████████| 25/25 [00:13<00:00,  1.89it/s]

best return seen: 500.0

Roll out the elite¶

Replay the best individual we've seen so far.

In [7]:

Copied!

rollout_reward = evaluate_one(best_params, jax.random.PRNGKey(123))
print("elite episode return:", float(rollout_reward))
rollout_reward = evaluate_one(best_params, jax.random.PRNGKey(123))
print("elite episode return:", float(rollout_reward))

elite episode return: 500.0

Next steps¶

Stack additional LIF layers in Controller.core for a deeper network.
Swap evosax.CMA_ES for OpenES or PGPE to compare strategies.
Move to a richer environment (Brax, Procgen) and let jax.vmap parallelise across both candidates and episodes per candidate.