cartpole_fast

JAX-accelerated cart-pole environment — 4,736x faster than gymnasium on GPU.

Izhikevich SNN evolved with CMA-ES — EASY mode, fitness 14,852 (target: 14,250)

A pure-JAX implementation of the Plank et al. (2024) neuromorphic CartPole benchmark, designed by TENNLab at the University of Tennessee as a standardised evaluation suite for neuromorphic computing systems. The benchmark defines four difficulty levels that systematically test a controller's ability to balance a pole under increasingly constrained observation and action spaces — from full-state feedback down to position-only observations that require temporal integration to infer velocity. cartpole_fast provides the complete benchmark (all four modes, spike encoders, fitness scoring) as a JAX-native environment that can be composed with jit, vmap, and lax.scan for hardware-accelerated evaluation.

Why JAX for Evolutionary Neuromorphic Benchmarks

Evolutionary optimisation of spiking neural networks requires evaluating large populations across many stochastic trials — hundreds of agents, each tested on dozens of random initial conditions. In a typical setup (200 agents x 20 trials), that's 4,000 independent rollouts per generation.

JAX makes this tractable through functional transformations:

• jax.vmap vectorises a single-agent rollout across the entire population and trial batch, turning 4,000 sequential rollouts into a single batched operation with no Python loop overhead.
• jax.jit compiles the full evaluation (environment steps, spike encoding, SNN forward pass, fitness scoring) into a single fused kernel — the entire generation evaluates in one JIT call.
• jax.lax.scan replaces the per-timestep Python loop with a compiled scan, critical for SNN simulations that run 20+ spike timesteps per environment step.
• Pure-functional design means environment state is an immutable pytree with no side effects. This is what makes vmap and scan composition possible — every function takes state in and returns state out.

The result: a full generation of 200 agents x 20 trials x 15,000 steps evaluates in ~9 ms on an L4 GPU, compared to ~44 seconds with gymnasium's vectorized environment. On CPU, JIT compilation alone gives 7-33x over gymnasium.

Features

• All 4 Plank benchmark modes with correct physics, fitness scoring, and activity thresholds
• Pure JAX — fully JIT-compiled, vmappable, and scannable; no Python loops in the hot path
• Neuromorphic encoding — Spike-FF-2/4 and Argyle-3/4 encoders from the Plank paper
• Population evaluation — evaluate 200 agents x 20 trials in a single jax.jit call (~9 ms on L4 GPU)
• GIF rendering for episode visualization

Installation

# Core (CPU — works on any platform)
pip install -e .

# With NVIDIA GPU acceleration
pip install -e .
pip install jax[cuda12]   # overwrites CPU-only jaxlib

# All extras (rendering, examples, dev tools)
pip install -e ".[all]"

Platform	Backend Install	Expected Speedup
CPU (any OS)	included by default	7-33x vs gymnasium
NVIDIA GPU	pip install jax[cuda12]	1,000-5,000x vs gymnasium
Cloud (Modal)	see benchmarks/modal_bench.py	L4 GPU results below

Quick Start

import jax
import cartpole_fast as cp

config = cp.make_config("EASY", max_steps=10000)

def random_policy(env_state, key):
    return jax.random.randint(key, (), 0, config.n_actions)

# Single episode
key = jax.random.PRNGKey(0)
final_state, total_reward = cp.rollout_episode(config, random_policy, key)

# Population evaluation (200 agents x 20 trials)
fitness = cp.evaluate_population(config, make_policy_fn, genotypes, key, n_trials=20)

Performance

Compared against gymnasium SyncVectorEnv (vectorized). Random policy, median of 10 runs.

NVIDIA L4 GPU

Scenario	cartpole_fast	gymnasium	Speedup
Single env (10K steps)	122 ms	98 ms	0.8x
10 envs x 1K steps	9.8 ms	121 ms	12x
100 envs x 1K steps	9.3 ms	928 ms	100x
1,000 envs x 1K steps	9.6 ms	9.7 s	1,017x
10,000 envs x 1K steps	10.6 ms	~97 s*	~9,200x
100,000 envs x 1K steps	23.9 ms	~16 min*	~40,800x
Pop eval (4K rollouts x 1K steps)	9 ms	43.6 s	4,736x

Scaling (L4 GPU)

Time (ms, log scale)                    gymnasium scales linearly ↗
1,000,000 |                                                    ·
          |                                               ·
  100,000 |                                          ·
          |                                     ·
   10,000 |                              · ·
          |                         ·
    1,000 |                    ·
          |               ·
      100 |          ·  ·
          |     ·
       10 |──●────●────●────●────●────●     cartpole_fast stays flat
          |
        1 ┼─────┼─────┼─────┼─────┼─────┼
          1    10   100   1K   10K  100K
                  Parallel environments

cartpole_fast processes 1 to 10,000 environments in ~10 ms flat. Gymnasium scales linearly — every 10x more environments costs 10x more time. On GPU, parallelism is essentially free until you saturate memory.

CPU (Modal, no GPU)

Scenario	cartpole_fast	gymnasium	Speedup
Single env (10K steps)	18 ms	125 ms	6.8x
100 envs x 1K steps	61 ms	1.4 s	22x
1,000 envs x 1K steps	438 ms	13.5 s	31x
Pop eval (4K rollouts x 1K steps)	1.6 s	52 s	33x

_{*Extrapolated linearly from N=1,000 (gymnasium scales linearly with env count).}

Key takeaways:

• 0.002 secs/1M steps on L4, competitive with gymnax (0.05 secs/1M on A100)
• Single-env GPU is slower (0.8x) due to kernel launch overhead — the value is batched parallelism
• Even on CPU, JIT compilation alone gives 7-33x over gymnasium

Run benchmarks yourself:

# Local (auto-detects CPU or GPU)
uv run --with gymnasium benchmarks/bench_throughput.py

# Modal (L4 GPU + CPU in parallel)
uv run --with modal,gymnasium modal run benchmarks/modal_bench.py

Benchmark Modes (Plank et al. 2024)

The Plank benchmark defines four difficulty levels that progressively constrain what the controller can observe and do. All modes share the same underlying physics (identical to OpenAI Gym), but with longer episodes (15,000 steps vs 200/500), wider initial positions (x ∈ [-1.2, 1.2] vs [-0.05, 0.05]), and difficulty-specific observation/action spaces.

Mode	Obs	Actions	Target Fitness	What It Tests
EASY	x, x_dot, theta, theta_dot	left, right	14,250	Basic balancing with full state
MEDIUM	x, x_dot, theta, theta_dot	left, right, nothing	12,000	Energy efficiency — must be idle ≥75% of steps
HARD	x, theta	left, right, nothing	9,000	Temporal integration — must infer velocity from position history
HARDEST	x, theta	left, right	6,000	Both constraints: partial obs + no idle action

EASY

Full 4D state observation, 2 actions. The baseline — any competent controller should approach the 14,250 target. This mode is useful for verifying that your SNN architecture and encoding pipeline work before moving to harder conditions.

MEDIUM

Same observations as EASY, but adds a third "do nothing" action and an activity threshold: fitness is only awarded at the full step count if the agent uses "do nothing" at least 75% of the time. Otherwise, fitness reduces to do_nothing_count / 0.75. This tests whether the controller can learn energy-efficient behaviour — intervening only when necessary, a property that matters for neuromorphic hardware where spike activity directly corresponds to energy consumption.

HARD

Observations are reduced to position only — x and theta, with velocities hidden. The agent must infer velocity from the temporal pattern of observations across successive steps. This is where spiking neural networks have a natural advantage: membrane dynamics and spike traces act as implicit memory, allowing the network to integrate temporal information without explicit recurrence. 3 actions (left, right, nothing).

HARDEST

Position-only observations with only 2 actions (left, right) — no "do nothing" option. The agent must both infer dynamics from partial observations and act at every timestep. Target fitness of 6,000 is deliberately low; this mode is designed to be a challenge even for well-tuned neuromorphic systems.

Differences from Standard CartPole

The Plank benchmark modifies the classic OpenAI Gym CartPole in several ways to create a more demanding and discriminating evaluation:

Feature	Standard CartPole (Gym)	Plank Benchmark (cartpole_fast)
Max steps	200 or 500	15,000
Initial cart position	x ∈ [-0.05, 0.05]	x ∈ [-1.2, 1.2]
Initial pole angle	θ ∈ [-0.05, 0.05]	θ ∈ [-0.10475, 0.10475]
Difficulty modes	1	4 (EASY / MEDIUM / HARD / HARDEST)
Observation masking	No	Yes — HARD/HARDEST hide velocities
3-action support	No	Yes — MEDIUM/HARD add "do nothing"
Activity threshold	No	Yes — MEDIUM requires ≥75% inactivity
Fitness scoring	Episode return	Mode-dependent (step count or activity-penalised)

The wider initial conditions and longer episodes make the task significantly harder than standard CartPole — a random policy scores ~9 steps (vs ~20 in Gym), and reaching the target fitness requires a controller that can recover from large initial displacements.

Neuromorphic Encoding

Encoder	Neurons/Dim	Recommended For
Spike-FF-2	2	EASY, MEDIUM, HARD
Spike-FF-4	4	General purpose
Argyle-3	3	Alternative to FF-2
Argyle-4	4	HARDEST

Pipeline: normalize obs → compute rates → Bernoulli spike sampling over T timesteps.

import cartpole_fast as cp

# Encode: obs → [T, encoded_dim] binary spikes
obs_max = cp.OBS_MAX[cp.Difficulty.EASY]
spikes = cp.encode(obs, cp.EncoderType.SPIKE_FF_4, T=20, max_spikes=20.0, obs_max=obs_max, key=key)

# Decode: output spike counts → action
action = cp.decode_vote(output_spike_counts, T=20)

Custom encoders work too — just produce [T, N] binary spike trains and use any policy_fn(env_state, key) → action.

Rendering

final_state, reward, trajectory = cp.rollout_episode_with_trajectory(config, policy_fn, key)
render_cfg = cp.RenderConfig(width=600, height=400, fps=50, frame_skip=2)
cp.render_episode_gif(trajectory, config, "episode.gif", render_cfg)

Examples

See examples/ for complete, runnable scripts:

• random_policy.py — Minimal demo: rollout a random agent and render a GIF.

  uv run --with Pillow python examples/random_policy.py

• evolve_snn.py — Full neuromorphic baseline: Izhikevich SNN evolved with CMA-ES on EASY mode. Self-contained — all SNN components are inlined, no external dependencies beyond evosax.

  uv run --with evosax --with Pillow python examples/evolve_snn.py

API Reference

Core

Symbol	Description
Difficulty	Enum: EASY, MEDIUM, HARD, HARDEST
CartPoleState	NamedTuple: state, obs, done, reward, step_count, do_nothing_count
CartPoleConfig	Frozen dataclass with all physics params
make_config(difficulty, max_steps, **overrides)	Create config for a benchmark mode
reset(config, key)	→ CartPoleState
step(config, state, action)	→ CartPoleState
compute_fitness(config, state)	→ scalar
TARGET_FITNESS	Dict of target fitness per difficulty
OBS_MAX	Dict of normalization bounds per difficulty

Neuromorphic

Symbol	Description
EncoderType	Enum: SPIKE_FF_2, SPIKE_FF_4, ARGYLE_3, ARGYLE_4
encode(obs, encoder_type, T, max_spikes, obs_max, key)	→ [T, encoded_dim]
encode_spike_ff4_raw(obs, T, max_spikes, obs_max, key)	→ [T, obs_dim*4] (bin-then-normalize order)
encoded_dim(obs_dim, encoder_type)	→ int
decode_vote(output_spike_counts, T)	→ action
decode_rate_difference(output_spike_counts, T)	→ action

Utilities

Symbol	Description
rollout_episode(config, policy_fn, key)	→ (final_state, total_reward)
rollout_episode_with_trajectory(config, policy_fn, key)	→ (state, reward, trajectory)
evaluate_trials(config, policy_fn, key, n_trials)	→ mean_fitness
evaluate_population(config, make_policy_fn, genotypes, key, n_trials)	→ fitness[pop]
RenderConfig	Dataclass: width, height, fps, frame_skip
render_frame(state_4d, config, render_config)	→ PIL Image
render_episode(trajectory, config, render_config)	→ [PIL Image]
render_episode_gif(trajectory, config, path, render_config)	→ path
save_gif(frames, path, render_config)	→ path
render_top_k(trajectories, fitnesses, config, render_config, k, generation)	→ [paths]

Citation

cartpole_fast is based on an initial implementation in MLX developed for the DevNCA paper, and grew out of the requirements for larger-scale extensions of that work. It is used in ongoing research and is released separately as a standalone tool for the community.

If you use cartpole_fast in your research, please cite:

@inproceedings{gaskin2025devnca,
  title     = {{DevNCA}: Co-Evolving Developmental Patterns and Plasticity Rules for Self-Organising Spiking Neural Networks},
  author    = {Gaskin, Benjamin},
  booktitle = {Artificial Life Conference Proceedings 37},
  volume    = {2025},
  number    = {1},
  pages     = {12},
  year      = {2025},
  publisher = {MIT Press},
  doi       = {10.1162/ISAL.a.840}
}

References

• Plank, J., Rizzo, J., White, A., & Schuman, C. (2024). "The Cart-Pole Application as a Benchmark for Neuromorphic Computing." TENNLab, University of Tennessee, Knoxville.
• Gaskin, B. (2025). "DevNCA: Co-Evolving Developmental Patterns and Plasticity Rules for Self-Organising Spiking Neural Networks." ALife 2025. MIT Press.
• Izhikevich, E. M. (2003). "Simple Model of Spiking Neurons." IEEE Transactions on Neural Networks, 14(6), 1569-1572.