Growing Neural Cellular Automata with Evolution Strategies ¶

Installation¶

You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:

In [1]:

%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]:

%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]:

import time

import jax
import jax.numpy as jnp
import mediapy
import optax
from flax import nnx
from jax import Array

from cax.core import ComplexSystem
from cax.core.perceive import ConvPerceive, grad_kernel, identity_kernel
from cax.core.update import NCAUpdate
from cax.utils import clip_and_uint8, get_emoji_array, rgba_to_rgb

import time import jax import jax.numpy as jnp import mediapy import optax from flax import nnx from jax import Array from cax.core import ComplexSystem from cax.core.perceive import ConvPerceive, grad_kernel, identity_kernel from cax.core.update import NCAUpdate from cax.utils import clip_and_uint8, get_emoji_array, rgba_to_rgb

Configuration¶

In [4]:

seed = 0

channel_size = 16
num_kernels = 3
hidden_size = 128
cell_dropout_rate = 0.5

num_steps = 128
population_size = 512
batch_size = 1

emoji = "🦎"
size = 40
pad_width = 16

key = jax.random.key(seed)
rngs = nnx.Rngs(seed)

seed = 0 channel_size = 16 num_kernels = 3 hidden_size = 128 cell_dropout_rate = 0.5 num_steps = 128 population_size = 512 batch_size = 1 emoji = "🦎" size = 40 pad_width = 16 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

In [5]:

y = get_emoji_array(emoji, size, pad_width)

mediapy.show_image(y)

y = get_emoji_array(emoji, size, pad_width) mediapy.show_image(y)

Instantiate system¶

In [6]:

class GrowingNCA(ComplexSystem):
    """Growing Neural Cellular Automata class."""

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize Growing NCA.

        Args:
            rngs: rng key.

        """
        self.perceive = ConvPerceive(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            feature_group_count=channel_size,
            rngs=rngs,
        )
        self.update = NCAUpdate(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            hidden_layer_sizes=(hidden_size,),
            cell_dropout_rate=cell_dropout_rate,
            step_size=0.5,
            zeros_init=True,
            rngs=rngs,
        )

        # Initialize kernel with sobel filters
        kernel = jnp.concatenate(
            [identity_kernel(num_dims=2), grad_kernel(num_dims=2)], axis=-1
        )
        kernel = jnp.expand_dims(
            jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2
        )
        self.perceive.conv.kernel[...] = kernel

    def _step(self, state: Array, input: Array | None = None) -> Array:
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        return next_state

    @nnx.jit
    def render(self, state):
        """Render state to RGB."""
        rgba = state[..., -4:]
        rgb = rgba_to_rgb(rgba)

        # Clip values to valid range and convert to uint8
        return clip_and_uint8(rgb)

    @nnx.jit
    def render_rgba(self, state):
        """Render state to RGBA."""
        rgba = state[..., -4:]

        # Clip values to valid range and convert to uint8
        return clip_and_uint8(rgba)

class GrowingNCA(ComplexSystem): """Growing Neural Cellular Automata class.""" def __init__(self, *, rngs: nnx.Rngs): """Initialize Growing NCA. Args: rngs: rng key. """ self.perceive = ConvPerceive( channel_size=channel_size, perception_size=num_kernels * channel_size, feature_group_count=channel_size, rngs=rngs, ) self.update = NCAUpdate( channel_size=channel_size, perception_size=num_kernels * channel_size, hidden_layer_sizes=(hidden_size,), cell_dropout_rate=cell_dropout_rate, step_size=0.5, zeros_init=True, rngs=rngs, ) # Initialize kernel with sobel filters kernel = jnp.concatenate( [identity_kernel(num_dims=2), grad_kernel(num_dims=2)], axis=-1 ) kernel = jnp.expand_dims( jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2 ) self.perceive.conv.kernel[...] = kernel def _step(self, state: Array, input: Array | None = None) -> Array: perception = self.perceive(state) next_state = self.update(state, perception, input) return next_state @nnx.jit def render(self, state): """Render state to RGB.""" rgba = state[..., -4:] rgb = rgba_to_rgb(rgba) # Clip values to valid range and convert to uint8 return clip_and_uint8(rgb) @nnx.jit def render_rgba(self, state): """Render state to RGBA.""" rgba = state[..., -4:] # Clip values to valid range and convert to uint8 return clip_and_uint8(rgba)

In [7]:

cs = GrowingNCA(rngs=rngs)

cs = GrowingNCA(rngs=rngs)

In [8]:

params = nnx.state(cs, nnx.Param)
print("Number of params:", sum(x.size for x in jax.tree.leaves(params)))

params = nnx.state(cs, nnx.Param) print("Number of params:", sum(x.size for x in jax.tree.leaves(params)))

Number of params: 8768

Sample initial state¶

In [9]:

def sample_state():
    """Sample a state with a single alive cell."""
    spatial_dims = y.shape[:2]

    # Init state
    state = jnp.zeros(spatial_dims + (channel_size,))

    # Set the center cell to alive, with hidden channels at one
    mid = tuple(size // 2 for size in spatial_dims)
    state = state.at[mid[0], mid[1], -1].set(1.0)
    return state.at[mid[0], mid[1], :-4].set(1.0)

def sample_state(): """Sample a state with a single alive cell.""" spatial_dims = y.shape[:2] # Init state state = jnp.zeros(spatial_dims + (channel_size,)) # Set the center cell to alive, with hidden channels at one mid = tuple(size // 2 for size in spatial_dims) state = state.at[mid[0], mid[1], -1].set(1.0) return state.at[mid[0], mid[1], :-4].set(1.0)

Train¶

Evolution Strategy¶

In [10]:

trainable_filter = nnx.All(nnx.Param, nnx.PathContains("update"))
solution = nnx.state(cs, trainable_filter)

trainable_filter = nnx.All(nnx.Param, nnx.PathContains("update")) solution = nnx.state(cs, trainable_filter)

In [11]:

import time

from evosax.algorithms import Open_ES as ES

learning_rate = 0.002
warmup_generations = 1_000
std = 0.008

es = ES(
    population_size=population_size,
    solution=solution,
    optimizer=optax.adam(
        learning_rate=optax.linear_schedule(0.0, learning_rate, warmup_generations)
    ),
    std_schedule=optax.constant_schedule(std),
)

es_params = es.default_params

import time from evosax.algorithms import Open_ES as ES learning_rate = 0.002 warmup_generations = 1_000 std = 0.008 es = ES( population_size=population_size, solution=solution, optimizer=optax.adam( learning_rate=optax.linear_schedule(0.0, learning_rate, warmup_generations) ), std_schedule=optax.constant_schedule(std), ) es_params = es.default_params

In [12]:

key, subkey = jax.random.split(key)
es_state = es.init(subkey, solution, es_params)

key, subkey = jax.random.split(key) es_state = es.init(subkey, solution, es_params)

Loss¶

In [13]:

def mse(state):
    """Mean Squared Error."""
    return jnp.mean(jnp.square(state[..., -4:] - y))

def mse(state): """Mean Squared Error.""" return jnp.mean(jnp.square(state[..., -4:] - y))

In [14]:

def loss_fn(cs, state, key):
    """Loss function."""
    state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
    _, state = nnx.split_rngs(splits=batch_size)(
        nnx.vmap(
            lambda cs, state: cs(state, num_steps=num_steps, return_states=True),
            in_axes=(state_axes, None),
        )
    )(cs, state)

    # Match the target over the second half of the rollout
    loss = mse(state[:, num_steps // 2 :])
    return loss

def loss_fn(cs, state, key): """Loss function.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) _, state = nnx.split_rngs(splits=batch_size)( nnx.vmap( lambda cs, state: cs(state, num_steps=num_steps, return_states=True), in_axes=(state_axes, None), ) )(cs, state) # Match the target over the second half of the rollout loss = mse(state[:, num_steps // 2 :]) return loss

Train step¶

In [15]:

@nnx.jit
def train_step(cs, es_state, key):
    """Train step."""
    key, key_ask, key_eval, key_tell = jax.random.split(key, 4)

    state = sample_state()

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

    # Evaluate the fitness of the population
    nnx.update(cs, population)

    state_axes = nnx.StateAxes({trainable_filter: 0, nnx.RngState: 0, ...: None})
    fitness = nnx.split_rngs(splits=population_size)(
        nnx.vmap(
            loss_fn,
            in_axes=(state_axes, None, None),
        )
    )(cs, state, key_eval)

    # Update the evolution strategy
    fitness = jnp.nan_to_num(fitness, nan=1e3, posinf=1e3)
    es_state, metrics = es.tell(key_tell, population, fitness, es_state, es_params)

    return es_state, metrics

@nnx.jit def train_step(cs, es_state, key): """Train step.""" key, key_ask, key_eval, key_tell = jax.random.split(key, 4) state = sample_state() # Generate a set of candidate solutions to evaluate population, es_state = es.ask(key_ask, es_state, es_params) # Evaluate the fitness of the population nnx.update(cs, population) state_axes = nnx.StateAxes({trainable_filter: 0, nnx.RngState: 0, ...: None}) fitness = nnx.split_rngs(splits=population_size)( nnx.vmap( loss_fn, in_axes=(state_axes, None, None), ) )(cs, state, key_eval) # Update the evolution strategy fitness = jnp.nan_to_num(fitness, nan=1e3, posinf=1e3) es_state, metrics = es.tell(key_tell, population, fitness, es_state, es_params) return es_state, metrics

Main loop¶

In [16]:

num_generations = 5000
print_interval = 500

losses = []
start = time.perf_counter()
for i in range(num_generations):
    key, subkey = jax.random.split(key)
    es_state, metrics = train_step(cs, es_state, subkey)

    losses.append(metrics["best_fitness_in_generation"])
    if i % print_interval == 0 or i == num_generations - 1:
        avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:])
        elapsed = time.perf_counter() - start
        print(
            f"Generation {i:>5}/{num_generations} | {elapsed:6.1f}s "
            f"| Loss {avg_loss:.3e}"
        )

print(
    f"✨ Evolved for {num_generations} generations "
    f"in {time.perf_counter() - start:.0f}s"
)

num_generations = 5000 print_interval = 500 losses = [] start = time.perf_counter() for i in range(num_generations): key, subkey = jax.random.split(key) es_state, metrics = train_step(cs, es_state, subkey) losses.append(metrics["best_fitness_in_generation"]) if i % print_interval == 0 or i == num_generations - 1: avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:]) elapsed = time.perf_counter() - start print( f"Generation {i:>5}/{num_generations} | {elapsed:6.1f}s " f"| Loss {avg_loss:.3e}" ) print( f"✨ Evolved for {num_generations} generations " f"in {time.perf_counter() - start:.0f}s" )

W0905 03:57:58.508877  647484 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.

W0905 03:57:59.711056  647464 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.

Generation     0/5000 |    5.5s | Loss 2.213e-02

Generation   500/5000 |  656.2s | Loss 2.411e-02

Generation  1000/5000 | 1305.1s | Loss 1.234e-02

Generation  1500/5000 | 1954.1s | Loss 9.912e-03

Generation  2000/5000 | 2603.3s | Loss 7.635e-03

Generation  2500/5000 | 3252.7s | Loss 5.685e-03

Generation  3000/5000 | 3902.0s | Loss 4.591e-03

Generation  3500/5000 | 4550.8s | Loss 4.038e-03

Generation  4000/5000 | 5199.8s | Loss 3.815e-03

Generation  4500/5000 | 5849.1s | Loss 3.625e-03

Generation  4999/5000 | 6497.1s | Loss 3.428e-03
✨ Evolved for 5000 generations in 6497s

In [17]:

trainable_params = es.get_mean(es_state)
nnx.update(cs, trainable_params)

trainable_params = es.get_mean(es_state) nnx.update(cs, trainable_params)

Run¶

In [18]:

num_examples = 8

state_init = jax.vmap(lambda _: sample_state())(jnp.zeros(num_examples))

state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
state_final, states = nnx.split_rngs(splits=num_examples)(
    nnx.vmap(
        lambda cs, state_init: cs(
            state_init, num_steps=2 * num_steps, return_states=True
        ),
        in_axes=(state_axes, 0),
    )
)(cs, state_init)

num_examples = 8 state_init = jax.vmap(lambda _: sample_state())(jnp.zeros(num_examples)) state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state_final, states = nnx.split_rngs(splits=num_examples)( nnx.vmap( lambda cs, state_init: cs( state_init, num_steps=2 * num_steps, return_states=True ), in_axes=(state_axes, 0), ) )(cs, state_init)

Visualize¶

In [19]:

frames_final = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, state_final)
frames_final_rgba = nnx.vmap(
    lambda cs, state: cs.render_rgba(state),
    in_axes=(None, 0),
)(cs, state_final)

mediapy.show_images(frames_final.repeat(3, axis=-3).repeat(3, axis=-2))
mediapy.show_images(frames_final_rgba.repeat(3, axis=-3).repeat(3, axis=-2))

frames_final = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, state_final) frames_final_rgba = nnx.vmap( lambda cs, state: cs.render_rgba(state), in_axes=(None, 0), )(cs, state_final) mediapy.show_images(frames_final.repeat(3, axis=-3).repeat(3, axis=-2)) mediapy.show_images(frames_final_rgba.repeat(3, axis=-3).repeat(3, axis=-2))

In [20]:

states = jnp.concatenate([state_init[:, None], states], axis=1)
frames = nnx.vmap(
    lambda cs, states: cs.render(states),
    in_axes=(None, 0),
)(cs, states)

mediapy.show_videos(frames.repeat(3, axis=-3).repeat(3, axis=-2))

states = jnp.concatenate([state_init[:, None], states], axis=1) frames = nnx.vmap( lambda cs, states: cs.render(states), in_axes=(None, 0), )(cs, states) mediapy.show_videos(frames.repeat(3, axis=-3).repeat(3, axis=-2))