evosax

Certain population-based algorithms, such as SimpleGA, require an explicit initial population and fitness values during initialization rather than starting from a single template solution.

To initialize these algorithms:

1. Create the algorithm instance (e.g., SimpleGA).
2. Prepare an initial population (e.g., by replicating a template solution) and its corresponding fitness values.
3. Use ga.init(key, population_init, fitness_init, params) to seed the state. The fitness_init argument is used to seed the stored population ranking for the first ask(...) call.

import jax
import jax.numpy as jnp
from evosax.algorithms import SimpleGA

key = jax.random.key(0)
solution = dummy_solution
ga = SimpleGA(population_size=32, solution=solution)
params = ga.default_params

# Replicate a template solution (or provide your own evaluated population).
population_init = jax.tree.map(
    lambda x: jnp.repeat(x[None, ...], ga.population_size, axis=0),
    solution,
)
fitness_init = ...

# Initialize state from the initial population and its fitness.
state = ga.init(key, population_init, fitness_init, params)

Most evolution strategies in evosax follow a standard ask-eval-tell cycle. This pattern allows you to generate candidate solutions, evaluate their fitness, and update the algorithm state in a loop compatible with JAX transformations.

1. Instantiate: Create the algorithm instance (e.g., CMA_ES) with a population_size and a solution template.
2. Initialize: Use es.init(key, solution, params) to create the initial state.
3. Ask: Call es.ask(key_ask, state, params) to generate a population of candidate solutions.
4. Evaluate: Compute the fitness for the generated population.
5. Tell: Call es.tell(key, population, fitness, state, params) to update the algorithm state and receive metrics.

import jax
from evosax.algorithms import CMA_ES

# Instantiate the search strategy
es = CMA_ES(population_size=32, solution=dummy_solution)
params = es.default_params

# Initialize state
key = jax.random.key(0)
state = es.init(key, dummy_solution, params)

# Ask-Eval-Tell loop
for i in range(num_generations):
    key, key_ask, key_eval = jax.random.split(key, 3)

    # Generate a set of candidate solutions to evaluate
    population, state = es.ask(key_ask, state, params)

    # Evaluate the fitness of the population
    fitness = ...

    # Update the evolution strategy
    state, metrics = es.tell(key, population, fitness, state, params)

# Get best solution
state.best_solution, state.best_fitness

The core evolution loop follows the ask, eval, and tell pattern. To implement restarts (e.g., for IPOP-style algorithms), you can monitor a condition (like fitness convergence) and re-initialize the algorithm's state using the current mean solution when the condition is met.

Note: When calling es.tell, fitness is typically passed as a negative value (-fitness) if the algorithm expects minimization, or vice versa depending on the specific implementation requirements.

# 1. Instantiate algorithm
from evosax.algorithms import Open_ES as ES
import optax

es = ES(
    population_size=16,
    solution=solution,
    optimizer=optax.adam(learning_rate=0.01),
    std_schedule=optax.constant_schedule(0.1),
)
params = es.default_params

# 2. Initialize state
state = es.init(subkey, solution, params)

# 3. Run loop with restart logic
def fitness_std_cond(population, fitness, state, params):
    return jnp.std(fitness) < 0.001

for i in range(num_generations):
    # Ask
    population, state = es.ask(key_ask, state, params)
    
    # Eval
    fitness, problem_state, info = problem.eval(key_eval, population, problem_state)
    
    # Tell
    state, metrics = es.tell(key_tell, population, -fitness, state, params)

    # Restart Condition
    if fitness_std_cond(population, fitness, state, params):
        mean = es.get_mean(state)
        state = es.init(subkey, mean, params)

The standard workflow for running an evolution strategy in evosax follows the Ask-Eval-Tell pattern. For SV_CMA_ES, you initialize the algorithm with a population size and a template solution, then use jax.lax.scan to iterate through generations.

In each step:

1. Ask: Use es.ask(key, state, params) to generate a population of candidate solutions.
2. Eval: Use problem.eval(key, population, problem_state) to compute the fitness of those candidates.
3. Tell: Use es.tell(key, population, fitness, state, params) to update the algorithm's state based on the evaluated fitness.

from evosax.algorithms import SV_CMA_ES as ES

# 1. Instantiate the strategy
es = ES(
    population_size=128,
    num_populations=128,
    solution=initial_solution,
)

# 2. Configure parameters
params = es.default_params.replace(alpha=2.0)

# 3. Initialize state
state = es.init(subkey, initial_solutions, params)

# 4. Define the loop step
def step(carry, key):
    state, params, problem_state = carry
    key_ask, key_eval, key_tell = jax.random.split(key, 3)

    # Ask
    population, state = es.ask(key_ask, state, params)
    
    # Eval
    fitness, problem_state, _ = problem.eval(key_eval, population, problem_state)

    # Tell
    state, metrics = es.tell(key_tell, population, fitness, state, params)

    return (state, params, problem_state), metrics

# 5. Run the loop using jax.lax.scan
_, metrics = jax.lax.scan(step, (state, params, problem_state), keys)

You can parallelize evosax algorithms across multiple JAX devices by using jax.sharding.Mesh and NamedSharding.

To achieve efficient parallelization, use two types of sharding:

1. Replicate Sharding: Use NamedSharding(mesh, PartitionSpec()) for parameters (params) and algorithm state (state) so that every device has a copy of the current solution.
2. Parallel Sharding: Use NamedSharding(mesh, PartitionSpec("devices")) for the population. This shards the candidate solutions across devices, allowing each device to evaluate a subset of the population in parallel.

When calling es.ask, specify out_shardings=(parallel_sharding, replicate_sharding) to ensure the returned population is distributed and the state is replicated.

import jax
from jax.sharding import Mesh, NamedSharding, PartitionSpec

# Setup mesh
devices = jax.devices()[:num_devices]
mesh = Mesh(devices, ("devices",))

# Define sharding strategies
replicate_sharding = NamedSharding(mesh, PartitionSpec())
parallel_sharding = NamedSharding(mesh, PartitionSpec("devices"))

# Example: Parallelized Ask
# population is sharded, state is replicated
population, state = jax.jit(
    es.ask,
    out_shardings=(parallel_sharding, replicate_sharding)
)(key_ask, state, params)

The standard evolution loop in evosax follows the ask -> evaluate -> tell pattern. When using PersistentES, you should monitor state.inner_step_counter to reset your inner problem parameters (like the current position xs) when the strategy resets its internal state.

1. Ask: population, state = strategy.ask(key_ask, state, params) generates candidate solutions.
2. Evaluate: Compute fitness for the population (often using jax.vmap for efficiency).
3. Tell: state, metrics = strategy.tell(key_tell, population, fitness, state, params) updates the strategy with the results.

for i in range(5000):
    key, key_ask, key_tell = jax.random.split(key, 3)

    # Reset inner problem if the ES resets its inner step counter
    if state.inner_step_counter == 0:
        xs = jnp.ones((population_size, 2)) * jnp.array([1.0, 1.0])

    # 1. Ask
    population, state = strategy.ask(key_ask, state, params)

    # 2. Evaluate (e.g., using vmap over an unroll function)
    fitness, xs = jax.vmap(unroll, in_axes=(0, 0, None, None, None))(
        xs, population, state.inner_step_counter, params.T, params.K
    )

    # 3. Tell
    state, metrics = strategy.tell(key_tell, population, fitness, state, params)

For algorithms like CMA_ES, evosax provides built-in restart conditions in evosax.restarts.restart_conds. You can use cma_cond and spread_cond to trigger restarts based on the internal state of the evolution strategy.

from evosax.algorithms import CMA_ES as ES
from evosax.restarts.restart_conds import cma_cond, spread_cond

es = ES(population_size=16, solution=solution)
params = es.default_params
state = es.init(subkey, solution, params)

# Inside the loop...
if spread_cond(population, fitness, state, params) | cma_cond(population, fitness, state, params):
    mean = es.get_mean(state)
    state = es.init(subkey, mean, params)

from evosax.algorithms import CMA_ES as ES
from evosax.restarts.restart_conds import cma_cond, spread_cond

es = ES(population_size=16, solution=solution)
params = es.default_params
state = es.init(subkey, solution, params)

# ... loop ...
if spread_cond(population, fitness, state, params) | cma_cond(population, fitness, state, params):
    mean = es.get_mean(state)
    state = es.init(subkey, mean, params)