Growing Neural Cellular Automata ¶

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.nn.pool import Pool
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.nn.pool import Pool 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
pool_size = 1_024
batch_size = 8
learning_rate = 2e-3

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 pool_size = 1_024 batch_size = 8 learning_rate = 2e-3 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,
            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, 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¶

Pool¶

In [10]:

state = jax.vmap(lambda _: sample_state())(jnp.zeros(pool_size))

pool = Pool.create({"state": state})

state = jax.vmap(lambda _: sample_state())(jnp.zeros(pool_size)) pool = Pool.create({"state": state})

Optimizer¶

In [11]:

lr_sched = optax.linear_schedule(
    init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=2_000
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=lr_sched),
)

update_params = nnx.All(nnx.Param, nnx.PathContains("update"))
optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

lr_sched = optax.linear_schedule( init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=2_000 ) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(learning_rate=lr_sched), ) update_params = nnx.All(nnx.Param, nnx.PathContains("update")) optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

Loss¶

In [12]:

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 [13]:

@nnx.jit
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, 0),
        )
    )(cs, state)

    # Sample a random step
    idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps)
    state = state[jnp.arange(batch_size), idx]

    loss = mse(state)
    return loss, state

@nnx.jit 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, 0), ) )(cs, state) # Sample a random step idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps) state = state[jnp.arange(batch_size), idx] loss = mse(state) return loss, state

Train step¶

In [14]:

@nnx.jit
def train_step(cs, optimizer, pool, key):
    """Train step."""
    sample_key, loss_key = jax.random.split(key)

    # Sample from pool
    pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False)
    current_state = batch["state"]

    # Sort by descending loss
    sort_idx = jnp.argsort(jax.vmap(mse)(current_state), descending=True)
    pool_idx = pool_idx[sort_idx]
    current_state = current_state[sort_idx]

    # Sample a new state to replace the worst
    new_state = sample_state()
    current_state = current_state.at[0].set(new_state)

    (loss, current_state), grad = nnx.value_and_grad(
        loss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params)
    )(cs, current_state, loss_key)
    optimizer.update(cs, grad)

    pool = pool.update(pool_idx, {"state": current_state})
    return loss, pool

@nnx.jit def train_step(cs, optimizer, pool, key): """Train step.""" sample_key, loss_key = jax.random.split(key) # Sample from pool pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False) current_state = batch["state"] # Sort by descending loss sort_idx = jnp.argsort(jax.vmap(mse)(current_state), descending=True) pool_idx = pool_idx[sort_idx] current_state = current_state[sort_idx] # Sample a new state to replace the worst new_state = sample_state() current_state = current_state.at[0].set(new_state) (loss, current_state), grad = nnx.value_and_grad( loss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params) )(cs, current_state, loss_key) optimizer.update(cs, grad) pool = pool.update(pool_idx, {"state": current_state}) return loss, pool

Main loop¶

In [15]:

num_train_steps = 8_192
print_interval = 128

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    loss, pool = train_step(cs, optimizer, pool, subkey)
    losses.append(loss)

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

print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

num_train_steps = 8_192 print_interval = 128 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) loss, pool = train_step(cs, optimizer, pool, subkey) losses.append(loss) if i % print_interval == 0 or i == num_train_steps - 1: avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:]) elapsed = time.perf_counter() - start print(f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e}") print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

Step    0/8192 |    3.8s | Loss 3.043e-02

Step  128/8192 |   10.0s | Loss 2.636e-02

Step  256/8192 |   16.2s | Loss 1.345e-02

Step  384/8192 |   22.4s | Loss 1.191e-02

Step  512/8192 |   28.6s | Loss 9.642e-03

Step  640/8192 |   34.7s | Loss 7.470e-03

Step  768/8192 |   40.9s | Loss 5.797e-03

Step  896/8192 |   47.1s | Loss 4.204e-03

Step 1024/8192 |   53.3s | Loss 2.987e-03

Step 1152/8192 |   59.4s | Loss 2.346e-03

Step 1280/8192 |   65.6s | Loss 1.815e-03

Step 1408/8192 |   71.7s | Loss 1.591e-03

Step 1536/8192 |   77.9s | Loss 1.200e-03

Step 1664/8192 |   84.1s | Loss 9.193e-04

Step 1792/8192 |   90.2s | Loss 6.969e-04

Step 1920/8192 |   96.4s | Loss 5.833e-04

Step 2048/8192 |  102.5s | Loss 5.069e-04

Step 2176/8192 |  108.7s | Loss 4.668e-04

Step 2304/8192 |  114.9s | Loss 4.335e-04

Step 2432/8192 |  121.0s | Loss 4.131e-04

Step 2560/8192 |  127.2s | Loss 3.866e-04

Step 2688/8192 |  133.3s | Loss 3.783e-04

Step 2816/8192 |  139.5s | Loss 3.590e-04

Step 2944/8192 |  145.7s | Loss 3.546e-04

Step 3072/8192 |  151.8s | Loss 3.432e-04

Step 3200/8192 |  158.0s | Loss 3.186e-04

Step 3328/8192 |  164.2s | Loss 3.224e-04

Step 3456/8192 |  170.4s | Loss 2.983e-04

Step 3584/8192 |  176.5s | Loss 3.161e-04

Step 3712/8192 |  182.7s | Loss 3.068e-04

Step 3840/8192 |  188.9s | Loss 2.861e-04

Step 3968/8192 |  195.1s | Loss 2.849e-04

Step 4096/8192 |  201.2s | Loss 2.630e-04

Step 4224/8192 |  207.4s | Loss 2.719e-04

Step 4352/8192 |  213.6s | Loss 2.409e-04

Step 4480/8192 |  219.8s | Loss 2.490e-04

Step 4608/8192 |  226.0s | Loss 2.536e-04

Step 4736/8192 |  232.1s | Loss 2.350e-04

Step 4864/8192 |  238.3s | Loss 2.314e-04

Step 4992/8192 |  244.5s | Loss 2.215e-04

Step 5120/8192 |  250.6s | Loss 2.180e-04

Step 5248/8192 |  256.8s | Loss 2.046e-04

Step 5376/8192 |  263.0s | Loss 2.073e-04

Step 5504/8192 |  269.1s | Loss 1.904e-04

Step 5632/8192 |  275.3s | Loss 2.070e-04

Step 5760/8192 |  281.5s | Loss 1.820e-04

Step 5888/8192 |  287.6s | Loss 1.786e-04

Step 6016/8192 |  293.8s | Loss 1.800e-04

Step 6144/8192 |  299.9s | Loss 1.670e-04

Step 6272/8192 |  306.1s | Loss 1.757e-04

Step 6400/8192 |  312.2s | Loss 1.626e-04

Step 6528/8192 |  318.4s | Loss 1.593e-04

Step 6656/8192 |  324.6s | Loss 1.502e-04

Step 6784/8192 |  330.7s | Loss 1.615e-04

Step 6912/8192 |  336.9s | Loss 1.604e-04

Step 7040/8192 |  343.0s | Loss 1.437e-04

Step 7168/8192 |  349.2s | Loss 1.424e-04

Step 7296/8192 |  355.3s | Loss 1.265e-04

Step 7424/8192 |  361.5s | Loss 1.366e-04

Step 7552/8192 |  367.6s | Loss 1.183e-04

Step 7680/8192 |  373.8s | Loss 1.390e-04

Step 7808/8192 |  380.0s | Loss 1.348e-04

Step 7936/8192 |  386.1s | Loss 1.258e-04

Step 8064/8192 |  392.3s | Loss 1.264e-04

Step 8191/8192 |  398.4s | Loss 1.138e-04
✨ Trained for 8192 steps in 398s

Run¶

In [16]:

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 [17]:

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(2, axis=-3).repeat(2, axis=-2))
mediapy.show_images(frames_final_rgba.repeat(2, axis=-3).repeat(2, 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(2, axis=-3).repeat(2, axis=-2)) mediapy.show_images(frames_final_rgba.repeat(2, axis=-3).repeat(2, axis=-2))

In [18]:

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

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

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