Neural Particle Automata ¶

A growing neural cellular automaton lives on a lattice: every cell has an address, and its neighbours are the cells next to it forever. Here the cells are particles. Each carries a position as well as a state, and the rule moves both, so who is next to whom is something the automaton decides as it goes.

That one change costs the convolution. With no lattice there is no stencil, so the neighbourhood is gathered instead by smoothed particle hydrodynamics: sums over whatever particles lie within a radius, weighted by a kernel that falls smoothly to zero at the edge. The smoothness is what keeps it differentiable while neighbours come and go.

Everything else is 40 - Growing NCA unchanged --- the same target, the same channel count, the same stochastic half-updates, the same pool.

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 Particles, SPHPerceive
from cax.nn.pool import Pool
from cax.utils import (
    clip_and_uint8,
    get_emoji_array,
    render_states,
    rgba_to_rgb,
    safe_norm,
)

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 Particles, SPHPerceive from cax.nn.pool import Pool from cax.utils import ( clip_and_uint8, get_emoji_array, render_states, rgba_to_rgb, safe_norm, )

Configuration¶

fused chooses how the perception is computed. Both routes return the same numbers, so it changes nothing about what the automaton learns --- only what it costs. The array route builds an array with an entry per pair; the fused route accumulates the same sums a tile at a time inside a Pallas kernel and never builds one, which is faster on a GPU and is the only route that reaches cloud sizes the first cannot fit. It is GPU-only, so the default follows the backend and this notebook runs anywhere.

In [4]:

seed = 0

# The perception is the expensive part, and on a GPU it can be computed by fused
# kernels instead of array operations. Same numbers, less memory traffic.
fused = jax.default_backend() == "gpu"

num_particles = 1024
channel_size = 16
hidden_size = 128
support_radius = 0.2
seed_radius = 0.1
cell_dropout_rate = 0.5
displacement_scale = 0.5

num_steps = 96
min_steps = 64
num_eval_steps = 256
pool_size = 512
batch_size = 8
learning_rate = 5e-4
overflow_weight = 100.0

num_train_steps = 4_000
grid_size = 96
density_weight = 1.0
colour_weight = 5.0

emoji = "🦎"
size = 64

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

seed = 0 # The perception is the expensive part, and on a GPU it can be computed by fused # kernels instead of array operations. Same numbers, less memory traffic. fused = jax.default_backend() == "gpu" num_particles = 1024 channel_size = 16 hidden_size = 128 support_radius = 0.2 seed_radius = 0.1 cell_dropout_rate = 0.5 displacement_scale = 0.5 num_steps = 96 min_steps = 64 num_eval_steps = 256 pool_size = 512 batch_size = 8 learning_rate = 5e-4 overflow_weight = 100.0 num_train_steps = 4_000 grid_size = 96 density_weight = 1.0 colour_weight = 5.0 emoji = "🦎" size = 64 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

The target is a picture, but the automaton is a cloud of points, so the two are compared by drawing the cloud: each particle is splatted onto a grid, and the resulting density and colour are matched against the emoji.

Comparing per particle instead --- asking each one to match the target where it happens to stand --- would let a particle sitting on empty background simply learn to be transparent. Nothing would ever make the cloud take the shape.

In [5]:

y = get_emoji_array(emoji, size, 0)

target_density = jax.image.resize(y[..., 3:], (grid_size, grid_size, 1), "linear")
target_density = target_density / jnp.mean(target_density)
target_colour = jax.image.resize(y[..., :3], (grid_size, grid_size, 3), "linear")

mediapy.show_image(rgba_to_rgb(y))

y = get_emoji_array(emoji, size, 0) target_density = jax.image.resize(y[..., 3:], (grid_size, grid_size, 1), "linear") target_density = target_density / jnp.mean(target_density) target_colour = jax.image.resize(y[..., :3], (grid_size, grid_size, 3), "linear") mediapy.show_image(rgba_to_rgb(y))

Drawing the cloud¶

Twice, for two purposes. draw rasterizes onto a coarse grid for the loss: cheap, and smooth enough that a gradient reaches each particle's position. render draws a proper picture.

In [6]:

def blur(image: Array) -> Array:
    """Blur with a separable binomial kernel."""
    weights = jnp.array([1.0, 4.0, 6.0, 4.0, 1.0]) / 16.0
    for axis in (0, 1):
        pad = [(0, 0)] * image.ndim
        pad[axis] = (2, 2)
        padded = jnp.pad(image, pad, mode="wrap")
        taps = jnp.stack(
            [
                jax.lax.dynamic_slice_in_dim(padded, i, image.shape[axis], axis)
                for i in range(weights.size)
            ]
        )
        image = jnp.tensordot(weights, taps, axes=(0, 0))
    return image


def splat(position: Array, value: Array) -> Array:
    """Rasterize particles onto a grid, each spread over the pixels it falls between."""
    pixel = position * grid_size
    corner = jnp.floor(pixel).astype(jnp.int32)
    frac = pixel - corner

    image = jnp.zeros((grid_size, grid_size, value.shape[-1]))
    for row_offset in (0, 1):
        for column_offset in (0, 1):
            weight = (frac[:, 0] if row_offset else 1.0 - frac[:, 0]) * (
                frac[:, 1] if column_offset else 1.0 - frac[:, 1]
            )
            row = (corner[:, 0] + row_offset) % grid_size
            column = (corner[:, 1] + column_offset) % grid_size
            image = image.at[row, column].add(weight[:, None] * value)
    return blur(image)


def render(particles: Particles, *, image_size: int = 384, sigma: float = 2.5) -> Array:
    """Draw the cloud for viewing: a gaussian per particle, composited over white.

    Separate from `draw` above, which the loss uses. The loss wants a cheap grid and is
    content with colour left premultiplied by density; a picture wants neither. Shown
    premultiplied, every region thinner than average comes out darkened, which reads as
    a pale and sparse cloud rather than the one that is actually there.
    """
    pixel = particles.position * image_size
    corner = jnp.floor(pixel).astype(jnp.int32)
    offset = jnp.arange(-int(3 * sigma), int(3 * sigma) + 1)
    rows, columns = jnp.meshgrid(offset, offset, indexing="ij")

    distance_row = rows[None] - (pixel - corner)[:, 0, None, None]
    distance_column = columns[None] - (pixel - corner)[:, 1, None, None]
    stamp = jnp.exp(-(distance_row**2 + distance_column**2) / (2.0 * sigma**2))
    stamp = stamp / (1e-8 + jnp.sum(stamp, axis=(1, 2), keepdims=True))

    row = (corner[:, 0, None, None] + rows[None]) % image_size
    column = (corner[:, 1, None, None] + columns[None]) % image_size
    colour = jnp.clip(particles.state[..., -4:-1], 0.0, 1.0)

    weight = (
        jnp.zeros((image_size, image_size, 1)).at[row, column].add(stamp[..., None])
    )
    painted = (
        jnp.zeros((image_size, image_size, 3))
        .at[row, column]
        .add(stamp[..., None] * colour[:, None, None, :])
    )

    alpha = jnp.clip(weight / jnp.mean(weight, where=weight > 0.0), 0.0, 1.0)
    return jnp.clip(
        (1.0 - alpha) + alpha * painted / jnp.maximum(weight, 1e-8), 0.0, 1.0
    )


def draw(particles: Particles) -> tuple[Array, Array]:
    """Draw the cloud as a density and a colour picture.

    Both are divided by the same constant, so they stay on one scale. Colour is left
    premultiplied by density rather than divided by it: dividing gives a meaningless
    colour wherever the cloud is empty.
    """
    density = splat(particles.position, jnp.ones_like(particles.state[:, :1]))
    colour = splat(particles.position, jnp.clip(particles.state[..., -4:-1], 0.0, 1.0))

    scale = jnp.mean(density)
    return density / scale, colour / scale

def blur(image: Array) -> Array: """Blur with a separable binomial kernel.""" weights = jnp.array([1.0, 4.0, 6.0, 4.0, 1.0]) / 16.0 for axis in (0, 1): pad = [(0, 0)] * image.ndim pad[axis] = (2, 2) padded = jnp.pad(image, pad, mode="wrap") taps = jnp.stack( [ jax.lax.dynamic_slice_in_dim(padded, i, image.shape[axis], axis) for i in range(weights.size) ] ) image = jnp.tensordot(weights, taps, axes=(0, 0)) return image def splat(position: Array, value: Array) -> Array: """Rasterize particles onto a grid, each spread over the pixels it falls between.""" pixel = position * grid_size corner = jnp.floor(pixel).astype(jnp.int32) frac = pixel - corner image = jnp.zeros((grid_size, grid_size, value.shape[-1])) for row_offset in (0, 1): for column_offset in (0, 1): weight = (frac[:, 0] if row_offset else 1.0 - frac[:, 0]) * ( frac[:, 1] if column_offset else 1.0 - frac[:, 1] ) row = (corner[:, 0] + row_offset) % grid_size column = (corner[:, 1] + column_offset) % grid_size image = image.at[row, column].add(weight[:, None] * value) return blur(image) def render(particles: Particles, *, image_size: int = 384, sigma: float = 2.5) -> Array: """Draw the cloud for viewing: a gaussian per particle, composited over white. Separate from `draw` above, which the loss uses. The loss wants a cheap grid and is content with colour left premultiplied by density; a picture wants neither. Shown premultiplied, every region thinner than average comes out darkened, which reads as a pale and sparse cloud rather than the one that is actually there. """ pixel = particles.position * image_size corner = jnp.floor(pixel).astype(jnp.int32) offset = jnp.arange(-int(3 * sigma), int(3 * sigma) + 1) rows, columns = jnp.meshgrid(offset, offset, indexing="ij") distance_row = rows[None] - (pixel - corner)[:, 0, None, None] distance_column = columns[None] - (pixel - corner)[:, 1, None, None] stamp = jnp.exp(-(distance_row**2 + distance_column**2) / (2.0 * sigma**2)) stamp = stamp / (1e-8 + jnp.sum(stamp, axis=(1, 2), keepdims=True)) row = (corner[:, 0, None, None] + rows[None]) % image_size column = (corner[:, 1, None, None] + columns[None]) % image_size colour = jnp.clip(particles.state[..., -4:-1], 0.0, 1.0) weight = ( jnp.zeros((image_size, image_size, 1)).at[row, column].add(stamp[..., None]) ) painted = ( jnp.zeros((image_size, image_size, 3)) .at[row, column] .add(stamp[..., None] * colour[:, None, None, :]) ) alpha = jnp.clip(weight / jnp.mean(weight, where=weight > 0.0), 0.0, 1.0) return jnp.clip( (1.0 - alpha) + alpha * painted / jnp.maximum(weight, 1e-8), 0.0, 1.0 ) def draw(particles: Particles) -> tuple[Array, Array]: """Draw the cloud as a density and a colour picture. Both are divided by the same constant, so they stay on one scale. Colour is left premultiplied by density rather than divided by it: dividing gives a meaningless colour wherever the cloud is empty. """ density = splat(particles.position, jnp.ones_like(particles.state[:, :1])) colour = splat(particles.position, jnp.clip(particles.state[..., -4:-1], 0.0, 1.0)) scale = jnp.mean(density) return density / scale, colour / scale

Instantiate system¶

In [7]:

class ParticleUpdate(nnx.Module):
    """Predict a move and a state change for every particle, from its perception."""

    def __init__(self, *, perception_size: int, rngs: nnx.Rngs):
        """Initialize particle update.

        Args:
                perception_size: Width of the perception.
                rngs: rng key.

        """
        self.layer_1 = nnx.Linear(perception_size, hidden_size, rngs=rngs)
        self.layer_2 = nnx.Linear(
            hidden_size,
            channel_size + 2,
            use_bias=False,
            kernel_init=nnx.initializers.zeros_init(),
            rngs=rngs,
        )
        self.rngs = rngs

    def __call__(
        self, state: Particles, perception: Array, input: Array | None = None
    ) -> Particles:
        """Process the current state, perception, and input to produce a new state.

        Args:
                state: Current particles.
                perception: Current perception.
                input: Optional input.

        Returns:
                Next particles.

        """
        delta = self.layer_2(nnx.relu(self.layer_1(perception)))
        move, change = delta[:, :2], delta[:, 2:]

        # Bound the move to a fraction of the support radius. Unbounded, a particle can
        # leave its neighbourhood in a single step, and everything it perceived is gone.
        # safe_norm, because the last layer starts at zero and an ordinary norm has no
        # derivative at the origin.
        length = safe_norm(move, axis=-1, keepdims=True)
        move = displacement_scale * support_radius * move / (1.0 + length)

        alive = jax.random.uniform(self.rngs.dropout(), (state.position.shape[0], 1))
        mask = (alive < cell_dropout_rate).astype(jnp.float32)

        return Particles(
            position=(state.position + mask * move) % 1.0,
            state=state.state + mask * change,
        )

class ParticleUpdate(nnx.Module): """Predict a move and a state change for every particle, from its perception.""" def __init__(self, *, perception_size: int, rngs: nnx.Rngs): """Initialize particle update. Args: perception_size: Width of the perception. rngs: rng key. """ self.layer_1 = nnx.Linear(perception_size, hidden_size, rngs=rngs) self.layer_2 = nnx.Linear( hidden_size, channel_size + 2, use_bias=False, kernel_init=nnx.initializers.zeros_init(), rngs=rngs, ) self.rngs = rngs def __call__( self, state: Particles, perception: Array, input: Array | None = None ) -> Particles: """Process the current state, perception, and input to produce a new state. Args: state: Current particles. perception: Current perception. input: Optional input. Returns: Next particles. """ delta = self.layer_2(nnx.relu(self.layer_1(perception))) move, change = delta[:, :2], delta[:, 2:] # Bound the move to a fraction of the support radius. Unbounded, a particle can # leave its neighbourhood in a single step, and everything it perceived is gone. # safe_norm, because the last layer starts at zero and an ordinary norm has no # derivative at the origin. length = safe_norm(move, axis=-1, keepdims=True) move = displacement_scale * support_radius * move / (1.0 + length) alive = jax.random.uniform(self.rngs.dropout(), (state.position.shape[0], 1)) mask = (alive < cell_dropout_rate).astype(jnp.float32) return Particles( position=(state.position + mask * move) % 1.0, state=state.state + mask * change, )

In [8]:

class NeuralParticleAutomata(ComplexSystem):
    """Neural Particle Automata class."""

    remat = True

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize Neural Particle Automata.

        Args:
                rngs: rng key.

        """
        self.perceive = SPHPerceive(
            support_radius=support_radius,
            mass=1.0 / num_particles,
            period=1.0,
            fused=fused,
        )
        self.update = ParticleUpdate(
            perception_size=SPHPerceive.perception_size(
                channel_size=channel_size, num_spatial_dims=2
            ),
            rngs=rngs,
        )

    def _step(self, state: Particles, input: Array | None = None) -> Particles:
        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(render(state))

class NeuralParticleAutomata(ComplexSystem): """Neural Particle Automata class.""" remat = True def __init__(self, *, rngs: nnx.Rngs): """Initialize Neural Particle Automata. Args: rngs: rng key. """ self.perceive = SPHPerceive( support_radius=support_radius, mass=1.0 / num_particles, period=1.0, fused=fused, ) self.update = ParticleUpdate( perception_size=SPHPerceive.perception_size( channel_size=channel_size, num_spatial_dims=2 ), rngs=rngs, ) def _step(self, state: Particles, input: Array | None = None) -> Particles: 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(render(state))

In [9]:

cs = NeuralParticleAutomata(rngs=rngs)

cs = NeuralParticleAutomata(rngs=rngs)

In [10]:

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

Sample initial state¶

Every particle starts inside a small disc, and every state starts at zero. The cloud has to carry itself out into the shape, which is the whole of what makes this a particle model rather than a lattice one wearing different clothes.

It is worth asking where the first non-zero update comes from, since a blank state through a zero-initialized output layer gives nothing back. It comes from the density gradient: the states are uniform but the positions are not, so it is the one thing a featureless cloud still has to say.

In [11]:

def sample_state(key: Array) -> Particles:
    """Sample particles inside a disc, with no state at all."""
    angle_key, radius_key = jax.random.split(key)

    angle = jax.random.uniform(angle_key, (num_particles,), maxval=2.0 * jnp.pi)
    # The square root spreads the particles evenly over the area rather than crowding
    # them into the middle.
    radius = seed_radius * jnp.sqrt(jax.random.uniform(radius_key, (num_particles,)))
    position = 0.5 + jnp.stack(
        [radius * jnp.cos(angle), radius * jnp.sin(angle)], axis=-1
    )

    return Particles(position=position, state=jnp.zeros((num_particles, channel_size)))

def sample_state(key: Array) -> Particles: """Sample particles inside a disc, with no state at all.""" angle_key, radius_key = jax.random.split(key) angle = jax.random.uniform(angle_key, (num_particles,), maxval=2.0 * jnp.pi) # The square root spreads the particles evenly over the area rather than crowding # them into the middle. radius = seed_radius * jnp.sqrt(jax.random.uniform(radius_key, (num_particles,))) position = 0.5 + jnp.stack( [radius * jnp.cos(angle), radius * jnp.sin(angle)], axis=-1 ) return Particles(position=position, state=jnp.zeros((num_particles, channel_size)))

Train¶

Pool¶

In [12]:

key, subkey = jax.random.split(key)
particles = jax.vmap(sample_state)(jax.random.split(subkey, pool_size))

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

key, subkey = jax.random.split(key) particles = jax.vmap(sample_state)(jax.random.split(subkey, pool_size)) pool = Pool.create({"particles": particles})

Optimizer¶

In [13]:

lr_sched = optax.piecewise_constant_schedule(
    init_value=learning_rate, boundaries_and_scales={2_000: 0.3, 4_000: 0.3}
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=lr_sched, weight_decay=0.0),
)

optimizer = nnx.Optimizer(cs, optimizer, wrt=nnx.Param)

lr_sched = optax.piecewise_constant_schedule( init_value=learning_rate, boundaries_and_scales={2_000: 0.3, 4_000: 0.3} ) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adamw(learning_rate=lr_sched, weight_decay=0.0), ) optimizer = nnx.Optimizer(cs, optimizer, wrt=nnx.Param)

Loss¶

Two terms match the drawn cloud to the target, and a third keeps the state in range. Nothing else bounds it, and an unbounded state is what diverges.

In [14]:

def render_loss(particles: Particles) -> Array:
    """Squared error between the drawn cloud and the target."""
    density, colour = draw(particles)

    loss = density_weight * jnp.mean(jnp.square(density - target_density))
    return loss + colour_weight * jnp.mean(
        jnp.square(colour - target_colour * target_density)
    )


def overflow_loss(state: Array) -> Array:
    """Penalize values outside [-1, 1], over the whole rollout."""
    overflow = state - jnp.clip(state, -1.0, 1.0)
    return jnp.mean(jnp.square(overflow))

def render_loss(particles: Particles) -> Array: """Squared error between the drawn cloud and the target.""" density, colour = draw(particles) loss = density_weight * jnp.mean(jnp.square(density - target_density)) return loss + colour_weight * jnp.mean( jnp.square(colour - target_colour * target_density) ) def overflow_loss(state: Array) -> Array: """Penalize values outside [-1, 1], over the whole rollout.""" overflow = state - jnp.clip(state, -1.0, 1.0) return jnp.mean(jnp.square(overflow))

In [15]:

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

    # Sample a random step, so the shape has to be right over a span of time
    idx = jax.random.randint(key, (batch_size,), min_steps, num_steps)
    particles = jax.tree.map(
        lambda steps: steps[jnp.arange(batch_size), idx], trajectory
    )

    loss = jnp.mean(jax.vmap(render_loss)(particles))
    loss += overflow_weight * overflow_loss(trajectory.state)

    return loss, particles

def loss_fn(cs, particles, key): """Loss function.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) _, trajectory = nnx.split_rngs(splits=batch_size)( nnx.vmap( lambda cs, particles: cs( particles, num_steps=num_steps, return_states=True ), in_axes=(state_axes, 0), ) )(cs, particles) # Sample a random step, so the shape has to be right over a span of time idx = jax.random.randint(key, (batch_size,), min_steps, num_steps) particles = jax.tree.map( lambda steps: steps[jnp.arange(batch_size), idx], trajectory ) loss = jnp.mean(jax.vmap(render_loss)(particles)) loss += overflow_weight * overflow_loss(trajectory.state) return loss, particles

Train step¶

In [16]:

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

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

    # Replace the first sample with a fresh seed
    new_particles = sample_state(seed_key)
    particles = jax.tree.map(
        lambda batch, seed: batch.at[0].set(seed), particles, new_particles
    )

    (loss, particles), grad = nnx.value_and_grad(loss_fn, has_aux=True)(
        cs, particles, loss_key
    )
    optimizer.update(cs, grad)

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

@nnx.jit def train_step(cs, optimizer, pool, key): """Train step.""" sample_key, seed_key, loss_key = jax.random.split(key, 3) # Sample from pool pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False) particles = batch["particles"] # Replace the first sample with a fresh seed new_particles = sample_state(seed_key) particles = jax.tree.map( lambda batch, seed: batch.at[0].set(seed), particles, new_particles ) (loss, particles), grad = nnx.value_and_grad(loss_fn, has_aux=True)( cs, particles, loss_key ) optimizer.update(cs, grad) pool = pool.update(pool_idx, {"particles": particles}) return loss, pool

Main loop¶

In [17]:

print_interval = 100

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

print_interval = 100 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/4000 |    5.1s | Loss 2.744e+01

Step  100/4000 |   10.3s | Loss 3.028e+00

Step  200/4000 |   15.4s | Loss 2.357e+00

Step  300/4000 |   20.6s | Loss 2.326e+00

Step  400/4000 |   25.8s | Loss 2.180e+00

Step  500/4000 |   31.0s | Loss 2.008e+00

Step  600/4000 |   36.2s | Loss 1.816e+00

Step  700/4000 |   41.4s | Loss 1.726e+00

Step  800/4000 |   46.6s | Loss 1.596e+00

Step  900/4000 |   51.8s | Loss 1.404e+00

Step 1000/4000 |   57.0s | Loss 1.402e+00

Step 1100/4000 |   62.1s | Loss 1.263e+00

Step 1200/4000 |   67.3s | Loss 1.167e+00

Step 1300/4000 |   72.5s | Loss 1.177e+00

Step 1400/4000 |   77.7s | Loss 1.022e+00

Step 1500/4000 |   82.9s | Loss 9.487e-01

Step 1600/4000 |   88.1s | Loss 9.009e-01

Step 1700/4000 |   93.3s | Loss 8.318e-01

Step 1800/4000 |   98.5s | Loss 7.848e-01

Step 1900/4000 |  103.7s | Loss 7.377e-01

Step 2000/4000 |  108.9s | Loss 7.082e-01

Step 2100/4000 |  114.1s | Loss 6.044e-01

Step 2200/4000 |  119.3s | Loss 5.399e-01

Step 2300/4000 |  124.5s | Loss 5.126e-01

Step 2400/4000 |  129.6s | Loss 4.768e-01

Step 2500/4000 |  134.8s | Loss 4.615e-01

Step 2600/4000 |  140.0s | Loss 4.547e-01

Step 2700/4000 |  145.2s | Loss 4.472e-01

Step 2800/4000 |  150.4s | Loss 3.958e-01

Step 2900/4000 |  155.6s | Loss 3.561e-01

Step 3000/4000 |  160.8s | Loss 3.717e-01

Step 3100/4000 |  166.0s | Loss 3.436e-01

Step 3200/4000 |  171.2s | Loss 3.617e-01

Step 3300/4000 |  176.4s | Loss 3.385e-01

Step 3400/4000 |  181.6s | Loss 3.315e-01

Step 3500/4000 |  186.8s | Loss 3.291e-01

Step 3600/4000 |  191.9s | Loss 3.282e-01

Step 3700/4000 |  197.1s | Loss 3.178e-01

Step 3800/4000 |  202.3s | Loss 2.993e-01

Step 3900/4000 |  207.5s | Loss 2.975e-01

Step 3999/4000 |  212.7s | Loss 3.061e-01
✨ Trained for 4000 steps in 213s

Run¶

Rolled far past the horizon it was trained on. The loss only ever looked at steps [min_steps, num_steps), so nothing asked the rule to still hold a gecko at step num_eval_steps --- it either learned a shape that persists, or it learned one that arrives on schedule and then falls apart.

In [18]:

key, subkey = jax.random.split(key)
particles_init = sample_state(subkey)

particles_final, particles = cs(
    particles_init, num_steps=num_eval_steps, return_states=True
)

key, subkey = jax.random.split(key) particles_init = sample_state(subkey) particles_final, particles = cs( particles_init, num_steps=num_eval_steps, return_states=True )

Visualize¶

In [19]:

mediapy.show_image(cs.render(particles_final))

mediapy.show_image(cs.render(particles_final))

In [20]:

particles = jax.tree.map(
    lambda first, rest: jnp.concatenate([first[None], rest]), particles_init, particles
)
frames = render_states(cs, particles)

mediapy.show_video(frames)

particles = jax.tree.map( lambda first, rest: jnp.concatenate([first[None], rest]), particles_init, particles ) frames = render_states(cs, particles) mediapy.show_video(frames)