Leniabreeder 
¶
Installation¶
You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:
In [1]:
Copied!
%pip install -U "jax[cuda]"
%pip install -U "jax[cuda]"
/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip
Note: you may need to restart the kernel to use updated packages.
Then, install CAX from PyPi:
In [2]:
Copied!
%pip install -U "cax[examples]"
%pip install -U "cax[examples]"
/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip
Note: you may need to restart the kernel to use updated packages.
Import¶
In [3]:
Copied!
import time
from functools import partial
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import mediapy
import numpy as np
import optax
from flax import nnx
from cax.cs.lenia import (
Lenia,
LeniaGrowthParams,
LeniaKernelParams,
LeniaRuleParams,
center_state,
exponential_growth_fn,
gaussian_kernel_fn,
metrics_fn,
)
from cax.nn.vae import VAE, binary_cross_entropy_with_logits, kl_divergence
import time
from functools import partial
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import mediapy
import numpy as np
import optax
from flax import nnx
from cax.cs.lenia import (
Lenia,
LeniaGrowthParams,
LeniaKernelParams,
LeniaRuleParams,
center_state,
exponential_growth_fn,
gaussian_kernel_fn,
metrics_fn,
)
from cax.nn.vae import VAE, binary_cross_entropy_with_logits, kl_divergence
Configuration¶
In [4]:
Copied!
seed = 0
# Lenia parameters
spatial_dims = (64, 64)
channel_size = 3
R = 12
T = 1 / 3
state_scale = 2
num_steps = 128
# Parameter ranges
num_self_kernels = 3
num_cross_kernels = 1
kernel_rank = 3
kernel_r_range = (0.2, 1.0)
kernel_beta_range = (0.0, 1.0)
growth_mean_range = (0.05, 0.5)
growth_std_range = (0.005, 0.1)
state_init_range = (0.0, 1.0)
# QD parameters
population_size = 1024
batch_size = 256
num_generations = 256
k = 3
num_obs_steps = 32
mutation_std = 0.01
weight_concentration = 100.0
min_mass = 0.5
min_concentration = 0.5
# VAE parameters
latent_size = 2
vae_features = (3, 16, 32)
vae_learning_rate = 1e-3
vae_num_epochs = 64
vae_batch_size = 32
vae_retrain_interval = 32
key = jax.random.key(seed)
seed = 0
# Lenia parameters
spatial_dims = (64, 64)
channel_size = 3
R = 12
T = 1 / 3
state_scale = 2
num_steps = 128
# Parameter ranges
num_self_kernels = 3
num_cross_kernels = 1
kernel_rank = 3
kernel_r_range = (0.2, 1.0)
kernel_beta_range = (0.0, 1.0)
growth_mean_range = (0.05, 0.5)
growth_std_range = (0.005, 0.1)
state_init_range = (0.0, 1.0)
# QD parameters
population_size = 1024
batch_size = 256
num_generations = 256
k = 3
num_obs_steps = 32
mutation_std = 0.01
weight_concentration = 100.0
min_mass = 0.5
min_concentration = 0.5
# VAE parameters
latent_size = 2
vae_features = (3, 16, 32)
vae_learning_rate = 1e-3
vae_num_epochs = 64
vae_batch_size = 32
vae_retrain_interval = 32
key = jax.random.key(seed)
Lenia Genotype Sampling¶
In [5]:
Copied!
@nnx.dataclass
class Genotype(nnx.Pytree):
"""Genotype for Lenia evolution."""
rule_params: LeniaRuleParams = nnx.data()
state_init: jax.Array = nnx.data()
def sample_kernel_params(key: jax.Array):
"""Sample kernel parameters."""
key_r, key_beta, key_rank, key_max_idx = jax.random.split(key, 4)
# Spatial extent
r = jax.random.uniform(key_r, minval=kernel_r_range[0], maxval=kernel_r_range[1])
# Rank: uniform over {1, ..., kernel_rank}
rank = jax.random.randint(key_rank, (), minval=1, maxval=kernel_rank + 1)
# Pick which active ring will have the dominant peak height
max_idx = jax.random.randint(key_max_idx, (), minval=0, maxval=rank)
# Sample peak heights uniformly from kernel_beta_range
beta = jax.random.uniform(
key_beta,
(kernel_rank,),
minval=kernel_beta_range[0],
maxval=kernel_beta_range[1],
)
# Set the chosen peak to kernel_beta_range[1]
beta = beta.at[max_idx].set(kernel_beta_range[1])
# Mask: active rings keep their peak height, inactive become nan
active_mask = jnp.arange(kernel_rank) < rank
beta = jnp.where(active_mask, beta, jnp.nan)
return LeniaKernelParams(r=r, beta=beta)
def sample_growth_params(key: jax.Array):
"""Sample growth parameters."""
key_mean, key_std = jax.random.split(key)
mean = jax.random.uniform(
key_mean, minval=growth_mean_range[0], maxval=growth_mean_range[1]
)
std = jax.random.uniform(
key_std, minval=growth_std_range[0], maxval=growth_std_range[1]
)
return LeniaGrowthParams(mean=mean, std=std)
def sample_rule_params(key: jax.Array):
"""Sample rule parameters: structured channel connectivity, Dirichlet weights."""
key_weight, key_kernel_params, key_growth_params = jax.random.split(key, 3)
# Self-kernels: each channel maps to itself
self_sources = jnp.repeat(jnp.arange(channel_size), num_self_kernels)
self_targets = self_sources
# Cross-kernels: all ordered pairs (i, j) where i != j
cross_sources = jnp.array(
[i for i in range(channel_size) for j in range(channel_size) if i != j]
)
cross_targets = jnp.array(
[j for i in range(channel_size) for j in range(channel_size) if i != j]
)
cross_sources = jnp.repeat(cross_sources, num_cross_kernels)
cross_targets = jnp.repeat(cross_targets, num_cross_kernels)
# Build structured channel connectivity (following Chan 2020)
channel_source = jnp.concatenate([self_sources, cross_sources])
channel_target = jnp.concatenate([self_targets, cross_targets])
num_rules = channel_source.shape[0]
# Sample weights jointly via Dirichlet (sum to 1 by construction)
weight = jax.random.dirichlet(key_weight, alpha=jnp.ones(num_rules))
# Sample kernel and growth parameters
keys_kernel = jax.random.split(key_kernel_params, num_rules)
kernel_params = jax.vmap(sample_kernel_params)(keys_kernel)
keys_growth = jax.random.split(key_growth_params, num_rules)
growth_params = jax.vmap(sample_growth_params)(keys_growth)
return LeniaRuleParams(
channel_source=channel_source,
channel_target=channel_target,
weight=weight,
kernel_params=kernel_params,
growth_params=growth_params,
)
def sample_state_init(key: jax.Array):
"""Sample a random initial state with a centered blob."""
blob_size = 2 * R
blob = jax.random.uniform(key, (blob_size, blob_size, channel_size))
pad_height = (spatial_dims[0] - blob_size) // 2
pad_width = (spatial_dims[1] - blob_size) // 2
return jnp.pad(blob, ((pad_height, pad_height), (pad_width, pad_width), (0, 0)))
def sample_genotype(key: jax.Array):
"""Sample a random genotype."""
key_rule, key_state = jax.random.split(key)
return Genotype(
rule_params=sample_rule_params(key_rule),
state_init=sample_state_init(key_state),
)
@nnx.dataclass
class Genotype(nnx.Pytree):
"""Genotype for Lenia evolution."""
rule_params: LeniaRuleParams = nnx.data()
state_init: jax.Array = nnx.data()
def sample_kernel_params(key: jax.Array):
"""Sample kernel parameters."""
key_r, key_beta, key_rank, key_max_idx = jax.random.split(key, 4)
# Spatial extent
r = jax.random.uniform(key_r, minval=kernel_r_range[0], maxval=kernel_r_range[1])
# Rank: uniform over {1, ..., kernel_rank}
rank = jax.random.randint(key_rank, (), minval=1, maxval=kernel_rank + 1)
# Pick which active ring will have the dominant peak height
max_idx = jax.random.randint(key_max_idx, (), minval=0, maxval=rank)
# Sample peak heights uniformly from kernel_beta_range
beta = jax.random.uniform(
key_beta,
(kernel_rank,),
minval=kernel_beta_range[0],
maxval=kernel_beta_range[1],
)
# Set the chosen peak to kernel_beta_range[1]
beta = beta.at[max_idx].set(kernel_beta_range[1])
# Mask: active rings keep their peak height, inactive become nan
active_mask = jnp.arange(kernel_rank) < rank
beta = jnp.where(active_mask, beta, jnp.nan)
return LeniaKernelParams(r=r, beta=beta)
def sample_growth_params(key: jax.Array):
"""Sample growth parameters."""
key_mean, key_std = jax.random.split(key)
mean = jax.random.uniform(
key_mean, minval=growth_mean_range[0], maxval=growth_mean_range[1]
)
std = jax.random.uniform(
key_std, minval=growth_std_range[0], maxval=growth_std_range[1]
)
return LeniaGrowthParams(mean=mean, std=std)
def sample_rule_params(key: jax.Array):
"""Sample rule parameters: structured channel connectivity, Dirichlet weights."""
key_weight, key_kernel_params, key_growth_params = jax.random.split(key, 3)
# Self-kernels: each channel maps to itself
self_sources = jnp.repeat(jnp.arange(channel_size), num_self_kernels)
self_targets = self_sources
# Cross-kernels: all ordered pairs (i, j) where i != j
cross_sources = jnp.array(
[i for i in range(channel_size) for j in range(channel_size) if i != j]
)
cross_targets = jnp.array(
[j for i in range(channel_size) for j in range(channel_size) if i != j]
)
cross_sources = jnp.repeat(cross_sources, num_cross_kernels)
cross_targets = jnp.repeat(cross_targets, num_cross_kernels)
# Build structured channel connectivity (following Chan 2020)
channel_source = jnp.concatenate([self_sources, cross_sources])
channel_target = jnp.concatenate([self_targets, cross_targets])
num_rules = channel_source.shape[0]
# Sample weights jointly via Dirichlet (sum to 1 by construction)
weight = jax.random.dirichlet(key_weight, alpha=jnp.ones(num_rules))
# Sample kernel and growth parameters
keys_kernel = jax.random.split(key_kernel_params, num_rules)
kernel_params = jax.vmap(sample_kernel_params)(keys_kernel)
keys_growth = jax.random.split(key_growth_params, num_rules)
growth_params = jax.vmap(sample_growth_params)(keys_growth)
return LeniaRuleParams(
channel_source=channel_source,
channel_target=channel_target,
weight=weight,
kernel_params=kernel_params,
growth_params=growth_params,
)
def sample_state_init(key: jax.Array):
"""Sample a random initial state with a centered blob."""
blob_size = 2 * R
blob = jax.random.uniform(key, (blob_size, blob_size, channel_size))
pad_height = (spatial_dims[0] - blob_size) // 2
pad_width = (spatial_dims[1] - blob_size) // 2
return jnp.pad(blob, ((pad_height, pad_height), (pad_width, pad_width), (0, 0)))
def sample_genotype(key: jax.Array):
"""Sample a random genotype."""
key_rule, key_state = jax.random.split(key)
return Genotype(
rule_params=sample_rule_params(key_rule),
state_init=sample_state_init(key_state),
)
Lenia Genotype Mutation¶
In [6]:
Copied!
def reflect(x, *, lower=None, upper=None):
"""Reflect x into bounds using reflective boundary conditions.
Supports one-sided reflection (only lower or only upper) and two-sided
reflection (both lower and upper).
Args:
x: Array to reflect.
lower: Lower bound. If None, no lower reflection is applied.
upper: Upper bound. If None, no upper reflection is applied.
Returns:
Array reflected into the specified bounds.
"""
# One-sided or no reflection
if lower is None and upper is None:
return x
if lower is None:
return jnp.where(x > upper, 2 * upper - x, x)
if upper is None:
return jnp.where(x < lower, 2 * lower - x, x)
# Two-sided: periodic folding into [lower, upper]
span = upper - lower
x = x - lower
x = jnp.abs(x)
n = jnp.floor(x / span).astype(int)
x = x - n * span
x = jnp.where(n % 2 == 0, x, span - x)
return x + lower
def mutate_bounded(key, x, *, lower, upper, mutation_std):
"""Gaussian mutation with scale relative to range, reflected into bounds.
Args:
key: PRNG key.
x: Array to mutate.
lower: Lower bound.
upper: Upper bound.
mutation_std: Standard deviation relative to the range (upper - lower).
Returns:
Mutated array reflected into [lower, upper].
"""
noise = mutation_std * (upper - lower) * jax.random.normal(key, x.shape)
return reflect(x + noise, lower=lower, upper=upper)
def mutate_genotype(key, genotype, *, mutation_std, weight_concentration):
"""Mutate a genotype with proper per-parameter constraints.
Args:
key: PRNG key.
genotype: Genotype to mutate.
mutation_std: Relative standard deviation for bounded Gaussian mutations.
weight_concentration: Concentration parameter for Dirichlet weight
perturbation.
Returns:
Mutated Genotype.
"""
key_weight, key_r, key_beta, key_mean, key_std, key_state = jax.random.split(key, 6)
# Weights: Dirichlet perturbation (stays on simplex)
weight = jax.random.dirichlet(
key_weight, alpha=weight_concentration * genotype.rule_params.weight
)
# Kernel radius: bounded Gaussian
r = mutate_bounded(
key_r,
genotype.rule_params.kernel_params.r,
lower=kernel_r_range[0],
upper=kernel_r_range[1],
mutation_std=mutation_std,
)
# Beta: Gaussian + reflect at 0 + normalize so max active = 1.0
is_inactive = jnp.isnan(genotype.rule_params.kernel_params.beta)
beta = genotype.rule_params.kernel_params.beta + mutation_std * jax.random.normal(
key_beta, genotype.rule_params.kernel_params.beta.shape
)
beta = jnp.where(is_inactive, jnp.nan, reflect(beta, lower=0.0))
beta = beta / jnp.nanmax(beta)
# Growth mean: bounded Gaussian
mean = mutate_bounded(
key_mean,
genotype.rule_params.growth_params.mean,
lower=growth_mean_range[0],
upper=growth_mean_range[1],
mutation_std=mutation_std,
)
# Growth std: bounded Gaussian
std = mutate_bounded(
key_std,
genotype.rule_params.growth_params.std,
lower=growth_std_range[0],
upper=growth_std_range[1],
mutation_std=mutation_std,
)
# State init: bounded Gaussian
state_init = mutate_bounded(
key_state,
genotype.state_init,
lower=0.0,
upper=1.0,
mutation_std=mutation_std,
)
return Genotype(
rule_params=LeniaRuleParams(
channel_source=genotype.rule_params.channel_source,
channel_target=genotype.rule_params.channel_target,
weight=weight,
kernel_params=LeniaKernelParams(r=r, beta=beta),
growth_params=LeniaGrowthParams(mean=mean, std=std),
),
state_init=state_init,
)
def reflect(x, *, lower=None, upper=None):
"""Reflect x into bounds using reflective boundary conditions.
Supports one-sided reflection (only lower or only upper) and two-sided
reflection (both lower and upper).
Args:
x: Array to reflect.
lower: Lower bound. If None, no lower reflection is applied.
upper: Upper bound. If None, no upper reflection is applied.
Returns:
Array reflected into the specified bounds.
"""
# One-sided or no reflection
if lower is None and upper is None:
return x
if lower is None:
return jnp.where(x > upper, 2 * upper - x, x)
if upper is None:
return jnp.where(x < lower, 2 * lower - x, x)
# Two-sided: periodic folding into [lower, upper]
span = upper - lower
x = x - lower
x = jnp.abs(x)
n = jnp.floor(x / span).astype(int)
x = x - n * span
x = jnp.where(n % 2 == 0, x, span - x)
return x + lower
def mutate_bounded(key, x, *, lower, upper, mutation_std):
"""Gaussian mutation with scale relative to range, reflected into bounds.
Args:
key: PRNG key.
x: Array to mutate.
lower: Lower bound.
upper: Upper bound.
mutation_std: Standard deviation relative to the range (upper - lower).
Returns:
Mutated array reflected into [lower, upper].
"""
noise = mutation_std * (upper - lower) * jax.random.normal(key, x.shape)
return reflect(x + noise, lower=lower, upper=upper)
def mutate_genotype(key, genotype, *, mutation_std, weight_concentration):
"""Mutate a genotype with proper per-parameter constraints.
Args:
key: PRNG key.
genotype: Genotype to mutate.
mutation_std: Relative standard deviation for bounded Gaussian mutations.
weight_concentration: Concentration parameter for Dirichlet weight
perturbation.
Returns:
Mutated Genotype.
"""
key_weight, key_r, key_beta, key_mean, key_std, key_state = jax.random.split(key, 6)
# Weights: Dirichlet perturbation (stays on simplex)
weight = jax.random.dirichlet(
key_weight, alpha=weight_concentration * genotype.rule_params.weight
)
# Kernel radius: bounded Gaussian
r = mutate_bounded(
key_r,
genotype.rule_params.kernel_params.r,
lower=kernel_r_range[0],
upper=kernel_r_range[1],
mutation_std=mutation_std,
)
# Beta: Gaussian + reflect at 0 + normalize so max active = 1.0
is_inactive = jnp.isnan(genotype.rule_params.kernel_params.beta)
beta = genotype.rule_params.kernel_params.beta + mutation_std * jax.random.normal(
key_beta, genotype.rule_params.kernel_params.beta.shape
)
beta = jnp.where(is_inactive, jnp.nan, reflect(beta, lower=0.0))
beta = beta / jnp.nanmax(beta)
# Growth mean: bounded Gaussian
mean = mutate_bounded(
key_mean,
genotype.rule_params.growth_params.mean,
lower=growth_mean_range[0],
upper=growth_mean_range[1],
mutation_std=mutation_std,
)
# Growth std: bounded Gaussian
std = mutate_bounded(
key_std,
genotype.rule_params.growth_params.std,
lower=growth_std_range[0],
upper=growth_std_range[1],
mutation_std=mutation_std,
)
# State init: bounded Gaussian
state_init = mutate_bounded(
key_state,
genotype.state_init,
lower=0.0,
upper=1.0,
mutation_std=mutation_std,
)
return Genotype(
rule_params=LeniaRuleParams(
channel_source=genotype.rule_params.channel_source,
channel_target=genotype.rule_params.channel_target,
weight=weight,
kernel_params=LeniaKernelParams(r=r, beta=beta),
growth_params=LeniaGrowthParams(mean=mean, std=std),
),
state_init=state_init,
)
Lenia Genotype Evaluation¶
In [7]:
Copied!
def fitness_fn(metrics):
"""Compute fitness from simulation metrics. Lower is better.
Args:
metrics: Dictionary of time-series metrics from the simulation.
Returns:
Scalar fitness value (inf if degenerate).
"""
is_degenerate = jnp.any(metrics["mass"] < min_mass) | jnp.any(
metrics["concentration"] < min_concentration
)
mass = metrics["mass"] # noqa: F841
# Center of mass displacements (handling torus wrapping)
center_of_mass = metrics["center_of_mass"]
displacements = jnp.diff(center_of_mass, axis=0)
world_size = jnp.array([dim / R for dim in spatial_dims])
displacements = displacements - jnp.round(displacements / world_size) * world_size
# Linear velocity: instantaneous speed at each step
linear_velocity = jnp.linalg.norm(displacements, axis=-1) * T
# Angle: heading direction at each step (in [-1, 1], units of pi)
angle = jnp.arctan2(displacements[:, 1], displacements[:, 0]) / jnp.pi
# Angular velocity: rate of change of heading
angle_diff = jnp.diff(angle)
angle_diff = (angle_diff + 3) % 2 - 1
angle_diff = jnp.where(linear_velocity[1:] > 0.01, angle_diff, 0.0)
angular_velocity = angle_diff * T # noqa: F841
mean_speed = jnp.mean(linear_velocity)
return jnp.where(is_degenerate, jnp.inf, -mean_speed)
def descriptor_fn(vae, observations):
"""Compute behavioral descriptor from multiple observations via VAE.
Args:
vae: VAE model.
observations: RGB observation array (num_obs_steps, H, W, 3) uint8.
Returns:
Latent descriptor vector (averaged over steps).
"""
means, _ = jax.vmap(vae.encoder)(observations.astype(jnp.float32) / 255.0)
return jnp.mean(means, axis=0)
def evaluate_fn(vae, genotype):
"""Evaluate a single Lenia genotype.
Args:
vae: VAE model used to compute latent descriptor.
genotype: Genotype to evaluate.
Returns:
Tuple of (fitness, descriptor, info).
"""
cs = Lenia(
spatial_dims=spatial_dims,
channel_size=channel_size,
R=R,
T=T,
state_scale=state_scale,
kernel_fn=gaussian_kernel_fn,
growth_fn=exponential_growth_fn,
rule_params=genotype.rule_params,
)
_, states = cs(genotype.state_init, num_steps=num_steps, return_states=True)
metrics = jax.vmap(partial(metrics_fn, R=R))(states)
# Observations
last_states = states[-num_obs_steps:]
centered_states = jax.vmap(partial(center_state, R=R))(last_states)
observations = jax.vmap(cs.render)(centered_states)
# Fitness
fitness = fitness_fn(metrics)
# Descriptor
descriptor = descriptor_fn(vae, observations)
info = {
"observations": observations,
"metrics": metrics,
}
return fitness, descriptor, info
def evaluate(vae, genotypes):
"""Evaluate a batch of genotypes.
Args:
vae: VAE model used to compute latent descriptors.
genotypes: Batch of genotypes to evaluate.
Returns:
Tuple of (fitness, descriptor, info) with batch dimensions.
"""
return jax.vmap(evaluate_fn, in_axes=(None, 0))(vae, genotypes)
def evaluate_in_chunks(vae, genotypes, chunk_size):
"""Evaluate genotypes a chunk at a time, so not all trajectories coexist in memory.
Args:
vae: VAE model used to compute latent descriptors.
genotypes: Batch of genotypes to evaluate.
chunk_size: Genotypes evaluated at once.
Returns:
Tuple of (fitness, descriptor, info) with batch dimensions.
"""
num = jax.tree.leaves(genotypes)[0].shape[0]
chunks = []
for i in range(0, num, chunk_size):
chunk = jax.tree.map(lambda leaf, i=i: leaf[i : i + chunk_size], genotypes)
chunks.append(evaluate(vae, chunk))
return jax.tree.map(lambda *leaves: jnp.concatenate(leaves), *chunks)
def fitness_fn(metrics):
"""Compute fitness from simulation metrics. Lower is better.
Args:
metrics: Dictionary of time-series metrics from the simulation.
Returns:
Scalar fitness value (inf if degenerate).
"""
is_degenerate = jnp.any(metrics["mass"] < min_mass) | jnp.any(
metrics["concentration"] < min_concentration
)
mass = metrics["mass"] # noqa: F841
# Center of mass displacements (handling torus wrapping)
center_of_mass = metrics["center_of_mass"]
displacements = jnp.diff(center_of_mass, axis=0)
world_size = jnp.array([dim / R for dim in spatial_dims])
displacements = displacements - jnp.round(displacements / world_size) * world_size
# Linear velocity: instantaneous speed at each step
linear_velocity = jnp.linalg.norm(displacements, axis=-1) * T
# Angle: heading direction at each step (in [-1, 1], units of pi)
angle = jnp.arctan2(displacements[:, 1], displacements[:, 0]) / jnp.pi
# Angular velocity: rate of change of heading
angle_diff = jnp.diff(angle)
angle_diff = (angle_diff + 3) % 2 - 1
angle_diff = jnp.where(linear_velocity[1:] > 0.01, angle_diff, 0.0)
angular_velocity = angle_diff * T # noqa: F841
mean_speed = jnp.mean(linear_velocity)
return jnp.where(is_degenerate, jnp.inf, -mean_speed)
def descriptor_fn(vae, observations):
"""Compute behavioral descriptor from multiple observations via VAE.
Args:
vae: VAE model.
observations: RGB observation array (num_obs_steps, H, W, 3) uint8.
Returns:
Latent descriptor vector (averaged over steps).
"""
means, _ = jax.vmap(vae.encoder)(observations.astype(jnp.float32) / 255.0)
return jnp.mean(means, axis=0)
def evaluate_fn(vae, genotype):
"""Evaluate a single Lenia genotype.
Args:
vae: VAE model used to compute latent descriptor.
genotype: Genotype to evaluate.
Returns:
Tuple of (fitness, descriptor, info).
"""
cs = Lenia(
spatial_dims=spatial_dims,
channel_size=channel_size,
R=R,
T=T,
state_scale=state_scale,
kernel_fn=gaussian_kernel_fn,
growth_fn=exponential_growth_fn,
rule_params=genotype.rule_params,
)
_, states = cs(genotype.state_init, num_steps=num_steps, return_states=True)
metrics = jax.vmap(partial(metrics_fn, R=R))(states)
# Observations
last_states = states[-num_obs_steps:]
centered_states = jax.vmap(partial(center_state, R=R))(last_states)
observations = jax.vmap(cs.render)(centered_states)
# Fitness
fitness = fitness_fn(metrics)
# Descriptor
descriptor = descriptor_fn(vae, observations)
info = {
"observations": observations,
"metrics": metrics,
}
return fitness, descriptor, info
def evaluate(vae, genotypes):
"""Evaluate a batch of genotypes.
Args:
vae: VAE model used to compute latent descriptors.
genotypes: Batch of genotypes to evaluate.
Returns:
Tuple of (fitness, descriptor, info) with batch dimensions.
"""
return jax.vmap(evaluate_fn, in_axes=(None, 0))(vae, genotypes)
def evaluate_in_chunks(vae, genotypes, chunk_size):
"""Evaluate genotypes a chunk at a time, so not all trajectories coexist in memory.
Args:
vae: VAE model used to compute latent descriptors.
genotypes: Batch of genotypes to evaluate.
chunk_size: Genotypes evaluated at once.
Returns:
Tuple of (fitness, descriptor, info) with batch dimensions.
"""
num = jax.tree.leaves(genotypes)[0].shape[0]
chunks = []
for i in range(0, num, chunk_size):
chunk = jax.tree.map(lambda leaf, i=i: leaf[i : i + chunk_size], genotypes)
chunks.append(evaluate(vae, chunk))
return jax.tree.map(lambda *leaves: jnp.concatenate(leaves), *chunks)
Dominated Novelty Search (DNS)¶
DNS is a Quality-Diversity algorithm that implements local competition through dynamic fitness transformations. Instead of using a fixed grid or archive, it computes a "dominated novelty" score: for each individual, the mean distance in descriptor space to its k-nearest fitter neighbors. This rewards individuals that are either the fittest, or occupy unique regions of descriptor space relative to better-performing solutions.
In [8]:
Copied!
def dominated_novelty_fn(fitness, descriptor, k=3):
"""Compute dominated novelty for a population.
Args:
fitness: Array of shape (N,). Lower is better. Invalid = inf.
descriptor: Array of shape (N, D).
k: Number of fitter neighbors for dominated novelty.
Returns:
Dominated novelty array of shape (N,).
"""
valid = fitness < jnp.inf
# Neighbors
neighbor = valid[:, None] & valid[None, :]
neighbor = jnp.fill_diagonal(neighbor, False, inplace=False)
# Fitter
fitter = fitness[:, None] >= fitness[None, :]
fitter = jnp.where(neighbor, fitter, False)
# Distance to neighbors
distance = jnp.linalg.norm(descriptor[:, None, :] - descriptor[None, :, :], axis=-1)
distance = jnp.where(neighbor, distance, jnp.inf)
# Distance to fitter neighbors
distance_fitter = jnp.where(fitter, distance, jnp.inf)
# Dominated Novelty - distance to k-fitter-nearest neighbors
values, indices = jax.vmap(partial(jax.lax.top_k, k=k))(-distance_fitter)
dominated_novelty = jnp.mean(
-values, axis=-1, where=jnp.take_along_axis(fitter, indices, axis=-1)
) # only the best (min fitness) individual should be nan
return dominated_novelty
def dominated_novelty_fn(fitness, descriptor, k=3):
"""Compute dominated novelty for a population.
Args:
fitness: Array of shape (N,). Lower is better. Invalid = inf.
descriptor: Array of shape (N, D).
k: Number of fitter neighbors for dominated novelty.
Returns:
Dominated novelty array of shape (N,).
"""
valid = fitness < jnp.inf
# Neighbors
neighbor = valid[:, None] & valid[None, :]
neighbor = jnp.fill_diagonal(neighbor, False, inplace=False)
# Fitter
fitter = fitness[:, None] >= fitness[None, :]
fitter = jnp.where(neighbor, fitter, False)
# Distance to neighbors
distance = jnp.linalg.norm(descriptor[:, None, :] - descriptor[None, :, :], axis=-1)
distance = jnp.where(neighbor, distance, jnp.inf)
# Distance to fitter neighbors
distance_fitter = jnp.where(fitter, distance, jnp.inf)
# Dominated Novelty - distance to k-fitter-nearest neighbors
values, indices = jax.vmap(partial(jax.lax.top_k, k=k))(-distance_fitter)
dominated_novelty = jnp.mean(
-values, axis=-1, where=jnp.take_along_axis(fitter, indices, axis=-1)
) # only the best (min fitness) individual should be nan
return dominated_novelty
In [9]:
Copied!
@nnx.dataclass
class DNSState(nnx.Pytree):
"""State for Dominated Novelty Search."""
population: Genotype = nnx.data()
fitness: jax.Array = nnx.data()
descriptor: jax.Array = nnx.data()
observations: jax.Array = nnx.data()
best_solution: Genotype = nnx.data()
best_fitness: jax.Array = nnx.data()
class DNS:
"""Dominated Novelty Search algorithm."""
def __init__(
self,
population_size: int,
batch_size: int,
mutation_std: float,
weight_concentration: float,
k: int = 3,
):
"""Initialize DNS.
Args:
population_size: Size of the population.
batch_size: Number of offspring to generate per generation.
mutation_std: Relative standard deviation for bounded Gaussian
mutations.
weight_concentration: Concentration parameter for Dirichlet weight
perturbation.
k: Number of nearest fitter neighbors for DNS competition.
"""
self.population_size = population_size
self.batch_size = batch_size
self.mutation_std = mutation_std
self.weight_concentration = weight_concentration
self.k = k
def init(
self,
population: Genotype,
fitness: jax.Array,
descriptor: jax.Array,
observations: jax.Array,
) -> DNSState:
"""Initialize DNS state from an evaluated population."""
best_idx = jnp.argmin(fitness)
return DNSState(
population=population,
fitness=fitness,
descriptor=descriptor,
observations=observations,
best_solution=jax.tree.map(lambda x: x[best_idx], population),
best_fitness=fitness[best_idx],
)
def ask(self, key: jax.Array, state: DNSState) -> Genotype:
"""Select parents and mutate to produce candidate population."""
key_select, key_mutate = jax.random.split(key)
# Select parents uniformly from valid members
valid = state.fitness < jnp.inf
p = jnp.where(valid, 1.0, 0.0)
p = p / jnp.maximum(jnp.sum(p), 1.0)
indices = jax.random.choice(
key_select,
state.fitness.shape[0],
shape=(self.batch_size,),
p=p,
)
parents = jax.tree.map(lambda x: x[indices], state.population)
# Mutate
keys = jax.random.split(key_mutate, self.batch_size)
mutate_fn = partial(
mutate_genotype,
mutation_std=self.mutation_std,
weight_concentration=self.weight_concentration,
)
population = jax.vmap(mutate_fn)(keys, parents)
return population
def tell(
self,
state: DNSState,
population: Genotype,
fitness: jax.Array,
descriptor: jax.Array,
observations: jax.Array,
) -> DNSState:
"""Apply mu+lambda selection with dominated novelty."""
# Concatenate parents and candidates
all_population = jax.tree.map(
lambda parent, candidate: jnp.concatenate([parent, candidate]),
state.population,
population,
)
all_fitness = jnp.concatenate([state.fitness, fitness])
all_descriptor = jnp.concatenate([state.descriptor, descriptor])
all_observations = jnp.concatenate([state.observations, observations])
# Dominated novelty competition
dominated_novelty = dominated_novelty_fn(all_fitness, all_descriptor, self.k)
valid = all_fitness < jnp.inf
meta_fitness = jnp.where(valid, dominated_novelty, -jnp.inf)
# Select top-N by meta-fitness
indices = jnp.argsort(meta_fitness, descending=True)[: self.population_size]
new_population = jax.tree.map(lambda x: x[indices], all_population)
new_fitness = all_fitness[indices]
new_descriptor = all_descriptor[indices]
new_observations = all_observations[indices]
# Update best solution
best_idx = jnp.argmin(new_fitness)
best_solution_in_population = jax.tree.map(
lambda x: x[best_idx], new_population
)
best_fitness_in_population = new_fitness[best_idx]
condition = best_fitness_in_population < state.best_fitness
best_solution = jax.tree.map(
lambda new, old: jnp.where(condition, new, old),
best_solution_in_population,
state.best_solution,
)
best_fitness = jnp.where(
condition, best_fitness_in_population, state.best_fitness
)
return DNSState(
population=new_population,
fitness=new_fitness,
descriptor=new_descriptor,
observations=new_observations,
best_solution=best_solution,
best_fitness=best_fitness,
)
@nnx.dataclass
class DNSState(nnx.Pytree):
"""State for Dominated Novelty Search."""
population: Genotype = nnx.data()
fitness: jax.Array = nnx.data()
descriptor: jax.Array = nnx.data()
observations: jax.Array = nnx.data()
best_solution: Genotype = nnx.data()
best_fitness: jax.Array = nnx.data()
class DNS:
"""Dominated Novelty Search algorithm."""
def __init__(
self,
population_size: int,
batch_size: int,
mutation_std: float,
weight_concentration: float,
k: int = 3,
):
"""Initialize DNS.
Args:
population_size: Size of the population.
batch_size: Number of offspring to generate per generation.
mutation_std: Relative standard deviation for bounded Gaussian
mutations.
weight_concentration: Concentration parameter for Dirichlet weight
perturbation.
k: Number of nearest fitter neighbors for DNS competition.
"""
self.population_size = population_size
self.batch_size = batch_size
self.mutation_std = mutation_std
self.weight_concentration = weight_concentration
self.k = k
def init(
self,
population: Genotype,
fitness: jax.Array,
descriptor: jax.Array,
observations: jax.Array,
) -> DNSState:
"""Initialize DNS state from an evaluated population."""
best_idx = jnp.argmin(fitness)
return DNSState(
population=population,
fitness=fitness,
descriptor=descriptor,
observations=observations,
best_solution=jax.tree.map(lambda x: x[best_idx], population),
best_fitness=fitness[best_idx],
)
def ask(self, key: jax.Array, state: DNSState) -> Genotype:
"""Select parents and mutate to produce candidate population."""
key_select, key_mutate = jax.random.split(key)
# Select parents uniformly from valid members
valid = state.fitness < jnp.inf
p = jnp.where(valid, 1.0, 0.0)
p = p / jnp.maximum(jnp.sum(p), 1.0)
indices = jax.random.choice(
key_select,
state.fitness.shape[0],
shape=(self.batch_size,),
p=p,
)
parents = jax.tree.map(lambda x: x[indices], state.population)
# Mutate
keys = jax.random.split(key_mutate, self.batch_size)
mutate_fn = partial(
mutate_genotype,
mutation_std=self.mutation_std,
weight_concentration=self.weight_concentration,
)
population = jax.vmap(mutate_fn)(keys, parents)
return population
def tell(
self,
state: DNSState,
population: Genotype,
fitness: jax.Array,
descriptor: jax.Array,
observations: jax.Array,
) -> DNSState:
"""Apply mu+lambda selection with dominated novelty."""
# Concatenate parents and candidates
all_population = jax.tree.map(
lambda parent, candidate: jnp.concatenate([parent, candidate]),
state.population,
population,
)
all_fitness = jnp.concatenate([state.fitness, fitness])
all_descriptor = jnp.concatenate([state.descriptor, descriptor])
all_observations = jnp.concatenate([state.observations, observations])
# Dominated novelty competition
dominated_novelty = dominated_novelty_fn(all_fitness, all_descriptor, self.k)
valid = all_fitness < jnp.inf
meta_fitness = jnp.where(valid, dominated_novelty, -jnp.inf)
# Select top-N by meta-fitness
indices = jnp.argsort(meta_fitness, descending=True)[: self.population_size]
new_population = jax.tree.map(lambda x: x[indices], all_population)
new_fitness = all_fitness[indices]
new_descriptor = all_descriptor[indices]
new_observations = all_observations[indices]
# Update best solution
best_idx = jnp.argmin(new_fitness)
best_solution_in_population = jax.tree.map(
lambda x: x[best_idx], new_population
)
best_fitness_in_population = new_fitness[best_idx]
condition = best_fitness_in_population < state.best_fitness
best_solution = jax.tree.map(
lambda new, old: jnp.where(condition, new, old),
best_solution_in_population,
state.best_solution,
)
best_fitness = jnp.where(
condition, best_fitness_in_population, state.best_fitness
)
return DNSState(
population=new_population,
fitness=new_fitness,
descriptor=new_descriptor,
observations=new_observations,
best_solution=best_solution,
best_fitness=best_fitness,
)
VAE for Unsupervised Descriptors¶
We train a VAE on rendered Lenia states to learn a latent descriptor space. The latent mean serves as the behavioral descriptor for DNS. The VAE is periodically retrained as the population evolves to adapt the descriptor space.
In [10]:
Copied!
def create_vae(key):
"""Create a fresh VAE instance."""
rngs = nnx.Rngs(key)
return VAE(
spatial_dims=spatial_dims,
features=vae_features,
latent_size=latent_size,
rngs=rngs,
)
def train_vae(vae, observations, key):
"""Train the VAE on observations.
Args:
vae: VAE instance.
observations: Array of shape (N, num_obs_steps, H, W, 3) uint8 in [0, 255].
key: Random key.
Returns:
Trained VAE.
"""
# Flatten to (N * num_obs_steps, H, W, 3) for training
flat_observations = observations.reshape(-1, *observations.shape[2:])
optimizer = nnx.Optimizer(vae, optax.adam(vae_learning_rate), wrt=nnx.Param)
@nnx.jit
def train_step(vae, optimizer, batch):
def loss_fn(vae):
logits, mean, logvar = vae(batch)
bce_loss = jnp.mean(binary_cross_entropy_with_logits(logits, batch))
kld_loss = jnp.mean(kl_divergence(mean, logvar))
return bce_loss + kld_loss
loss, grads = nnx.value_and_grad(loss_fn)(vae)
optimizer.update(vae, grads)
return loss
num_observations = flat_observations.shape[0]
for epoch in range(vae_num_epochs):
key, subkey = jax.random.split(key)
permutation = jax.random.permutation(subkey, num_observations)
for i in range(0, num_observations, vae_batch_size):
batch_idxs = permutation[i : i + vae_batch_size]
if batch_idxs.shape[0] < vae_batch_size:
continue
batch = flat_observations[batch_idxs].astype(jnp.float32) / 255.0
train_step(vae, optimizer, batch)
return vae
@nnx.jit
def encode_observations(vae, observations):
"""Encode a batch of observations to latent descriptors.
Args:
vae: VAE model.
observations: Array of shape (N, num_obs_steps, H, W, 3) uint8 in [0, 255].
Returns:
Array of latent descriptor vectors with shape (N, latent_size).
"""
return jnp.concatenate(
[
jax.vmap(descriptor_fn, in_axes=(None, 0))(vae, chunk)
for chunk in jnp.array_split(
observations, max(1, observations.shape[0] // 1024)
)
]
)
def create_vae(key):
"""Create a fresh VAE instance."""
rngs = nnx.Rngs(key)
return VAE(
spatial_dims=spatial_dims,
features=vae_features,
latent_size=latent_size,
rngs=rngs,
)
def train_vae(vae, observations, key):
"""Train the VAE on observations.
Args:
vae: VAE instance.
observations: Array of shape (N, num_obs_steps, H, W, 3) uint8 in [0, 255].
key: Random key.
Returns:
Trained VAE.
"""
# Flatten to (N * num_obs_steps, H, W, 3) for training
flat_observations = observations.reshape(-1, *observations.shape[2:])
optimizer = nnx.Optimizer(vae, optax.adam(vae_learning_rate), wrt=nnx.Param)
@nnx.jit
def train_step(vae, optimizer, batch):
def loss_fn(vae):
logits, mean, logvar = vae(batch)
bce_loss = jnp.mean(binary_cross_entropy_with_logits(logits, batch))
kld_loss = jnp.mean(kl_divergence(mean, logvar))
return bce_loss + kld_loss
loss, grads = nnx.value_and_grad(loss_fn)(vae)
optimizer.update(vae, grads)
return loss
num_observations = flat_observations.shape[0]
for epoch in range(vae_num_epochs):
key, subkey = jax.random.split(key)
permutation = jax.random.permutation(subkey, num_observations)
for i in range(0, num_observations, vae_batch_size):
batch_idxs = permutation[i : i + vae_batch_size]
if batch_idxs.shape[0] < vae_batch_size:
continue
batch = flat_observations[batch_idxs].astype(jnp.float32) / 255.0
train_step(vae, optimizer, batch)
return vae
@nnx.jit
def encode_observations(vae, observations):
"""Encode a batch of observations to latent descriptors.
Args:
vae: VAE model.
observations: Array of shape (N, num_obs_steps, H, W, 3) uint8 in [0, 255].
Returns:
Array of latent descriptor vectors with shape (N, latent_size).
"""
return jnp.concatenate(
[
jax.vmap(descriptor_fn, in_axes=(None, 0))(vae, chunk)
for chunk in jnp.array_split(
observations, max(1, observations.shape[0] // 1024)
)
]
)
Initialize Population¶
In [11]:
Copied!
# Sample initial population of genotypes
key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, 8 * population_size)
genotypes = jax.vmap(sample_genotype)(keys)
# Create initial VAE
key, subkey = jax.random.split(key)
vae = create_vae(subkey)
# Evaluate initial population
fitness, descriptor, info = evaluate_in_chunks(
vae, genotypes, chunk_size=population_size
)
print(f"Fitness range: [{fitness.min():.4f}, {fitness.max():.4f}]")
print(f"Non-degenerate: {jnp.sum(fitness < jnp.inf)}/{population_size}")
print(f"Descriptor shape: {descriptor.shape}")
# Train VAE on initial observations and re-encode
key, subkey = jax.random.split(key)
vae = train_vae(vae, info["observations"], subkey)
descriptor = encode_observations(vae, info["observations"])
# Sample initial population of genotypes
key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, 8 * population_size)
genotypes = jax.vmap(sample_genotype)(keys)
# Create initial VAE
key, subkey = jax.random.split(key)
vae = create_vae(subkey)
# Evaluate initial population
fitness, descriptor, info = evaluate_in_chunks(
vae, genotypes, chunk_size=population_size
)
print(f"Fitness range: [{fitness.min():.4f}, {fitness.max():.4f}]")
print(f"Non-degenerate: {jnp.sum(fitness < jnp.inf)}/{population_size}")
print(f"Descriptor shape: {descriptor.shape}")
# Train VAE on initial observations and re-encode
key, subkey = jax.random.split(key)
vae = train_vae(vae, info["observations"], subkey)
descriptor = encode_observations(vae, info["observations"])
Fitness range: [-0.0058, inf] Non-degenerate: 52/1024 Descriptor shape: (8192, 2)
In [12]:
Copied!
# Initialize DNS
dns = DNS(
population_size=population_size,
batch_size=batch_size,
mutation_std=mutation_std,
weight_concentration=weight_concentration,
k=k,
)
dns_state = dns.init(genotypes, fitness, descriptor, info["observations"])
# Initialize DNS
dns = DNS(
population_size=population_size,
batch_size=batch_size,
mutation_std=mutation_std,
weight_concentration=weight_concentration,
k=k,
)
dns_state = dns.init(genotypes, fitness, descriptor, info["observations"])
Evolution Loop¶
In [13]:
Copied!
@nnx.jit
def step(key, vae, dns_state):
"""Perform one generation: ask, evaluate, tell."""
population = dns.ask(key, dns_state)
fitness, descriptor, info = evaluate(vae, population)
dns_state = dns.tell(
dns_state, population, fitness, descriptor, info["observations"]
)
return dns_state
best_fitnesses = []
start = time.perf_counter()
for generation in range(num_generations):
key, subkey = jax.random.split(key)
dns_state = step(subkey, vae, dns_state)
best_fitnesses.append(float(dns_state.best_fitness))
# --- Periodic VAE retraining (AURORA) ---
if (generation + 1) % vae_retrain_interval == 0:
key, subkey = jax.random.split(key)
vae = create_vae(subkey)
key, subkey = jax.random.split(key)
vae = train_vae(vae, dns_state.observations, subkey)
# Re-encode with new VAE (no re-evaluation needed)
dns_state.descriptor = encode_observations(vae, dns_state.observations)
if generation % 16 == 0 or generation == num_generations - 1:
elapsed = time.perf_counter() - start
print(
f"Generation {generation:>3}/{num_generations} | {elapsed:6.1f}s "
f"| Best Fitness {float(dns_state.best_fitness):.4f} "
f"| Valid {int(jnp.sum(dns_state.fitness < jnp.inf))}/{population_size}"
)
print(
f"✨ Evolved for {num_generations} generations "
f"in {time.perf_counter() - start:.0f}s"
)
@nnx.jit
def step(key, vae, dns_state):
"""Perform one generation: ask, evaluate, tell."""
population = dns.ask(key, dns_state)
fitness, descriptor, info = evaluate(vae, population)
dns_state = dns.tell(
dns_state, population, fitness, descriptor, info["observations"]
)
return dns_state
best_fitnesses = []
start = time.perf_counter()
for generation in range(num_generations):
key, subkey = jax.random.split(key)
dns_state = step(subkey, vae, dns_state)
best_fitnesses.append(float(dns_state.best_fitness))
# --- Periodic VAE retraining (AURORA) ---
if (generation + 1) % vae_retrain_interval == 0:
key, subkey = jax.random.split(key)
vae = create_vae(subkey)
key, subkey = jax.random.split(key)
vae = train_vae(vae, dns_state.observations, subkey)
# Re-encode with new VAE (no re-evaluation needed)
dns_state.descriptor = encode_observations(vae, dns_state.observations)
if generation % 16 == 0 or generation == num_generations - 1:
elapsed = time.perf_counter() - start
print(
f"Generation {generation:>3}/{num_generations} | {elapsed:6.1f}s "
f"| Best Fitness {float(dns_state.best_fitness):.4f} "
f"| Valid {int(jnp.sum(dns_state.fitness < jnp.inf))}/{population_size}"
)
print(
f"✨ Evolved for {num_generations} generations "
f"in {time.perf_counter() - start:.0f}s"
)
Generation 0/256 | 4.8s | Best Fitness -0.0058 | Valid 121/1024
Generation 16/256 | 9.2s | Best Fitness -0.0211 | Valid 1024/1024
Generation 32/256 | 262.1s | Best Fitness -0.0435 | Valid 1024/1024
Generation 48/256 | 263.6s | Best Fitness -0.0534 | Valid 1024/1024
Generation 64/256 | 516.4s | Best Fitness -0.0594 | Valid 1024/1024
Generation 80/256 | 517.9s | Best Fitness -0.0859 | Valid 1024/1024
Generation 96/256 | 768.0s | Best Fitness -0.0943 | Valid 1024/1024
Generation 112/256 | 769.5s | Best Fitness -0.0967 | Valid 1024/1024
Generation 128/256 | 1021.6s | Best Fitness -0.0967 | Valid 1024/1024
Generation 144/256 | 1023.0s | Best Fitness -0.0967 | Valid 1024/1024
Generation 160/256 | 1271.8s | Best Fitness -0.0967 | Valid 1024/1024
Generation 176/256 | 1273.3s | Best Fitness -0.0967 | Valid 1024/1024
Generation 192/256 | 1525.4s | Best Fitness -0.0967 | Valid 1024/1024
Generation 208/256 | 1526.9s | Best Fitness -0.0967 | Valid 1024/1024
Generation 224/256 | 1777.4s | Best Fitness -0.0967 | Valid 1024/1024
Generation 240/256 | 1778.8s | Best Fitness -0.0967 | Valid 1024/1024
Generation 255/256 | 2028.2s | Best Fitness -0.0967 | Valid 1024/1024 ✨ Evolved for 256 generations in 2028s
Visualization¶
Fitness over generations¶
In [14]:
Copied!
plt.figure(figsize=(10, 4))
plt.plot(-jnp.array(best_fitnesses))
plt.xlabel("Generation")
plt.ylabel("Best Speed")
plt.title("Best Fitness (Speed) over Generations")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
plt.figure(figsize=(10, 4))
plt.plot(-jnp.array(best_fitnesses))
plt.xlabel("Generation")
plt.ylabel("Best Speed")
plt.title("Best Fitness (Speed) over Generations")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
Population scatter plot in latent space¶
In [15]:
Copied!
valid_mask = dns_state.fitness < jnp.inf
valid_descriptor = dns_state.descriptor[valid_mask]
valid_fitnesses = dns_state.fitness[valid_mask]
plt.figure(figsize=(8, 6))
scatter = plt.scatter(
valid_descriptor[:, 0],
valid_descriptor[:, 1],
c=-valid_fitnesses,
cmap="viridis",
s=20,
alpha=0.8,
)
plt.colorbar(scatter, label="Speed")
plt.xlabel("Latent dim 1")
plt.ylabel("Latent dim 2")
plt.title("Population in VAE Latent Space (colored by speed)")
plt.tight_layout()
plt.show()
valid_mask = dns_state.fitness < jnp.inf
valid_descriptor = dns_state.descriptor[valid_mask]
valid_fitnesses = dns_state.fitness[valid_mask]
plt.figure(figsize=(8, 6))
scatter = plt.scatter(
valid_descriptor[:, 0],
valid_descriptor[:, 1],
c=-valid_fitnesses,
cmap="viridis",
s=20,
alpha=0.8,
)
plt.colorbar(scatter, label="Speed")
plt.xlabel("Latent dim 1")
plt.ylabel("Latent dim 2")
plt.title("Population in VAE Latent Space (colored by speed)")
plt.tight_layout()
plt.show()
Best patterns¶
In [16]:
Copied!
idx = 0
sorted_idxs = jnp.argsort(dns_state.fitness)[idx : idx + 8]
genotypes = jax.tree.map(lambda x: x[sorted_idxs], dns_state.population)
@jax.jit
def simulate_and_render(genotype):
"""Simulate a genotype and render all frames."""
cs = Lenia(
spatial_dims=spatial_dims,
channel_size=channel_size,
R=R,
T=T,
state_scale=state_scale,
kernel_fn=gaussian_kernel_fn,
growth_fn=exponential_growth_fn,
rule_params=genotype.rule_params,
)
_, states = cs(genotype.state_init, num_steps=num_steps, return_states=True)
states = jnp.concatenate([genotype.state_init[None], states])
return jax.vmap(cs.render)(states)
# Pull the frames to host memory: the upscale quadruples an array this large
frames = np.asarray(jax.vmap(simulate_and_render)(genotypes))
mediapy.show_videos(frames.repeat(2, axis=-3).repeat(2, axis=-2), fps=30, columns=8)
idx = 0
sorted_idxs = jnp.argsort(dns_state.fitness)[idx : idx + 8]
genotypes = jax.tree.map(lambda x: x[sorted_idxs], dns_state.population)
@jax.jit
def simulate_and_render(genotype):
"""Simulate a genotype and render all frames."""
cs = Lenia(
spatial_dims=spatial_dims,
channel_size=channel_size,
R=R,
T=T,
state_scale=state_scale,
kernel_fn=gaussian_kernel_fn,
growth_fn=exponential_growth_fn,
rule_params=genotype.rule_params,
)
_, states = cs(genotype.state_init, num_steps=num_steps, return_states=True)
states = jnp.concatenate([genotype.state_init[None], states])
return jax.vmap(cs.render)(states)
# Pull the frames to host memory: the upscale quadruples an array this large
frames = np.asarray(jax.vmap(simulate_and_render)(genotypes))
mediapy.show_videos(frames.repeat(2, axis=-3).repeat(2, axis=-2), fps=30, columns=8)