Growing Neural Cellular Automata with Reinforcement Learning ¶

Growing NCA is usually trained by backpropagating through the whole developmental process, which costs time and memory in proportion to its length.

Here the system is treated as a policy $\pi_\theta$ mapping a state to the next one, and every state is scored by its negative distance to the target, $r(s) = -\lVert s_\text{RGBA} - y \rVert^2$. The objective is the discounted return of a rollout of length $H$, with everything beyond the horizon summarized by a learned value function $V$:

$$J(\theta) = \mathbb{E}\Big[\sum_{t=0}^{H-1} \gamma^{t} \, r(s_t) \; + \; \gamma^{H} V(s_H)\Big]$$

Taking $H$ well below the full number of developmental steps is what makes this cheap, since cost scales with $H$ rather than with the length of the process. The cellular automaton is differentiable, so $\nabla_\theta J$ flows analytically through the rollout and through $V(s_H)$, instead of being estimated from samples as a policy gradient would be.

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
horizon = 64
discount = 0.999
lambda_ = 0.95

critic_features = (32, 64, 128)
critic_hidden_size = 128
critic_learning_rate = 3e-4
critic_num_updates = 4
target_step_size = 0.05
bootstrap_warmup = 200

pool_size = 1_024
batch_size = 32
num_resets = 16
num_train_steps = 16_384
learning_rate = 3e-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 horizon = 64 discount = 0.999 lambda_ = 0.95 critic_features = (32, 64, 128) critic_hidden_size = 128 critic_learning_rate = 3e-4 critic_num_updates = 4 target_step_size = 0.05 bootstrap_warmup = 200 pool_size = 1_024 batch_size = 32 num_resets = 16 num_train_steps = 16_384 learning_rate = 3e-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

Instantiate critic¶

The critic estimates the return of a state, which is what lets the policy be trained without unrolling the whole developmental process.

Every reward is a negative distance, so no state can have positive value. Writing $V_\phi = -\operatorname{softplus}(\cdot)$ builds that into the critic and keeps the policy from chasing states where an unconstrained critic would extrapolate upwards. Bootstrap values come from a slowly moving copy $V_{\phi^-}$, updated as $\phi^- \leftarrow (1 - \tau)\,\phi^- + \tau\,\phi$, so the target the critic regresses toward does not move with every gradient step.

In [9]:

class Critic(nnx.Module):
    """Critic estimating the return of a state."""

    def __init__(self, *, features: tuple[int, ...], hidden_size: int, rngs: nnx.Rngs):
        """Initialize the critic.

        Args:
            features: Number of features of each strided convolution.
            hidden_size: Size of the hidden layer.
            rngs: rng key.

        """
        self.convs = nnx.List(
            [
                nnx.Conv(
                    in_features=in_features,
                    out_features=out_features,
                    kernel_size=(3, 3),
                    strides=(2, 2),
                    padding="SAME",
                    rngs=rngs,
                )
                for in_features, out_features in zip(
                    (channel_size, *features[:-1]), features, strict=True
                )
            ]
        )
        self.linear_1 = nnx.Linear(features[-1], hidden_size, rngs=rngs)
        self.linear_2 = nnx.Linear(hidden_size, 1, rngs=rngs)

    def __call__(self, state: Array) -> Array:
        """Estimate the value of a state."""
        x = state
        for conv in self.convs:
            x = jax.nn.relu(conv(x))

        # Pool over space so the head does not depend on the size of the grid
        x = jnp.mean(x, axis=(-3, -2))
        x = jax.nn.relu(self.linear_1(x))
        value = jnp.squeeze(self.linear_2(x), axis=-1)

        # Rewards are negative distances, so values are negative too
        return -jax.nn.softplus(value)

class Critic(nnx.Module): """Critic estimating the return of a state.""" def __init__(self, *, features: tuple[int, ...], hidden_size: int, rngs: nnx.Rngs): """Initialize the critic. Args: features: Number of features of each strided convolution. hidden_size: Size of the hidden layer. rngs: rng key. """ self.convs = nnx.List( [ nnx.Conv( in_features=in_features, out_features=out_features, kernel_size=(3, 3), strides=(2, 2), padding="SAME", rngs=rngs, ) for in_features, out_features in zip( (channel_size, *features[:-1]), features, strict=True ) ] ) self.linear_1 = nnx.Linear(features[-1], hidden_size, rngs=rngs) self.linear_2 = nnx.Linear(hidden_size, 1, rngs=rngs) def __call__(self, state: Array) -> Array: """Estimate the value of a state.""" x = state for conv in self.convs: x = jax.nn.relu(conv(x)) # Pool over space so the head does not depend on the size of the grid x = jnp.mean(x, axis=(-3, -2)) x = jax.nn.relu(self.linear_1(x)) value = jnp.squeeze(self.linear_2(x), axis=-1) # Rewards are negative distances, so values are negative too return -jax.nn.softplus(value)

In [10]:

critic = Critic(features=critic_features, hidden_size=critic_hidden_size, rngs=rngs)

# A slow-moving copy of the critic provides every bootstrap value, so that the
# critic never regresses toward its own moving predictions.
critic_target = nnx.clone(critic)

critic = Critic(features=critic_features, hidden_size=critic_hidden_size, rngs=rngs) # A slow-moving copy of the critic provides every bootstrap value, so that the # critic never regresses toward its own moving predictions. critic_target = nnx.clone(critic)

Sample initial state¶

In [11]:

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

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

lr_sched = optax.warmup_cosine_decay_schedule(
    init_value=0.01 * learning_rate,
    peak_value=learning_rate,
    warmup_steps=200,
    decay_steps=num_train_steps,
    end_value=0.02 * learning_rate,
)


def normalize_by_norm(eps: float = 1e-8) -> optax.GradientTransformation:
    """Normalize each gradient tensor by its own norm."""

    def init_fn(params):
        del params
        return optax.EmptyState()

    def update_fn(updates, state, params=None):
        del params
        return jax.tree.map(lambda g: g / (jnp.linalg.norm(g) + eps), updates), state

    return optax.GradientTransformation(init_fn, update_fn)


optimizer = optax.chain(
    optax.zero_nans(),
    normalize_by_norm(),
    optax.clip_by_global_norm(0.5),
    optax.adam(learning_rate=lr_sched),
)

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

critic_optimizer = nnx.Optimizer(
    critic,
    optax.chain(
        optax.clip_by_global_norm(1.0),
        optax.adam(learning_rate=critic_learning_rate),
    ),
    wrt=nnx.Param,
)

lr_sched = optax.warmup_cosine_decay_schedule( init_value=0.01 * learning_rate, peak_value=learning_rate, warmup_steps=200, decay_steps=num_train_steps, end_value=0.02 * learning_rate, ) def normalize_by_norm(eps: float = 1e-8) -> optax.GradientTransformation: """Normalize each gradient tensor by its own norm.""" def init_fn(params): del params return optax.EmptyState() def update_fn(updates, state, params=None): del params return jax.tree.map(lambda g: g / (jnp.linalg.norm(g) + eps), updates), state return optax.GradientTransformation(init_fn, update_fn) optimizer = optax.chain( optax.zero_nans(), normalize_by_norm(), optax.clip_by_global_norm(0.5), optax.adam(learning_rate=lr_sched), ) update_params = nnx.All(nnx.Param, nnx.PathContains("update")) optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params) critic_optimizer = nnx.Optimizer( critic, optax.chain( optax.clip_by_global_norm(1.0), optax.adam(learning_rate=critic_learning_rate), ), wrt=nnx.Param, )

Reward¶

In [14]:

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


def reward_fn(state):
    """Reward a state for being close to the target."""
    return -mse(state)

def mse(state): """Mean Squared Error.""" return jnp.mean(jnp.square(state[..., -4:] - y)) def reward_fn(state): """Reward a state for being close to the target.""" return -mse(state)

Return¶

The return of a rollout mixes the rewards it collected with the value of the state it ended in. The $\lambda$-return does this recursively, backwards from $G_H = V_{\phi^-}(s_H)$:

$$G_t = r_t + \gamma \big[(1 - \lambda) \, V_{\phi^-}(s_{t+1}) + \lambda \, G_{t+1}\big]$$

so $\lambda = 0$ trusts a single step and its bootstrap, $\lambda = 1$ trusts the whole rollout, and values in between interpolate.

In [15]:

def lambda_return(rewards, values):
    """Compute the lambda-return of a trajectory.

    Args:
        rewards: Rewards of shape `(horizon,)`, collected on entering each state.
        values: Estimated values of shape `(horizon + 1,)`, of the initial state
            through the final state of the rollout.

    Returns:
        The lambda-returns of shape `(horizon,)`, one for each state but the last.

    """

    def scan_fn(carry, x):
        reward, value = x
        carry = reward + discount * ((1.0 - lambda_) * value + lambda_ * carry)
        return carry, carry

    _, returns = jax.lax.scan(scan_fn, values[-1], (rewards, values[1:]), reverse=True)
    return returns

def lambda_return(rewards, values): """Compute the lambda-return of a trajectory. Args: rewards: Rewards of shape `(horizon,)`, collected on entering each state. values: Estimated values of shape `(horizon + 1,)`, of the initial state through the final state of the rollout. Returns: The lambda-returns of shape `(horizon,)`, one for each state but the last. """ def scan_fn(carry, x): reward, value = x carry = reward + discount * ((1.0 - lambda_) * value + lambda_ * carry) return carry, carry _, returns = jax.lax.scan(scan_fn, values[-1], (rewards, values[1:]), reverse=True) return returns

Loss¶

In [16]:

def rollout(cs, state):
    """Roll the system out for `horizon` steps and reward every state it visits."""
    state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
    _, states = nnx.split_rngs(splits=batch_size)(
        nnx.vmap(
            lambda cs, state: cs(state, num_steps=horizon, return_states=True),
            in_axes=(state_axes, 0),
        )
    )(cs, state)

    rewards = jax.vmap(jax.vmap(reward_fn))(states)
    return states, rewards


def actor_loss_fn(cs, critic_target, state, bootstrap_scale):
    """Compute the negative bootstrapped return of a rollout.

    The gradient flows analytically through the rollout, and through the value of
    the state it ends in, which is what carries credit past the horizon.
    `bootstrap_scale` ramps up early in training so that the policy is not steered
    by the critic before the critic is accurate.
    """
    states, rewards = rollout(cs, state)

    returns = jnp.sum(rewards * discount ** jnp.arange(horizon), axis=-1)
    returns = returns + bootstrap_scale * discount**horizon * critic_target(
        states[:, -1]
    )

    return -jnp.mean(returns), (states, rewards)


def critic_loss_fn(critic, critic_target, states, rewards):
    """Regress the critic toward the lambda-returns of a rollout."""
    targets = jax.vmap(lambda_return)(rewards, critic_target(states))

    return jnp.mean(jnp.square(jax.lax.stop_gradient(targets) - critic(states[:, :-1])))

def rollout(cs, state): """Roll the system out for `horizon` steps and reward every state it visits.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) _, states = nnx.split_rngs(splits=batch_size)( nnx.vmap( lambda cs, state: cs(state, num_steps=horizon, return_states=True), in_axes=(state_axes, 0), ) )(cs, state) rewards = jax.vmap(jax.vmap(reward_fn))(states) return states, rewards def actor_loss_fn(cs, critic_target, state, bootstrap_scale): """Compute the negative bootstrapped return of a rollout. The gradient flows analytically through the rollout, and through the value of the state it ends in, which is what carries credit past the horizon. `bootstrap_scale` ramps up early in training so that the policy is not steered by the critic before the critic is accurate. """ states, rewards = rollout(cs, state) returns = jnp.sum(rewards * discount ** jnp.arange(horizon), axis=-1) returns = returns + bootstrap_scale * discount**horizon * critic_target( states[:, -1] ) return -jnp.mean(returns), (states, rewards) def critic_loss_fn(critic, critic_target, states, rewards): """Regress the critic toward the lambda-returns of a rollout.""" targets = jax.vmap(lambda_return)(rewards, critic_target(states)) return jnp.mean(jnp.square(jax.lax.stop_gradient(targets) - critic(states[:, :-1])))

Train step¶

In [17]:

@nnx.jit
def train_step(
    cs, critic, critic_target, optimizer, critic_optimizer, pool, key, scale
):
    """Train step."""
    # Sample from pool
    pool_idx, batch = pool.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 new states to replace the worst
    new_state = sample_state()
    current_state = current_state.at[:num_resets].set(new_state)

    # Update the policy through the rollout
    (loss, (states, rewards)), grad = nnx.value_and_grad(
        actor_loss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params)
    )(cs, critic_target, current_state, scale)
    optimizer.update(cs, grad)

    # Fit the critic on the states the rollout just visited
    states = jnp.concatenate([current_state[:, None], states], axis=1)
    for _ in range(critic_num_updates):
        critic_loss, critic_grad = nnx.value_and_grad(critic_loss_fn)(
            critic, critic_target, states, rewards
        )
        critic_optimizer.update(critic, critic_grad)

    # Let the target critic drift toward the critic
    nnx.update(
        critic_target,
        jax.tree.map(
            lambda t, p: (1.0 - target_step_size) * t + target_step_size * p,
            nnx.state(critic_target, nnx.Param),
            nnx.state(critic, nnx.Param),
        ),
    )

    pool = pool.update(pool_idx, {"state": states[:, -1]})
    return loss, critic_loss, pool

@nnx.jit def train_step( cs, critic, critic_target, optimizer, critic_optimizer, pool, key, scale ): """Train step.""" # Sample from pool pool_idx, batch = pool.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 new states to replace the worst new_state = sample_state() current_state = current_state.at[:num_resets].set(new_state) # Update the policy through the rollout (loss, (states, rewards)), grad = nnx.value_and_grad( actor_loss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params) )(cs, critic_target, current_state, scale) optimizer.update(cs, grad) # Fit the critic on the states the rollout just visited states = jnp.concatenate([current_state[:, None], states], axis=1) for _ in range(critic_num_updates): critic_loss, critic_grad = nnx.value_and_grad(critic_loss_fn)( critic, critic_target, states, rewards ) critic_optimizer.update(critic, critic_grad) # Let the target critic drift toward the critic nnx.update( critic_target, jax.tree.map( lambda t, p: (1.0 - target_step_size) * t + target_step_size * p, nnx.state(critic_target, nnx.Param), nnx.state(critic, nnx.Param), ), ) pool = pool.update(pool_idx, {"state": states[:, -1]}) return loss, critic_loss, pool

Main loop¶

In [18]:

print_interval = 128

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    scale = jnp.minimum(i / bootstrap_warmup, 1.0)
    loss, critic_loss, pool = train_step(
        cs, critic, critic_target, optimizer, critic_optimizer, pool, subkey, scale
    )

    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
        pool_mse = jnp.mean(jax.vmap(mse)(pool.data["state"]))
        print(
            f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s "
            f"| Loss {avg_loss:.3e} | Critic {critic_loss:.2e} "
            f"| Pool MSE {pool_mse:.3e}"
        )

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

print_interval = 128 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) scale = jnp.minimum(i / bootstrap_warmup, 1.0) loss, critic_loss, pool = train_step( cs, critic, critic_target, optimizer, critic_optimizer, pool, subkey, scale ) 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 pool_mse = jnp.mean(jax.vmap(mse)(pool.data["state"])) print( f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s " f"| Loss {avg_loss:.3e} | Critic {critic_loss:.2e} " f"| Pool MSE {pool_mse:.3e}" ) print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

W0904 23:56:42.126720  573306 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.
W0904 23:56:42.247375  573306 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 |   10.0s | Loss 1.887e+00 | Critic 1.93e-01 | Pool MSE 3.043e-02

Step  128/16384 |   23.5s | Loss 1.704e+00 | Critic 5.05e-03 | Pool MSE 1.568e-02

Step  256/16384 |   36.7s | Loss 2.215e+00 | Critic 1.14e-03 | Pool MSE 1.237e-02

Step  384/16384 |   49.9s | Loss 2.547e+00 | Critic 7.50e-04 | Pool MSE 9.959e-03

Step  512/16384 |   63.2s | Loss 2.657e+00 | Critic 7.26e-04 | Pool MSE 7.949e-03

Step  640/16384 |   76.4s | Loss 2.671e+00 | Critic 2.21e-04 | Pool MSE 5.742e-03

Step  768/16384 |   89.6s | Loss 2.529e+00 | Critic 7.75e-05 | Pool MSE 3.662e-03

Step  896/16384 |  102.9s | Loss 2.436e+00 | Critic 7.48e-05 | Pool MSE 2.917e-03

Step 1024/16384 |  116.1s | Loss 2.360e+00 | Critic 6.25e-05 | Pool MSE 2.301e-03

Step 1152/16384 |  129.3s | Loss 2.269e+00 | Critic 4.99e-05 | Pool MSE 1.952e-03

Step 1280/16384 |  142.5s | Loss 2.176e+00 | Critic 6.94e-05 | Pool MSE 1.783e-03

Step 1408/16384 |  155.8s | Loss 2.076e+00 | Critic 7.82e-05 | Pool MSE 1.515e-03

Step 1536/16384 |  169.0s | Loss 1.975e+00 | Critic 4.80e-05 | Pool MSE 1.386e-03

Step 1664/16384 |  182.2s | Loss 1.866e+00 | Critic 6.39e-05 | Pool MSE 1.180e-03

Step 1792/16384 |  195.4s | Loss 1.762e+00 | Critic 5.34e-05 | Pool MSE 1.131e-03

Step 1920/16384 |  208.7s | Loss 1.659e+00 | Critic 6.31e-05 | Pool MSE 1.049e-03

Step 2048/16384 |  221.9s | Loss 1.573e+00 | Critic 4.64e-05 | Pool MSE 9.645e-04

Step 2176/16384 |  235.1s | Loss 1.490e+00 | Critic 5.81e-05 | Pool MSE 9.057e-04

Step 2304/16384 |  248.3s | Loss 1.421e+00 | Critic 3.23e-05 | Pool MSE 8.407e-04

Step 2432/16384 |  261.6s | Loss 1.354e+00 | Critic 4.16e-05 | Pool MSE 8.098e-04

Step 2560/16384 |  274.8s | Loss 1.288e+00 | Critic 3.43e-05 | Pool MSE 7.454e-04

Step 2688/16384 |  288.0s | Loss 1.228e+00 | Critic 1.83e-04 | Pool MSE 7.371e-04

Step 2816/16384 |  301.2s | Loss 1.166e+00 | Critic 2.49e-05 | Pool MSE 6.945e-04

Step 2944/16384 |  314.5s | Loss 1.107e+00 | Critic 2.17e-05 | Pool MSE 6.604e-04

Step 3072/16384 |  327.7s | Loss 1.055e+00 | Critic 1.96e-05 | Pool MSE 6.640e-04

Step 3200/16384 |  340.9s | Loss 1.005e+00 | Critic 2.79e-05 | Pool MSE 6.639e-04

Step 3328/16384 |  354.1s | Loss 9.596e-01 | Critic 2.34e-05 | Pool MSE 6.195e-04

Step 3456/16384 |  367.3s | Loss 9.114e-01 | Critic 2.44e-05 | Pool MSE 5.680e-04

Step 3584/16384 |  380.6s | Loss 8.693e-01 | Critic 1.87e-05 | Pool MSE 5.453e-04

Step 3712/16384 |  393.8s | Loss 8.350e-01 | Critic 1.42e-05 | Pool MSE 5.844e-04

Step 3840/16384 |  407.0s | Loss 7.999e-01 | Critic 3.17e-05 | Pool MSE 1.745e-03

Step 3968/16384 |  420.2s | Loss 7.569e-01 | Critic 7.68e-05 | Pool MSE 1.560e-03

Step 4096/16384 |  433.5s | Loss 7.384e-01 | Critic 1.05e-05 | Pool MSE 9.781e-04

Step 4224/16384 |  446.7s | Loss 7.328e-01 | Critic 1.49e-05 | Pool MSE 8.637e-04

Step 4352/16384 |  459.9s | Loss 7.281e-01 | Critic 1.10e-05 | Pool MSE 8.055e-04

Step 4480/16384 |  473.1s | Loss 7.186e-01 | Critic 2.08e-05 | Pool MSE 9.025e-04

Step 4608/16384 |  486.4s | Loss 7.188e-01 | Critic 3.86e-05 | Pool MSE 7.698e-04

Step 4736/16384 |  499.7s | Loss 7.122e-01 | Critic 1.89e-05 | Pool MSE 9.899e-04

Step 4864/16384 |  512.9s | Loss 6.979e-01 | Critic 3.41e-05 | Pool MSE 9.580e-04

Step 4992/16384 |  526.1s | Loss 6.889e-01 | Critic 1.93e-05 | Pool MSE 8.691e-04

Step 5120/16384 |  539.3s | Loss 6.886e-01 | Critic 1.30e-05 | Pool MSE 7.161e-04

Step 5248/16384 |  552.6s | Loss 6.793e-01 | Critic 3.83e-05 | Pool MSE 7.956e-04

Step 5376/16384 |  565.8s | Loss 6.834e-01 | Critic 1.26e-05 | Pool MSE 8.543e-04

Step 5504/16384 |  579.0s | Loss 6.715e-01 | Critic 1.10e-05 | Pool MSE 7.246e-04

Step 5632/16384 |  592.3s | Loss 6.742e-01 | Critic 1.69e-05 | Pool MSE 7.886e-04

Step 5760/16384 |  605.5s | Loss 6.643e-01 | Critic 8.29e-06 | Pool MSE 6.969e-04

Step 5888/16384 |  618.7s | Loss 6.643e-01 | Critic 1.10e-05 | Pool MSE 7.035e-04

Step 6016/16384 |  631.9s | Loss 6.528e-01 | Critic 2.66e-05 | Pool MSE 6.613e-04

Step 6144/16384 |  645.1s | Loss 6.318e-01 | Critic 1.31e-05 | Pool MSE 6.209e-04

Step 6272/16384 |  658.4s | Loss 6.117e-01 | Critic 1.35e-05 | Pool MSE 5.391e-04

Step 6400/16384 |  671.6s | Loss 6.124e-01 | Critic 1.84e-05 | Pool MSE 6.377e-04

Step 6528/16384 |  684.8s | Loss 5.978e-01 | Critic 3.50e-05 | Pool MSE 5.469e-04

Step 6656/16384 |  698.0s | Loss 5.989e-01 | Critic 1.45e-05 | Pool MSE 6.038e-04

Step 6784/16384 |  711.3s | Loss 5.845e-01 | Critic 1.69e-05 | Pool MSE 5.565e-04

Step 6912/16384 |  724.5s | Loss 5.801e-01 | Critic 1.10e-05 | Pool MSE 6.660e-04

Step 7040/16384 |  737.7s | Loss 5.649e-01 | Critic 2.26e-05 | Pool MSE 7.558e-04

Step 7168/16384 |  750.9s | Loss 5.656e-01 | Critic 1.20e-05 | Pool MSE 5.491e-04

Step 7296/16384 |  764.2s | Loss 5.619e-01 | Critic 2.90e-05 | Pool MSE 5.802e-04

Step 7424/16384 |  777.4s | Loss 5.585e-01 | Critic 1.73e-05 | Pool MSE 6.135e-04

Step 7552/16384 |  790.6s | Loss 5.559e-01 | Critic 1.59e-05 | Pool MSE 7.117e-04

Step 7680/16384 |  803.8s | Loss 5.424e-01 | Critic 2.00e-05 | Pool MSE 5.973e-04

Step 7808/16384 |  817.1s | Loss 5.334e-01 | Critic 1.01e-05 | Pool MSE 5.037e-04

Step 7936/16384 |  830.3s | Loss 5.299e-01 | Critic 1.62e-05 | Pool MSE 5.813e-04

Step 8064/16384 |  843.5s | Loss 5.192e-01 | Critic 2.33e-05 | Pool MSE 4.612e-04

Step 8192/16384 |  856.8s | Loss 5.170e-01 | Critic 3.63e-05 | Pool MSE 7.862e-04

Step 8320/16384 |  870.0s | Loss 5.163e-01 | Critic 1.71e-05 | Pool MSE 5.266e-04

Step 8448/16384 |  883.2s | Loss 5.025e-01 | Critic 1.84e-05 | Pool MSE 1.156e-03

Step 8576/16384 |  896.4s | Loss 4.853e-01 | Critic 3.05e-05 | Pool MSE 6.339e-04

Step 8704/16384 |  909.7s | Loss 4.769e-01 | Critic 3.16e-05 | Pool MSE 5.164e-04

Step 8832/16384 |  922.9s | Loss 4.799e-01 | Critic 1.43e-05 | Pool MSE 4.293e-04

Step 8960/16384 |  936.1s | Loss 4.715e-01 | Critic 1.92e-05 | Pool MSE 4.464e-04

Step 9088/16384 |  949.4s | Loss 4.648e-01 | Critic 1.26e-05 | Pool MSE 4.233e-04

Step 9216/16384 |  962.6s | Loss 4.621e-01 | Critic 1.00e-05 | Pool MSE 4.035e-04

Step 9344/16384 |  975.8s | Loss 4.569e-01 | Critic 2.01e-05 | Pool MSE 3.823e-04

Step 9472/16384 |  989.0s | Loss 4.602e-01 | Critic 1.23e-05 | Pool MSE 3.359e-04

Step 9600/16384 | 1002.3s | Loss 4.492e-01 | Critic 5.09e-05 | Pool MSE 3.468e-04

Step 9728/16384 | 1015.5s | Loss 4.395e-01 | Critic 1.47e-05 | Pool MSE 4.078e-04

Step 9856/16384 | 1028.7s | Loss 4.233e-01 | Critic 1.42e-05 | Pool MSE 5.306e-04

Step 9984/16384 | 1041.9s | Loss 4.013e-01 | Critic 1.25e-05 | Pool MSE 4.189e-04

Step 10112/16384 | 1055.2s | Loss 3.936e-01 | Critic 1.12e-05 | Pool MSE 3.871e-04

Step 10240/16384 | 1068.4s | Loss 3.933e-01 | Critic 9.37e-06 | Pool MSE 3.229e-04

Step 10368/16384 | 1081.6s | Loss 3.875e-01 | Critic 1.68e-05 | Pool MSE 3.078e-04

Step 10496/16384 | 1094.8s | Loss 3.854e-01 | Critic 5.88e-06 | Pool MSE 3.376e-04

Step 10624/16384 | 1108.0s | Loss 3.812e-01 | Critic 1.32e-05 | Pool MSE 2.874e-04

Step 10752/16384 | 1121.3s | Loss 3.733e-01 | Critic 1.08e-05 | Pool MSE 3.606e-04

Step 10880/16384 | 1134.5s | Loss 3.782e-01 | Critic 1.22e-05 | Pool MSE 3.183e-04

Step 11008/16384 | 1147.7s | Loss 3.646e-01 | Critic 7.55e-06 | Pool MSE 2.800e-04

Step 11136/16384 | 1160.9s | Loss 3.573e-01 | Critic 6.93e-06 | Pool MSE 2.727e-04

Step 11264/16384 | 1174.1s | Loss 3.525e-01 | Critic 5.09e-06 | Pool MSE 2.765e-04

Step 11392/16384 | 1187.3s | Loss 3.351e-01 | Critic 1.60e-05 | Pool MSE 4.960e-04

Step 11520/16384 | 1200.6s | Loss 3.246e-01 | Critic 4.84e-04 | Pool MSE 3.536e-04

Step 11648/16384 | 1213.8s | Loss 3.252e-01 | Critic 7.42e-06 | Pool MSE 3.462e-04

Step 11776/16384 | 1227.0s | Loss 3.217e-01 | Critic 5.99e-06 | Pool MSE 2.348e-04

Step 11904/16384 | 1240.2s | Loss 3.197e-01 | Critic 7.72e-06 | Pool MSE 2.493e-04

Step 12032/16384 | 1253.4s | Loss 3.168e-01 | Critic 8.83e-06 | Pool MSE 2.280e-04

Step 12160/16384 | 1266.7s | Loss 3.125e-01 | Critic 1.42e-04 | Pool MSE 2.157e-04

Step 12288/16384 | 1279.9s | Loss 2.989e-01 | Critic 6.42e-06 | Pool MSE 2.779e-04

Step 12416/16384 | 1293.1s | Loss 2.901e-01 | Critic 1.03e-05 | Pool MSE 2.078e-04

Step 12544/16384 | 1306.3s | Loss 2.885e-01 | Critic 9.62e-06 | Pool MSE 2.420e-04

Step 12672/16384 | 1319.5s | Loss 2.798e-01 | Critic 7.86e-06 | Pool MSE 2.049e-04

Step 12800/16384 | 1332.7s | Loss 2.776e-01 | Critic 7.04e-06 | Pool MSE 1.774e-04

Step 12928/16384 | 1346.0s | Loss 2.756e-01 | Critic 8.77e-06 | Pool MSE 1.633e-04

Step 13056/16384 | 1359.2s | Loss 2.728e-01 | Critic 5.53e-06 | Pool MSE 2.053e-04

Step 13184/16384 | 1372.4s | Loss 2.473e-01 | Critic 1.11e-05 | Pool MSE 2.641e-04

Step 13312/16384 | 1385.6s | Loss 2.386e-01 | Critic 9.02e-06 | Pool MSE 3.179e-04

Step 13440/16384 | 1398.9s | Loss 2.360e-01 | Critic 7.22e-06 | Pool MSE 2.646e-04

Step 13568/16384 | 1412.1s | Loss 2.381e-01 | Critic 1.21e-05 | Pool MSE 3.042e-04

Step 13696/16384 | 1425.3s | Loss 2.388e-01 | Critic 6.99e-06 | Pool MSE 2.893e-04

Step 13824/16384 | 1438.5s | Loss 2.415e-01 | Critic 2.03e-05 | Pool MSE 2.876e-04

Step 13952/16384 | 1451.7s | Loss 2.460e-01 | Critic 8.35e-06 | Pool MSE 2.830e-04

Step 14080/16384 | 1464.9s | Loss 2.405e-01 | Critic 9.00e-05 | Pool MSE 3.085e-04

Step 14208/16384 | 1478.1s | Loss 2.388e-01 | Critic 8.57e-06 | Pool MSE 3.153e-04

Step 14336/16384 | 1491.4s | Loss 2.330e-01 | Critic 5.47e-06 | Pool MSE 2.952e-04

Step 14464/16384 | 1504.6s | Loss 2.378e-01 | Critic 9.34e-06 | Pool MSE 3.215e-04

Step 14592/16384 | 1517.8s | Loss 2.419e-01 | Critic 1.43e-05 | Pool MSE 2.924e-04

Step 14720/16384 | 1531.0s | Loss 2.326e-01 | Critic 9.03e-06 | Pool MSE 2.914e-04

Step 14848/16384 | 1544.2s | Loss 2.344e-01 | Critic 1.26e-05 | Pool MSE 3.062e-04

Step 14976/16384 | 1557.5s | Loss 2.388e-01 | Critic 7.16e-06 | Pool MSE 2.557e-04

Step 15104/16384 | 1570.7s | Loss 2.414e-01 | Critic 7.70e-06 | Pool MSE 2.429e-04

Step 15232/16384 | 1583.9s | Loss 2.434e-01 | Critic 1.47e-05 | Pool MSE 2.171e-04

Step 15360/16384 | 1597.1s | Loss 2.392e-01 | Critic 9.42e-06 | Pool MSE 2.055e-04

Step 15488/16384 | 1610.3s | Loss 2.404e-01 | Critic 1.03e-05 | Pool MSE 1.830e-04

Step 15616/16384 | 1623.6s | Loss 2.408e-01 | Critic 1.12e-05 | Pool MSE 1.954e-04

Step 15744/16384 | 1636.8s | Loss 2.268e-01 | Critic 8.04e-06 | Pool MSE 2.193e-04

Step 15872/16384 | 1650.0s | Loss 2.286e-01 | Critic 1.06e-05 | Pool MSE 2.147e-04

Step 16000/16384 | 1663.2s | Loss 2.322e-01 | Critic 6.79e-06 | Pool MSE 2.183e-04

Step 16128/16384 | 1676.4s | Loss 2.324e-01 | Critic 1.60e-05 | Pool MSE 2.336e-04

Step 16256/16384 | 1689.7s | Loss 2.359e-01 | Critic 9.89e-06 | Pool MSE 2.358e-04

Step 16383/16384 | 1702.8s | Loss 2.351e-01 | Critic 1.51e-05 | Pool MSE 2.380e-04
✨ Trained for 16384 steps in 1703s

Run¶

In [19]:

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

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

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