Isotropic Neural Cellular Automata ¶

A growing neural cellular automaton reads its neighborhood through Sobel filters fixed to the pixel grid. Those filters answer "how does this change as I move right?", so the rule inherits a sense of direction from the lattice it happens to live on. Rotate the world and it behaves differently.

An isotropic automaton must not. Its update has to commute with rotation,

$$\mathrm{step}(R_\theta \cdot s) = R_\theta \cdot \mathrm{step}(s),$$

which means the perception may only measure quantities a rotation leaves alone.

That constraint has an immediate consequence. If the rule cannot tell which way is up and the seed is a single cell, nothing in the system distinguishes one orientation from another, so the pattern is grown at whatever angle the automaton settles on. The target must therefore be scored at every angle, and the automaton must break the symmetry itself. It does, using the randomness in its own stochastic updates.

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, Perceive, grad2_kernel, grad_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, Perceive, grad2_kernel, grad_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
hidden_size = 128
cell_dropout_rate = 0.5
num_directions = 4

num_steps = 96
pool_size = 1_024
batch_size = 8
learning_rate = 1e-3
overflow_weight = 1.0
diff_weight = 10.0

num_train_steps = 16_384
num_angles = 256  # at least pi * width, one sample per pixel on the outer ring

emoji = "🦎"
size = 48
pad_width = 12

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

seed = 0 channel_size = 16 hidden_size = 128 cell_dropout_rate = 0.5 num_directions = 4 num_steps = 96 pool_size = 1_024 batch_size = 8 learning_rate = 1e-3 overflow_weight = 1.0 diff_weight = 10.0 num_train_steps = 16_384 num_angles = 256 # at least pi * width, one sample per pixel on the outer ring emoji = "🦎" size = 48 pad_width = 12 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

Colors are premultiplied by alpha so that transparent regions carry no color to match.

An elongated body nearly matches itself under a half turn, so the loss below --- which scores the result at every rotation and keeps the best --- has a second, spurious minimum at 180 degrees. Left alone the automaton settles between the two and grows a gecko with neither head nor tail.

One extra target channel separates them: which half of the body a cell belongs to. It is not symmetric under a half turn, so the two minima stop being interchangeable.

In [5]:

y = get_emoji_array(emoji, size, pad_width)

spatial_dims = y.shape[:2]

half = jnp.sign(jnp.linspace(-1.0, 1.0, spatial_dims[0]))[:, None] * y[..., 3] * 0.5
y = jnp.concatenate([half[..., None], y], axis=-1)

mediapy.show_image(rgba_to_rgb(y[..., 1:]))

y = get_emoji_array(emoji, size, pad_width) spatial_dims = y.shape[:2] half = jnp.sign(jnp.linspace(-1.0, 1.0, spatial_dims[0]))[:, None] * y[..., 3] * 0.5 y = jnp.concatenate([half[..., None], y], axis=-1) mediapy.show_image(rgba_to_rgb(y[..., 1:]))

Instantiate system¶

Perception¶

The rotation invariants of a collection of plane vectors are generated by their pairwise inner products and their pairwise determinants,

$$\langle \nabla c_i, \nabla c_j \rangle \qquad\text{and}\qquad \det(\nabla c_i, \nabla c_j) = \partial_x c_i \, \partial_y c_j - \partial_y c_i \, \partial_x c_j .$$

They behave differently under a reflection: the inner products are unchanged, while the determinants change sign. So the two carry different information, and the difference is exactly handedness.

An isotropic perception built only from gradient magnitudes $\lVert \nabla c_i \rVert$ — the diagonal inner products — is therefore blind to a mirror image, and cannot prefer a lizard to its reflection. Adding a few determinants makes the rule chiral, which lets the objective ask for the shape we actually want rather than forgiving either.

The Laplacian is the Moore-neighborhood stencil rather than the default five-point one, whose leading error favors the grid axes — a bias a system asked to behave the same in every direction would find and use.

In [6]:

class IsotropicPerceive(Perceive[Array, Array]):
    """Perceive only what survives a rotation of the world."""

    def __init__(self, *, channel_size: int, num_directions: int, rngs: nnx.Rngs):
        """Initialize isotropic perception.

        Args:
            channel_size: Number of state channels.
            num_directions: Number of channels whose gradients also contribute the
                determinants that make the perception sensitive to handedness.
            rngs: rng key.

        """
        self.channel_size = channel_size
        self.num_directions = num_directions
        self.conv = ConvPerceive(
            channel_size=channel_size,
            perception_size=3 * channel_size,
            feature_group_count=channel_size,
            padding="CIRCULAR",
            rngs=rngs,
        )

        # Initialize kernel with sobel filters and an isotropic laplacian
        kernel = jnp.concatenate(
            [grad_kernel(num_dims=2), grad2_kernel(num_dims=2, neighborhood="moore")],
            axis=-1,
        )
        kernel = jnp.expand_dims(
            jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2
        )
        self.conv.conv.kernel[...] = kernel

    def __call__(self, state: Array) -> Array:
        """Perceive the neighborhood of every cell.

        Args:
            state: State of the cellular automaton.

        Returns:
            The state, its gradient magnitudes, its laplacian, and the determinants
            between the gradients of the first `num_directions` channels.

        """
        filtered = self.conv(state).reshape(*state.shape[:-1], self.channel_size, 3)
        grad_x, grad_y, laplacian = filtered[..., 0], filtered[..., 1], filtered[..., 2]
        grad_norm = jnp.sqrt(jnp.square(grad_x) + jnp.square(grad_y) + 1e-8)

        # Determinants between gradients, which change sign under a reflection.
        # Taken between state channels rather than learned combinations of them, so the
        # overflow penalty that holds the state in range keeps them bounded too.
        cut = slice(None, self.num_directions)
        det = (
            grad_x[..., cut, None] * grad_y[..., None, cut]
            - grad_y[..., cut, None] * grad_x[..., None, cut]
        )
        rows, cols = jnp.triu_indices(self.num_directions, k=1)

        return jnp.concatenate(
            [state, grad_norm, laplacian, det[..., rows, cols]], axis=-1
        )

class IsotropicPerceive(Perceive[Array, Array]): """Perceive only what survives a rotation of the world.""" def __init__(self, *, channel_size: int, num_directions: int, rngs: nnx.Rngs): """Initialize isotropic perception. Args: channel_size: Number of state channels. num_directions: Number of channels whose gradients also contribute the determinants that make the perception sensitive to handedness. rngs: rng key. """ self.channel_size = channel_size self.num_directions = num_directions self.conv = ConvPerceive( channel_size=channel_size, perception_size=3 * channel_size, feature_group_count=channel_size, padding="CIRCULAR", rngs=rngs, ) # Initialize kernel with sobel filters and an isotropic laplacian kernel = jnp.concatenate( [grad_kernel(num_dims=2), grad2_kernel(num_dims=2, neighborhood="moore")], axis=-1, ) kernel = jnp.expand_dims( jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2 ) self.conv.conv.kernel[...] = kernel def __call__(self, state: Array) -> Array: """Perceive the neighborhood of every cell. Args: state: State of the cellular automaton. Returns: The state, its gradient magnitudes, its laplacian, and the determinants between the gradients of the first `num_directions` channels. """ filtered = self.conv(state).reshape(*state.shape[:-1], self.channel_size, 3) grad_x, grad_y, laplacian = filtered[..., 0], filtered[..., 1], filtered[..., 2] grad_norm = jnp.sqrt(jnp.square(grad_x) + jnp.square(grad_y) + 1e-8) # Determinants between gradients, which change sign under a reflection. # Taken between state channels rather than learned combinations of them, so the # overflow penalty that holds the state in range keeps them bounded too. cut = slice(None, self.num_directions) det = ( grad_x[..., cut, None] * grad_y[..., None, cut] - grad_y[..., cut, None] * grad_x[..., None, cut] ) rows, cols = jnp.triu_indices(self.num_directions, k=1) return jnp.concatenate( [state, grad_norm, laplacian, det[..., rows, cols]], axis=-1 )

System¶

In [7]:

class IsotropicNCA(ComplexSystem):
    """Isotropic Neural Cellular Automata class."""

    remat = True

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

        Args:
            rngs: rng key.

        """
        self.perceive = IsotropicPerceive(
            channel_size=channel_size, num_directions=num_directions, rngs=rngs
        )
        self.update = NCAUpdate(
            channel_size=channel_size,
            perception_size=3 * channel_size
            + num_directions * (num_directions - 1) // 2,
            hidden_layer_sizes=(hidden_size,),
            cell_dropout_rate=cell_dropout_rate,
            zeros_init=True,
            rngs=rngs,
        )

    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."""
        return clip_and_uint8(rgba_to_rgb(state[..., -4:]))

class IsotropicNCA(ComplexSystem): """Isotropic Neural Cellular Automata class.""" remat = True def __init__(self, *, rngs: nnx.Rngs): """Initialize Isotropic NCA. Args: rngs: rng key. """ self.perceive = IsotropicPerceive( channel_size=channel_size, num_directions=num_directions, rngs=rngs ) self.update = NCAUpdate( channel_size=channel_size, perception_size=3 * channel_size + num_directions * (num_directions - 1) // 2, hidden_layer_sizes=(hidden_size,), cell_dropout_rate=cell_dropout_rate, zeros_init=True, rngs=rngs, ) 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.""" return clip_and_uint8(rgba_to_rgb(state[..., -4:]))

In [8]:

cs = IsotropicNCA(rngs=rngs)

cs = IsotropicNCA(rngs=rngs)

In [9]:

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: 9536

Sample initial state¶

A single live cell in the middle. It is as symmetric a starting point as there is, which is the whole difficulty: the automaton has to manufacture an orientation out of nothing.

In [10]:

def sample_state():
    """Sample a state with a single alive cell."""
    state = jnp.zeros(spatial_dims + (channel_size,))

    mid = tuple(dim // 2 for dim 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.""" state = jnp.zeros(spatial_dims + (channel_size,)) mid = tuple(dim // 2 for dim 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 [11]:

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

lr_sched = optax.linear_schedule(
    init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=10_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=10_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¶

The automaton grows the target at an angle nobody chose, so the objective scores it at every angle and keeps the best,

$$\mathcal{L}(s) = \min_{\theta} \, \bigl\lVert P[s] - R_\theta \, P[y] \bigr\rVert^2 = \min_{\theta} \, \lVert P[s] \rVert^2 + \lVert P[y] \rVert^2 - 2 \, \langle P[s], R_\theta P[y] \rangle ,$$

where $P$ resamples an image onto a polar grid, turning a rotation into a shift along the angular axis. Only the last term depends on $\theta$, and as a function of the shift it is a circular cross-correlation — so all $\theta$ are evaluated at once by a fast Fourier transform rather than one at a time.

In [13]:

num_radii = spatial_dims[0] // 2

radius = jnp.linspace(0.5 / spatial_dims[0], 1.0, num_radii)[:, None]
angle = jnp.linspace(0.0, 2.0 * jnp.pi, num_angles, endpoint=False)[None, :]
polar_grid = jnp.stack(
    [
        (radius * jnp.sin(angle) + 1.0) * (spatial_dims[0] - 1) / 2,
        (radius * jnp.cos(angle) + 1.0) * (spatial_dims[1] - 1) / 2,
    ]
)


def to_polar(image: Array) -> Array:
    """Resample an image onto the polar grid, so that a rotation becomes a shift."""
    return jnp.stack(
        [
            jax.scipy.ndimage.map_coordinates(
                image[..., c], polar_grid, order=1, mode="constant"
            )
            for c in range(image.shape[-1])
        ],
        axis=-1,
    )


target_polar = to_polar(y)
target_fft = jnp.conj(jnp.fft.rfft(target_polar, axis=-2))
target_square = jnp.sum(jnp.square(target_polar), axis=-2, keepdims=True)


def rotation_invariant_loss(rgba: Array) -> Array:
    """Squared error to the target, minimized over every rotation of the result."""
    polar = to_polar(rgba)
    spectrum = jnp.fft.rfft(polar, axis=-2)
    correlation = jnp.fft.irfft(spectrum * target_fft, n=num_angles, axis=-2)
    square = jnp.sum(jnp.square(polar), axis=-2, keepdims=True)

    per_rotation = jnp.mean(square + target_square - 2.0 * correlation, axis=(0, 2))
    return jnp.min(per_rotation)


def overflow_loss(trajectory: Array) -> Array:
    """Penalize values outside [-2, 2], over the whole rollout."""
    overflow = trajectory - jnp.clip(trajectory, -2.0, 2.0)
    return jnp.sum(jnp.square(overflow)) / batch_size


def difference_loss(history: Array) -> Array:
    """Penalize how much the state moves from one step to the next."""
    return jnp.sum(jnp.mean(jnp.abs(jnp.diff(history, axis=1)), axis=(0, 2, 3, 4)))

num_radii = spatial_dims[0] // 2 radius = jnp.linspace(0.5 / spatial_dims[0], 1.0, num_radii)[:, None] angle = jnp.linspace(0.0, 2.0 * jnp.pi, num_angles, endpoint=False)[None, :] polar_grid = jnp.stack( [ (radius * jnp.sin(angle) + 1.0) * (spatial_dims[0] - 1) / 2, (radius * jnp.cos(angle) + 1.0) * (spatial_dims[1] - 1) / 2, ] ) def to_polar(image: Array) -> Array: """Resample an image onto the polar grid, so that a rotation becomes a shift.""" return jnp.stack( [ jax.scipy.ndimage.map_coordinates( image[..., c], polar_grid, order=1, mode="constant" ) for c in range(image.shape[-1]) ], axis=-1, ) target_polar = to_polar(y) target_fft = jnp.conj(jnp.fft.rfft(target_polar, axis=-2)) target_square = jnp.sum(jnp.square(target_polar), axis=-2, keepdims=True) def rotation_invariant_loss(rgba: Array) -> Array: """Squared error to the target, minimized over every rotation of the result.""" polar = to_polar(rgba) spectrum = jnp.fft.rfft(polar, axis=-2) correlation = jnp.fft.irfft(spectrum * target_fft, n=num_angles, axis=-2) square = jnp.sum(jnp.square(polar), axis=-2, keepdims=True) per_rotation = jnp.mean(square + target_square - 2.0 * correlation, axis=(0, 2)) return jnp.min(per_rotation) def overflow_loss(trajectory: Array) -> Array: """Penalize values outside [-2, 2], over the whole rollout.""" overflow = trajectory - jnp.clip(trajectory, -2.0, 2.0) return jnp.sum(jnp.square(overflow)) / batch_size def difference_loss(history: Array) -> Array: """Penalize how much the state moves from one step to the next.""" return jnp.sum(jnp.mean(jnp.abs(jnp.diff(history, axis=1)), axis=(0, 2, 3, 4)))

Both regularizers are accumulated over the entire rollout rather than read off its end. A state that runs away in the middle is invisible to the endpoint, and is what makes training diverge.

In [14]:

def loss_fn(cs, state, key):
    """Loss function."""
    state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
    initial_state = state
    _, states = 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, so the pattern has to be right over a span of time
    idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps)
    state = states[jnp.arange(batch_size), idx]

    loss = jnp.mean(jax.vmap(rotation_invariant_loss)(state[..., -5:]))
    loss += overflow_weight * overflow_loss(states)
    loss += diff_weight * difference_loss(
        jnp.concatenate([initial_state[:, None], states], axis=1)
    )

    return loss, state

def loss_fn(cs, state, key): """Loss function.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) initial_state = state _, states = 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, so the pattern has to be right over a span of time idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps) state = states[jnp.arange(batch_size), idx] loss = jnp.mean(jax.vmap(rotation_invariant_loss)(state[..., -5:])) loss += overflow_weight * overflow_loss(states) loss += diff_weight * difference_loss( jnp.concatenate([initial_state[:, None], states], axis=1) ) return loss, state

Train step¶

A sample whose cells have all died carries no gradient, and a pool that fills up with them cannot recover. Any that die are replaced by a fresh seed.

In [15]:

@nnx.jit
def train_step(cs, optimizer, pool, key):
    """Train step."""
    # Sample from pool
    sample_key, loss_key = jax.random.split(key)
    pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False)
    current_state = batch["state"]

    # Replace the first sample, and any that died, with a fresh seed
    new_state = sample_state()
    current_state = current_state.at[0].set(new_state)
    dead = jnp.sum(current_state[..., -1:], axis=(1, 2, 3)) < 1e-6
    current_state = jnp.where(dead[:, None, None, None], new_state, current_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 from pool sample_key, loss_key = jax.random.split(key) pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False) current_state = batch["state"] # Replace the first sample, and any that died, with a fresh seed new_state = sample_state() current_state = current_state.at[0].set(new_state) dead = jnp.sum(current_state[..., -1:], axis=(1, 2, 3)) < 1e-6 current_state = jnp.where(dead[:, None, None, None], new_state, current_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 [16]:

print_interval = 200

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:>5}/{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")

print_interval = 200 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:>5}/{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")

W0904 23:40:09.439282  556847 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.

Step     0/16384 |    4.2s | Loss 3.213e+01

Step   200/16384 |   15.3s | Loss 1.353e+01

Step   400/16384 |   26.4s | Loss 1.078e+01

Step   600/16384 |   37.5s | Loss 1.038e+01

Step   800/16384 |   48.6s | Loss 1.024e+01

Step  1000/16384 |   59.7s | Loss 1.011e+01

Step  1200/16384 |   70.8s | Loss 9.829e+00

Step  1400/16384 |   81.9s | Loss 8.161e+00

Step  1600/16384 |   93.0s | Loss 6.852e+00

Step  1800/16384 |  104.1s | Loss 6.033e+00

Step  2000/16384 |  115.2s | Loss 5.322e+00

Step  2200/16384 |  126.3s | Loss 4.738e+00

Step  2400/16384 |  137.4s | Loss 4.388e+00

Step  2600/16384 |  148.5s | Loss 4.073e+00

Step  2800/16384 |  159.6s | Loss 3.786e+00

Step  3000/16384 |  170.8s | Loss 3.629e+00

Step  3200/16384 |  181.9s | Loss 3.296e+00

Step  3400/16384 |  193.0s | Loss 3.237e+00

Step  3600/16384 |  204.1s | Loss 3.034e+00

Step  3800/16384 |  215.2s | Loss 2.884e+00

Step  4000/16384 |  226.2s | Loss 2.817e+00

Step  4200/16384 |  237.3s | Loss 2.676e+00

Step  4400/16384 |  248.4s | Loss 2.692e+00

Step  4600/16384 |  259.5s | Loss 2.536e+00

Step  4800/16384 |  270.5s | Loss 2.512e+00

Step  5000/16384 |  281.6s | Loss 2.525e+00

Step  5200/16384 |  292.7s | Loss 2.464e+00

Step  5400/16384 |  303.7s | Loss 2.348e+00

Step  5600/16384 |  314.8s | Loss 2.366e+00

Step  5800/16384 |  325.9s | Loss 2.271e+00

Step  6000/16384 |  337.0s | Loss 2.314e+00

Step  6200/16384 |  348.0s | Loss 2.273e+00

Step  6400/16384 |  359.1s | Loss 2.269e+00

Step  6600/16384 |  370.2s | Loss 2.222e+00

Step  6800/16384 |  381.3s | Loss 2.251e+00

Step  7000/16384 |  392.4s | Loss 2.118e+00

Step  7200/16384 |  403.4s | Loss 2.099e+00

Step  7400/16384 |  414.5s | Loss 2.072e+00

Step  7600/16384 |  425.6s | Loss 2.014e+00

Step  7800/16384 |  436.7s | Loss 2.053e+00

Step  8000/16384 |  447.8s | Loss 2.044e+00

Step  8200/16384 |  459.0s | Loss 2.016e+00

Step  8400/16384 |  470.1s | Loss 1.965e+00

Step  8600/16384 |  481.2s | Loss 1.957e+00

Step  8800/16384 |  492.3s | Loss 1.941e+00

Step  9000/16384 |  503.5s | Loss 1.912e+00

Step  9200/16384 |  514.6s | Loss 1.886e+00

Step  9400/16384 |  525.7s | Loss 1.857e+00

Step  9600/16384 |  536.8s | Loss 1.856e+00

Step  9800/16384 |  548.0s | Loss 1.863e+00

Step 10000/16384 |  559.1s | Loss 1.821e+00

Step 10200/16384 |  570.3s | Loss 1.839e+00

Step 10400/16384 |  581.4s | Loss 1.812e+00

Step 10600/16384 |  592.6s | Loss 1.792e+00

Step 10800/16384 |  603.8s | Loss 1.819e+00

Step 11000/16384 |  614.9s | Loss 1.811e+00

Step 11200/16384 |  626.1s | Loss 1.788e+00

Step 11400/16384 |  637.2s | Loss 1.803e+00

Step 11600/16384 |  648.4s | Loss 1.786e+00

Step 11800/16384 |  659.5s | Loss 1.777e+00

Step 12000/16384 |  670.6s | Loss 1.780e+00

Step 12200/16384 |  681.7s | Loss 1.783e+00

Step 12400/16384 |  692.8s | Loss 1.790e+00

Step 12600/16384 |  703.9s | Loss 1.781e+00

Step 12800/16384 |  715.0s | Loss 1.785e+00

Step 13000/16384 |  726.1s | Loss 1.769e+00

Step 13200/16384 |  737.2s | Loss 1.769e+00

Step 13400/16384 |  748.3s | Loss 1.756e+00

Step 13600/16384 |  759.4s | Loss 1.757e+00

Step 13800/16384 |  770.4s | Loss 1.747e+00

Step 14000/16384 |  781.5s | Loss 1.769e+00

Step 14200/16384 |  792.6s | Loss 1.760e+00

Step 14400/16384 |  803.6s | Loss 1.754e+00

Step 14600/16384 |  814.7s | Loss 1.748e+00

Step 14800/16384 |  825.7s | Loss 1.755e+00

Step 15000/16384 |  836.8s | Loss 1.747e+00

Step 15200/16384 |  847.8s | Loss 1.775e+00

Step 15400/16384 |  858.9s | Loss 1.768e+00

Step 15600/16384 |  869.9s | Loss 1.776e+00

Step 15800/16384 |  881.0s | Loss 1.759e+00

Step 16000/16384 |  892.0s | Loss 1.738e+00

Step 16200/16384 |  903.0s | Loss 1.754e+00

Step 16383/16384 |  913.1s | Loss 1.762e+00
✨ Trained for 16384 steps in 913s

Run¶

Nothing in the rule or the seed picks an orientation, so each run settles on its own. The pattern is the same; the angle is not.

In [17]:

num_examples = 6

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=num_steps, return_states=True),
        in_axes=(state_axes, 0),
    )
)(cs, state_init)

num_examples = 6 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=num_steps, return_states=True), in_axes=(state_axes, 0), ) )(cs, state_init)

Visualize¶

In [18]:

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

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

In [19]:

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))