JAX-Tqdm

github.com/jeremiecoullon/jax-tqdm

This library allows you to add the popular Python tqdm progress bar to JAX compiled scans and loops. The original method was developed by Jeremie Coullon, and this repo packages it up, and it can be installed from pip

pip install jax-tqdm

Its usage is as simple as annotating JAX scans or loops, e.g.

from jax_tqdm import scan_tqdm
from jax import lax
import jax.numpy as jnp

n = 10_000

@scan_tqdm(n)
def step(carry, x):
    return carry + 1, carry + 1

last_number, all_numbers = lax.scan(step, 0, jnp.arange(n))

see the README for more details.

Jaxpr-Viz

github.com/zombie-einstein/jaxpr-viz

This library is designed to visualise JAX computation graphs. JAX has built-in methods to visualise the HLO graph produced by JAX, but when I’ve used this, I’ve found it somewhat too low-level and hard to parse.

The intention of Jaxpr-viz it to generate higher-level representations of the computation graph with more information about the structure of the program, i.e. how jax.jit annotated sub-function are connected. It does this by parsing the jaxpr generated by JAX when it parses the Python input.

So for example a JAX program with nested functions

@jax.jit
def foo(x):
    return 2 * x

@jax.jit
def bar(x):
    x = foo(x)
    return x - 1

when visualised by Jaxpr-viz produces:

By default, it will collapse sub-graphs that only contain built in primitives, to make the structure of the overall program clearer. This can be toggled to show the full details of the computation graph:

It can also visualise more complex primitives like conditional statements

@jax.jit
def conditional(arg):
    return jax.lax.cond(
        arg >= 0.0,
        lambda x_true: x_true + 3.0,
        lambda x_false: x_false - 3.0,
        arg
    )

producing

Switch statements and scans/loops are also supported.

See the repo for more details and installation instructions.

At the moment it used pydot/graphviz in the backend to produce static renders of the computation graph. In future, it might be nice to have a more interactive visualisation, for example to allow regions of the graph to be dynamically collapsed and expanded.

Introduction

This post extends previous work on probabilistic cellular automata (CA) that can be found in this post so have a read of that first.

TLDR: The dynamics of a CA are described by its update rule that maps the previous state of a cell, and its neighbours, to its updated state. In this extension cells are a probability distribution over possible states. The update rules then map probability distributions to probability distributions and the update rules themselves can be probabilistic.

Project Improvements

There have been two main updates:

Log Probability Distributions

As noted in the previous post, due to the recursive nature of the CA, using state probabilities directly often results in numerical underflow as probabilities decay to small values. This can be avoided using log probabilities and techniques like the “log-sum-exp trick”.

The conversion to log probabilities turned out to be fairly straightforward, instead of the direct probabilistic update

we can use

Once we expand $\text{log}P\left(q_{i}^{t}\right)$, we find we need to only store

as the state of the CA.

This allows for much better numerical stability (and larger executions as shown below), at the cost of not being able to represent absolute values like zero.

JAX Implementation

JAX is a Python high performance numerical computation library that has been around for a few years but seems to have recently gained a lot of popularity. There’s a lot to be said about it (I particularly really love the functional API) but it has two particularly killer features:

• Performance: JAX compiles to high performance code via XLA. The compilation and optimisation stage results in high performance CPU code, but can also compile to GPU or even TPU without changes to the code. In particular for this project this allows very fast execution of CA at large scales on GPU (speeding up the optimisation process detailed below).
• Gradients: Programs written using JAX can then (usually) be differentiated, and their gradients found (see here for more details). This has numerous applications across many areas of ML and mathematical modelling, but in this particular case we can look at differentiation of probabilistic cellular automata.

A notebook with examples of using this implementation can be found here.

Differentiability and Optimisation

The probabilistic CA effectively maps an update rule/distribution $R$ and initial distribution/state $S_{0}$ to an output series of states $S_{t}$:

Once implemented in JAX we can then calculate derivatives of the output with respect to inputs e.g.

In this example we will look at differentiating with respect to the update rule, and using this to optimise the rule using gradient descent.

For binary states the update rules is designated by 8 values, $R_{j}$, giving the probability of a previous state mapping to a new state

where $j$ just maps the possible permutations of the previous states to integers.

We can then calculate $\partial S_{t}\mathbin{/}\partial R_{j}$, and use gradient descent to calculate iteratively update the rule

where $L(S_{t})$ here is some loss function on the output state.

Example

The code for this example can be found in this notebook.

As a simple optimisation task we will generate a target outputs state, $S^{\prime}_{t}$ from a known ruleset and then starting from a random initial rule, move the rule towards the known rule using the loss function

The output generated by the CA over the course of training is shown below, along with the target distribution. We can see how the initial random state produces a mostly random output (as might be expected) but the target behaviour is revealed as we move towards the target ruleset.

The change in loss, and gradients show some interesting behaviour (as shown below), training slows initially, before quickly converging after ~20,000 steps which is also reflected in the gradients of the individual rule components $R_{j}$

Conclusion & Next Steps

The combination of a proper log probability implementation and JAXs ability to differentiate (plus performance gain from JAX) turned out really nicely.

It’d now be nice to take this and see if optimisation can be used to find interesting probabilistic CA rules. From initial experimentation the hard part of this was designing a good loss function. The example here relies on having a known target, but what aggregate loss function will generate interesting CA rules? I suspect some interesting entropy measure, but need to do more work on what this might look like.

One are where this might be fruitful is for larger state spaces. The space of binary states is small enough to explore manually, but three states already becomes a much bigger space, being able to search this space effectively using gradient based methods may be an interesting result.

Introduction

With this project I wanted to look at applying RL to a flocking model but also to see if the flock as a whole can be driven by a single RL agent interacting with itself, despite treating the members of the flock as individuals.

The boid model contains rules designed to emerge flocking behavior (as you might see in large bird flocks, or fish shoals) but it’d be interesting to see if these behaviours could be learnt from the bottom up, i.e. the RL agent not learning a policy for the flock as a whole, but a policy at the individual agent level. This might then be a nice route to generating behaviours for other agent based models.

Boids Model

The boid model was developed by Craig Reynolds in 1986, and now you will likely find lots of implementations as it’s a nice hobby project (I’ve written it myself a couple of times). Despite its simplicity the model demonstrates rich emergent behaviour and produces flock formations reminiscent of flocks and shoals seen in the natural world.

The model is formed from a set agents referred to as “boids” each moving about the simulated space with the ability to rotate their trajectory (steer). Each boid follows 3 basic rules referred to as separation, alignment and cohesion

• Separation: The boid steers away from any crowding flockmates to avoid collisions
• Alignment: The boid steers toward the average heading of local flockmates
• Cohesion: The boid steers towards the centre of mass of local flockmates

The model progresses stepwise, at each step the boids steering according to the above rules, then updating boid positions in parallel. These generally form the basic flocking rules, though additional rules can be added for more complex behaviour such as avoiding environmental obstacles or seeking goals.

Motivation

In this specific case having coded up boids before, the rules can be hard to optimize, especially as more complex situations are added (for example environmental objects) and different contributions to the steering vector need to be weighed. So it is interesting to investigate if ML can generate nice optimizations of the existing rule, or policies in this RL case.

More generally, working with agent based models in many cases requires modelling and optimizing parameters for large numbers of homogenous (or homogenous subsets of) agents which can serve as a background for more complex parts of the simulation. For example

• An agent based stock market model might contain a large numbers of simple strategy traders that form the background noise of the market
• A traffic simulator may contain large numbers of simple background agents filling out traffic

As model complexity increases it becomes increasingly hard to both manually program robust complex behaviours and optimize parameters that control those behaviours. In this case RL could be a powerful tool to generate complex policies for these agents from simpler goals.

Large scale multi-agent RL comes at the obvious computational cost of running and training large numbers of agents. As such I thought it would be interesting to investigate whether a single RL agent could be used to design a policy for a set of interacting homogenous agents, and can it learn emergent or cooperative behaviours?

Implementation

Training Environment

The environment consists of a flock of agents (boids), their phase space stored as positions $x_{i}$, heading $\theta_{i}$ and speed $s_{i}$ indexed by agent $i$. The agents live on a torus (this avoids the need to track information on the boundaries) such that $x_{i}=0=l$ where $l$ is the width/height of the space (usually normalized to $1.0$).

The action space of the agents are discrete rotations, for example we might use the values [-π/10, 0, π/10].

The environment treats each agent individually, as such the bulk of the computational work of the environment is generating local views of the flock for each agent, done by shifting and rotating the co-ordinates to centre on each agent, and also relative headings between agents (with the added complication of working on the toroidal space)

• Generate the component matrices $x_{ij} = x_{i} \rightarrow x_{j}$ where $i\neq j$. These are the shortest vectors from agent $i$ to agent $j$ on the torus
• From the components generate the distance matrix $d_{ij}$ the (shortest) Euclidean distance between agents $i$ and $j$
• Generate $\theta_{ij}$ the smallest angle between the headings of pairs of agents
• Sort each agents neighbours by distance from that agent, then only include observations from the nearest $n$ neighbouring agents
• Rotate the shifted vectors to align the axes with each boids heading
• Return the concatenated relative vectors and headings of $n$ nearest neighbours to each agent returning the $n_{\text{agents}}\times 3n$ matrix of observations for each agent

The sorting step ensures that each agents local view should have common features with other agents (as opposed to features arranged according to agent indices).

The rewards signal is based purely on the distances between neighbouring agents $r_{i} = \sum_{j}f(d_{ij})$. Choosing $f(d_{ij})$ was one of the more challenging aspects, an initial choice was as simple binary $f(d_{ij})=1$ if $d_{ij}$ is less than some threshold and $0$ otherwise but this seems not encourage the boids to move closer to each other. A continuous function $\exp(-\alpha d)$ encourages agents to move closer, but due to the toroidal space, and the nature of the flock it seems there are good solutions where a boid is evenly distanced rather than close to other boids.

Along with a penalty for being to close to other boids, the environment currently uses

where $p$ is a large penalty value, and $\alpha$ controls how the rewards scale with distance.

The step(actions) function of the environment, as per the Open-AI API accepts actions and advances the model one step. The environment expects actions for each agent’s, in this case for discrete actions, this would be a 1d array length $n_agents$ indexing the possible rotations. The step function in turn returns local observations and rewards for each boid i.e. the $n_{\text{agents}}\times 3n$ local observation matrix and $n_{\text{agent}}$ rewards.

Agent Based Buffer

To facilitate the multiple boids of the training environment, I’ve expanded the experience buffer to store the transition values for each agent at every step. The buffer acts as a queue with new values replacing the oldest entries, indexed by the current step and agent index. This is not strictly necessary, a single queue could be used (just pushing the all the agent values in order) but this format should allow for histories to be recalled for each agent as might be required for RL agents using recursive networks.

The training loop then proceeds pretty much as a regular DQN agent with the addition that

• Experience samples are drawn uniformly from the agent and steps
• The actions of the dqn are generated from a matrix of local observations from each agent, generating an array of actions for each agent

Results

Training

Training with this model revealed a number of difficulties:

• Local Minima: The model seemed to have a few local policy solutions the agent can settle on. In particular all agents moving in straight lines or all agents always steering in one direction, both of which clearly produce consistent rewards, and in particularly can produce excellent rewards if agents randomly start close to each other. This seemed to be best mitigated by the appropriate choice of rewards function, or potentially penalising simple behaviour patterns.
• Feedback: Since the agent is only interacting with its own actions as training progresses and the exploration parameter decreases the actions produced by the agent can become highly localised on a few actions as agents only gain experience of the limited (and predictable) actions of the flock as a whole. It may be beneficial to either maintain a subset of boids that act in a random manner, or add randomness to the application of the steering vectors.
• Over-training: Linked to the feedback issue, the model quite easily overtrains. The agent seemingly developing flocking behaviours at an optimal number of episodes, before then loosing a lot of reactive policies past that point.

Results are very sensitive to the parameters of the training environment, in particular the choice of velocity, steering angles and reward function. Despite this some nice results can be produced, demonstrating flocking similar to that produced by the designed boids rules despite the simple reward function. Some nice example are shown below, for increasing number of episodes also showing the results of overtraining.

Further Work

I’ve labelled this post part 1 as I’ve I feel there are a number of interesting directions to take this that I want to follow up on:

• Continuous action space: A discrete action space was chose in this case to make use of a DQN agent, but the extension to a continuous space is simple to implement; allowing for continuous steering values, and potentially changes in velocity.
• Environmental Obstacles: Basic obstacles should be a simple extension, in particular spherical objects, with an associated penalty for interception, should be simple to add.
• Competitive Agents: It may be interesting to add adversarial agents into the model perhaps representing predators, or competition for resources. This agents or agent(s) could then also be driven by an RL agent in the same manner. This may also have the benefit of driving the flock agent to generate more novel and robust policies compared to the simple flocking policy.
• Vision Model: Part of the difficulty in designing the environment was creating observations of the flock with fixed dimensions to be passed to the RL agent. This may actually be easier if this is done using a view model for each agent, i.e. each boids generates pixels representing the field of vision of each boid. This would create a standard observation format that would more easily accommodate additional complexity in the model. In practice though this would likely require a ray tracing for each boid which could be potentially very expensive for large flocks of agents.

Modules for the environment and buffer, as well as examples of usage can be found on the github repo. The environment follows the Open AI gym API so should be compatible with RL agents using this that format!

Introduction

I read this paper a while ago and thought the problem it looked at was really interesting. It seems like it might be nice to see if RL could be applied to generating synchronization. This is an interim post looking at the implementation of the model, which will hopefully be used as a RL training environment.

Theory

As per the linked paper, this model models firefly swarms which synchronize their light flashes in a distributed manner. In the model nodes represent fireflies, and each firefly fires/flashes based on its phase $\phi(t)$ and threshold phase $\phi_{t}$. In isolation a fireflies phase increases linearly over time until it reaches $\phi_{t}$ where the firefly fires and resets its phase to 0 (i.e. in isolation the firefly will oscillate firing at regular intervals).

In an effort to co-ordinate their flashes the fireflies react to flashes from other fireflies, updating their phase as $\phi\rightarrow\phi+\Delta\phi$ where the update depends on the current phase

where

In the case the signals are instantaneous the synchronization eventually always occurs. The more interesting case is where there are delays in the signal (or this could be thought of as a finite propagation speed).

Implementation

Given the longer term goal of using this as an RL training environment time is discrete, with all nodes updated at each step. In the case that signals propagate instantly this model would be almost trivial to implement. Simulating delayed signals is done by placing events in the future of each node, so they are processed with a delay as the model steps forward.

The model is initialized for a fixed number of steps $s$ and $n$ nodes along with:

• A $s\times n$ array $P$ storing the phase of each node at time step $t$ of the model
• A $n\times n$ distance matrix $D$ representing the transmission time between each pair of nodes
• A $(s+\max(D))\times n$ array that will track events observed by each node. The additional $max(D)$ rows are required to always allow events to set in the future of each node.

The threshold phase $\phi_{t}$ is fixed at $1$, so the size of a time step is effectively controlled by a parameter $\delta\phi$ which is the amount the phase is increase for each node when no signal is observed

At each step the model then advances using (in pseudocode)

for each node
    if phase >= threshold -> phase = 0

for each node
    if phase == 0
        for each node x (not including this node)
            increment the number of events for x at time t+distance

step <- step+1

for each node
    if events[t][node] > 0 
        node-phase = phase_update(phase)
    else
        node-phase = node-phase + delta-phase

effectively at each step, the nodes fires if its phase is at the threshold value. This firing places an event at step $t+\text{distance}$ for each node (the events are then effectively in the future of each node as we step the model forward). We then advance time and check if any events have been observed and update each nodes phase accordingly.

Note: This does have the drawback that event distances have to be integer values (as it then informs where in the array to add the event). Events have to be placed at minimum size of $1$ to be seen. This does mean a distance $0$ can be used to implement a lack of communication between nodes.

This implementation has been chosen with a couple of things in mind:

• Using this as an RL environment will be easier updates are step based.
• Implemented in numpy with pre-allocated arrays for the state of the model is pretty fast. I think this could be done for continuous time with appropriate scheduling of the event observation, but feel it would be a lot more computation to handle queues of events for each node.

Conclusion

The real trick in this model was modelling the delay/travel time of signals by placing events in the future of each node. It’s interesting to see how the co-ordination breaks down as signal delays increase, or gaps are created in the network.

As noted in the introduction I’d like to use this to train an RL model at a node level, that is the agent should choose how to update it’s phase given the current phase and observed events, and if this results in global coordination.

The code to run the current model,and examples of usage can be found on the github repo.

Introduction

This project was off-shoot of the work on probabilistic cellular automata the intention was to investigate how causal dependencies propagate through a 1d cellular automata (CA), and how these might inform identification of the behaviour of cellular automata.

Theory

If we consider cellular automata rules that only consider a cells nearest neighbours the ca update rule maps triples (of cells) to the state of the cell at the next step. We can represent the state of the cellular automata as triples, on graph that represents adjacent and overlapping triples, where an edge represent shifting a triple of cells to the left and appending the next state (the generalization of this concept is the de Bruijn Graph) i.e. if we move along the state array from left to right we are walking along the corresponding directed graph.

For different rules we can then look at how the rule maps triple to triples, for example rule 110 is defined by the mapping from triple to states:

then the possible updates for a triple-to-triples (as opposed to triples-to-states) from applying rule 110 can be represented by adjacency matrices:

From this it should be clear that the state of a triple can be causally dependent on the previous state in 4 ways:

• The triple always maps to the same value, independent of the neighbourhood
• The next state of the triple depends on both it’s neighbours
• The next state of the triple depends on either it’s left or right neighbours only

We can then use this to examine the causal region that precedes a cell, that is if the state $s$ of triple at position $i$ at step $t$ of the model is $s_{i}^{t}$ then we say the causal region are preceding triples that the triple $s_{i}^{t}$ depends on. For example if the state of a triple $s_{i}^{t}$ is independent of its neighbours then it’s causal region contains $s_{i}^{t-1}$ (and we can then recursively follow this backward). If $s_{i}^{t}$ depends on only on it’s preceding left neighbour then the causal region contains $s_{i-1}^{t-1}$ and $s_{i}^{t-1}$.

Implementation

As usual this was quick to implement in numpy with some judicious use of slicing operations. The algorithm to generate plots showing the size of causal regions followed the following steps:

• Generated the phase space array for the basic CA model (i.e. the space-time state of cells).
• For each cell look-up it’s dependency on its neighbourhood. For each cell assign it a pair containing it’s left and right dependency, where 0 indicates the cell is only dependent on the previous cell and $\pm 1$ indicates the cell is dependent on it’s neighbour. For example:
  • ( 0,0) indicates the cell only depends on the previous cell
  • (-1,0) indicates the cell is dependent on the left (but not the right)
  • ( 0,1) indicates the cell is dependent on the right (but not the left)
  • (-1,1) indicates the cell is dependent on the left and right
• Iterate forward over rows (i.e. time) and for each row, add the contributions from the previous row dependent on the causal dependence.

In the final step we are applying something like the following to each cell:

s[t][i][0] = s[t][i][0] + decay*s[t-1][s+s[t][i][0]][0]
s[t][i][1] = s[t][i][1] + decay*s[t-1][s+s[t][i][1]][1]

the decay term applies more weight to contributions from more recent rows, and also means the size of the causal region does not just explode over time.

The result of this algorithm is a 3d array in the shape [steps][width][2] where the final index is the size of the casual region in the left and right directions on the array respectively. To be able to plot this means flattening using some function that helps represent the underlying dynamics (this was also a similar issue in the probabilistic cellular automata project). I attempted to try and capture both the magnitude of the causal region and the imbalance between left and right causality and so settled on

where $l_{i}^{t}$ and $r_{i}^{t}$ are the left and right dependencies respectively.

Results

I’ve chery-picked some of the interesting examples here, in most cases where the rule has very simple behaviour (i.e. where the rule evolves to a fixed or oscillating pattern) the causal region plot doesn’t reveal much over the phase space-diagram. The most interesting cases seem to be where the rule generates propagating structures over complex/random regions, in these cases the causal regions seem to nicely pick out these structures from the background noise:

In the case of rules that produce totally chaotic behaviour, the causal region plot itself seems to contain the same amount of noise.

Conclusion

Like the probabilistic cellular automata project I feel the algorithm and implementation turned out really nicely but there was not a strong conclusion to be drawn from the results. As noted above in a few cases it seems to nicely pick out patterns from the background noise, but this is not general across the board.

I think part of this might be down to choice of function to flatten the causal information, this is one place where a better choice (or choice of plotting method might reveal more information) though it still seems like it would need to contain information on both the size of the preceding causal region and the dependence on direction.

As usual the code to produce these plots is available here with examples.

Introduction

Most agent based modeling (ABM) frameworks tend to make use of an object-oriented (OOP) principals, and probably with good reason. The agents
in an ABM are generally stateful and thus natural candidates for representation by classes encapsulating state and functionality; and inheritance (or composition) allows common functionality to be reused.

This project was an experiment in creating a “functional” ABM API inspired in part by computation graph building APIs. Generally speaking these allow a developer to construct a computation graph from function definitions with the work of linking everything together done in the background (something like dagster is a good example).

It was also motivated by a couple of things I’ve noticed when writing ABMs using OOP patterns (though these are very subjective!):

• Python has some really great OOP features, but it sometimes feels that the flexibility and dynamic nature of python makes using OOP patterns a bit tedious when writing an ABM; where you often want interfaces and encapsulation to be followed quite strictly. This is obviously less of a concern in other typed/compiled languages, but then you lose the flexibility and speed of work in python.
• I also feel I end up writing a lot of boilerplate code, and ABM APIs tend not to offer much apart from templates classes, or the patterns they do recommend tend to be quite restrictive.

So the aim was to try and create and API that did a lot of work in the background, and offered a clean and flexible API to build a model.

Disclaimer The design I ended up with did not turn out strictly “functional” as it still sometimes relies on updating objects in-place for reasons that will be outlined in this post. Though I feel a proper functional approach is possible with some tweaks.

Theory

Very an broadly an agent based model consists of agents that interact in some manner (possibly inside a simulated environment). As the model progress the state of agents are updated depending on the state of other agents and the environment. It’s usually the job of the model designer to dictate how the agents behave and also how agents interact within the model/environment.

At the core of this is the concept of representing the time evolution of the model as a causal graph.

On the graph agents are represented by nodes, and edges represent causal dependence between agents. Each agent owns its own state, so altogether the agents represent the model/simulation environment, and the graph how the components of the environment interact over time.

Each node (i.e. agent) then has:

• Nodes that precede it causally: The state of these nodes then act as inputs to the update function of the node. These nodes could also be thought of as being the observed state of the model when the node updates
• Downstream nodes that the node can update. This allows the node to alter the environment outside the state it owns.

If you are familiar with DAGS (or computational graphs, neural networks etc.)
this will likely seem familiar with the differences being that the causal graph can be recursive (as agents can update multiple times/repeatedly) and the agents are stateful.

To represent this in a functional framework we break up an agent definition into two components

• The state of the agent
• An update function that is called when the agent is updated

When the agent is updated, its update function is called with the states of the preceding agents and the current state of the agent as arguments. The update function then returns the new state of the node, and any updates to the state of downstream nodes.

Implementation

@agent Decorator

An agents behaviour can be defined by decorating an update function with the @agent decorator. The function should have the signature

@agent(scheduler=scheduler)
def foo(t, antecedents, state, descendants):
    # Update state and descendants
    ...
    return next_event_time

When the update function is called the arguments will be

• t: The current model time
• antecedents: Data structure containing the states of preceding nodes
• state: The state of the node that is updating
• descendants: Data structure containing states of descendant nodes that this node can update

and it returns one value, the time this agent will next update.

The decorator argument scheduler is a class that controls when agents are called and should be provided when the agent is defined (provided as part of the package).

In this example the decorator would create a new type (class) called foo wrapping the update function and allowing it to be used as a component of the model.

As noted above this is where the ‘functional’ approach breaks down as a bit as the update function is expected to update the state and descendants in place. This should be possible by instead having the function return new state and descendants This was done here as in some cases it’s not possible to store a reference and assign a value to it. For example for numpy arrays if you try something

x = np.array([1,2,3])
y = x[:1]
y = np.array([10])

this will not update the slice of x that y refers to, it will just change what y refers to. So in practice it’s easier to pass states by reference and update them in place

Model Initialization

Initializing the model is then done by creating instances of agent types following the same signature of the update function, for example for the foo function we decorated above we might do

# Initial states of the nodes
agent_state_1 = [0]
agent_state_2 = [1]

# Initialize instances of agent nodes
foo(0, agent_state_2, agent_state_1, {})
foo(1, agent_state_1, agent_state_2, {})

Where the t argument is the time of the agents first event. In the background this would initialize agent nodes with agent_state_2 be a precedent of agent_state_1 and vice versa (and both nodes have no downstream dependents).

Example

This is probably better explained with a more in depth example. We can implement a function that initializes and runs Conway’s game of life using this API. In full this looks like:

def gol(steps, initial_state):
    
    scheduler = StepBasedScheduler(steps)
    history = []
    
    @agent(scheduler=scheduler)
    def cell(t, antecedents, state, descendants):
        
        live_neighbours = np.sum(antecedents) - antecedents[1,1]
        
        if live_neighbours < 2:
            new_state = 0
        elif live_neighbours == 2:
            new_state = state[0,0]
        elif live_neighbours == 3:
            new_state = 1
        elif live_neighbours > 3:
            new_state = 0
        
        state[0:1,0:1] = new_state
        
        return t + 1
    
    # Initialize the nodes on a grid        
    for i in range(1, initial_state.shape[0]-1):
        for j in range(1, initial_state.shape[1]-1):
            cell(0, 
                 initial_state[i-1: i+2, j-1: j+2],
                 initial_state[i:i+1, j:j+1],
                 {})
    
    # Run the model for the requested number of steps and 
    # at each step store a copy of the array
    while not scheduler.finished:
        history.append(initial_state.copy())
        scheduler.step()
        
    return history 

This function accepts the initial state of the model (a numpy array) and the number of steps to run the model for. Breaking it down

• This line initializes a scheduler that will run the model for a fixed number of steps

scheduler = StepBasedScheduler(steps)

• The agent behaviour (an agent is a cell on the array in this case) is defined using the decorator

  @agent(scheduler=scheduler)
      def cell(t, antecedents, state, descendants):
          live_neighbours = np.sum(antecedents) - antecedents[1,1]
        
        if live_neighbours < 2:
            new_state = 0
        elif live_neighbours == 2:
            new_state = state[0,0]
        elif live_neighbours == 3:
            new_state = 1
        elif live_neighbours > 3:
            new_state = 0
        
        state[0:1,0:1] = new_state
        
        return t + 1

as expected it looks at the previous state of the cells surrounding it (its antecedents) and counts the number of live cells then updates its own state accordingly. It returns t+1 which is the next of the model (all the cells will use this function so this will all the cells update each step as desired).

• We then initialize the cells with the relevant antecedents and state. In this case the antecedents are slices of the numpy array representing the neighbourhood surrounding each cell; and the state the relevant cell. It’s then just case of iterating over the array and assigning each cell to an agent. We also set the first step of each agent to fire at 0:

  for i in range(1, initial_state.shape[0]-1):
        for j in range(1, initial_state.shape[1]-1):
            cell(0, 
                 initial_state[i-1: i+2, j-1: j+2],
                 initial_state[i:i+1, j:j+1],
                 {})

• The model can then be run using the scheduler, and additionally we store a copy of the state at each step ot track the history of the model

  while not scheduler.finished:
        history.append(initial_state.copy())
        scheduler.step()

This example used numpy as a data-structure to store the state of the model, but the API is very general, and as long as it can be passed by reference any data structure can be used to track the state and pass the relevant parts of it to the update function.

The scheduling is also very general, in this case the model was designed to update all the agents at fixed steps, but you could just as well use datetime as a variable, or have agents act at different intervals.

For more examples see here in the repo.

Conclusion

For a short project I feel this turned out quite neatly. I think one of the nice features that came out as a side effect is being able to make use of different backends to store the state of the model. An approach using classes would require storing state as part of the class structure but with this API we can make more efficient use of data structures to track state. For small models this results in quite neat code.

This does come with some drawbacks though.

• You need to pass everything required as part of the arguments to the update function (where as part of an OOP pattern you would be able to look up attributes on the class). For example, you might want to pass in a global object like a random number generator (the API might benefit from having an additional context argument to make this easier)
• Allowing nodes to alter other nodes can cause some conflicts. This seems like a necessary feature to have to allow for a broader range of models (agents should be able to effect their environments right?) but in a model where agents update on the same step, this can cause issues where an agents state is updated by itself and another node. Though this can be avoided with the appropriate design choices.

A python package to be able to use this API can be found on it’s github repo along with usage examples. Please try it out, it’d be interesting to see how this API holds up across more model implementations.

Introduction

I have to admit I find cellular automaton (CA) endlessly fascinating. They have never seemed to have found that killer real-world applications or ultimate deep insight; but I find the questions they raise about emergent behaviour and self organization incredibly engaging, that always keep me coming back to try out some new approach to exploring them.

In this post I’ll be looking at an extension to CA to allow for cells to be in a mixed state and probabilistic update rules to be applied.

1D Cellular Automata (Briefly)

To keep this post brief I’ll not go into too much depth of the theory on CA, if you’ve not encountered them before there’s a ton of stuff out there and as always wikipedia is a good place to start.

This project focused on 1-dimensional CA. This model can be imagined as a 1d array of cells, with each cell in a (usually discrete) state $s\in S$. At each update of the model, a cells state is updated based on its own state and local neighbourhood of cells. For example if each cell only considers it’s nearest neighbours then an update rule of a CA is the mapping

where $i$ indexes the cells position in the array and $t$ the step. The array is usually wrapped to form a closed loop (i.e. the end-cells are neighbours), and the evolution of the model visualized as a 2d space-time array where each row is one step of the model.

CA can have multiple states, and rules can be defined for various neighbourhoods around a cell, but in this case we will concentrate on 2-state and size 3 neighbourhood rules.

In many cases it’s convenient to index rules using the Wolfram system where the index is derived from the mapping represented in base $\vert S\vert$. For example if $S=\{0,1\}$ then full specifying the update rule requires a mapping for the $\vert S\vert^{3}=2^3$ possible triple states resulting in $2^8=256$ possible rules (See here for more details).

For example rule 110 is defined by the mapping

A CA is then parameterized by an update rule, an initial state (common choice are a single live cell, or a random initial state) and topology of the array (i.e. a closed loop of fixed boundary conditions).

Motivation

This project was motivated by a couple of broad questions that have motivated a great deal of research on 1d CA:

• Is there some way to classify the long term behaviour of CA? If one looks at all the possible rules configurations (for a chosen CA configuration) the dynamics of rules seem to roughly fall into 4 categories

  • Evolution to a static homogeneous state
  • Periodic behaviour where cells oscillate between states at each step or stable inhomogeneous structures
  • Evolution to chaotic or seemingly noisy patterns
  • Formation of localised persistent structures that can interact

  though some rules can overlap some of these behaviors, and can obviously also depend on the initial state chosen. How can these classifications be made rigorous (or is there an underlying statistics or feature that classifies rules) and what forms the boundary between these behaviours?

  Of particular interests are rules where stable structures form against a chaotic background, and such states can propagate and interact. For exmample it has been shown that rule 110 is turing complete.

• In the case where structures do form, how do they form and persist against the background of chaotic noise? Is there some necessary condition for their formation and how do they propagate information and interact?

Note: These are a very brief outline of some of the interesting topics on the subject. There’s really a lot of interesting work covering these topic and more!!

Interesting work has been done on how changes in the initial state change the long term behaviour of the model (an analogue of the Lyapunov exponent used to study chaotic systems) and how information propagates through the array. It would seem expedient then to be able to express the initial state as a distribution across initial states or as something like a random walk that evolves over time as the update rule is applied to it.

This motivates being able to model a CA where the state of a cell can be a probability distribution on the state space (or maybe this could be thought of a superposition of states as in QM, though without any phase information). As the update rule is applied to a neighbourhood of cells, this also requires that we are able to model the joint or conditional probability of neighbouring cells.

Theory

Note: In the literature “Probabilistic CA” seems to refer to a model where the update rule is probabilistic but each cell still always has a fixed discrete state. So for now the name “Continuous Probabilistic CA” seems like a good name to distinguish from this

Probabilistic States

We’ll consider a model where each cell has a probability of being in a state $s$ denoted as

The update rule applies to a cell, and its left and right neighbours. The probability of the cell being in a state at the next step is then given by

where the r.h.s is the joint probability of a state of the preceding triple and the summation is performed over the triples will result in the state $s_{i}^{t+1}$, i.e. when

for all combinations of the preceding states and $f(\dots)$ is the CA update function.

Note: The joint probability is important here as for certain rules neighbouring cells cannot be in certain configurations and as such should not be treated independently

This approach is ok for one step of the model but to repeatedly perform this update for an arbitrary number of steps we need the joint probability of two neighbouring cells at each site, and its neighbour to the right for each time step

extending the above approach for the update of a single cell gives

the summation now over the overlapping triple cell states that will create the joint state. Then using the chain rule of probabilities we can decompose this to

where

This form is then useful as we need only store the joint probabilities for each cell (and it’s neighbour) and can obtain $h$ from this and the marginal probabilities.

Finally for the starting state at $t=0$ we assume that the initial probabilities are independent such that the initial array of joint probabilities can be found using

Probabilistic Update Rules

In the above model, any uncertainty in the model can only arise from uncertainty in the initial state (if all the cells are in only one state initially i.e. $P(s)\in\{0,1\}$ it behaves like a regular CA), as the update rule still maps discrete states to discrete states.

We can extend the model to include probabilistic updates, represented by the conditional distribution

where $q_{i}^{t}$ is the preceding quadruple of states that inform the updated joint probability

this actually somewhat simplifies the form of the summation to give

where the summation is over all the possible patterns of the preceding quadruples cells.

In the case that

then a deterministic CA rule is recovered.

Implementation

This was straightforward enough to implement in numpy. As with most cellular automata models like this (where the update is done in parallel for all cells) the real trick is shifting the state array left and right to be able to vectorize the update step.

The model is specified by 3 parameters (as in a regular CA):

• The initial state of the cells, in this case it has the dimensions $\text{steps}\times\vert S\vert$ specifying the initial probability distribution for each cell

• The update rule, provided as a 2d array mapping the permutation of triples to the probability of the update state. This represents:

  \[P\left(s_{i}^{t+1}\:\vert\: s_{i-1}^{t}, s_{i}^{t}, s_{i+1}^{t}\right)\]

• The number of steps to run the model for

In brief the models then follows the following steps:

• Use the rule array to generate an array representing the update of joint probabilities conditioned on the preceding quadruples of cells

  \[P\left(s_{i}^{t+1}, s_{i+1}^{t+1}\:\vert\: q_{i}^{t}\right)\]

• Initialize an empty array for joint probabilities with the shape

  \[\text{steps}\times \text{width}\times \vert S \vert\times\vert S \vert\]

  (i.e. a joint probability distribution for each cell and for each step of the model)

• Set the initial joint probability row from the initial state, using the independence of the initial states
• At each step calculate the marginal probabilities required to calculate $h\left(s_{i}^{t+1},s_{i+1}^{t+1}\right)$ along with the left and right shifted rows for $P\left(s_{i-1}^{t},s_{i}^{t}\right)$ and $P\left(s_{i+1}^{t},s_{i+2}^{t}\right)$
• For each joint probability entry sum over all the contributions from the combinations of the preceding quadruple of cells

Complexity and Numerical Underflow

Computationally there are a couple of potential drawbacks

• Computational complexity: Increasing the number of states increases both the storage space required and the complexity of the update calculation. The number of states means scaling the array of joint probabilities like $\vert S\vert^2$ for each cell. When then need to then perform the update for each of these new entities as well as including contributions for the additional permutation of states which scale as $\vert S\vert^4$.
• Numerical Underflow: As with many models where probabilities are repeatedly multiplied numerical underflow can occur. A common approach is to work with log probabilities and use techniques like the log-sum-exp trick. In this case there a couple of things that make this tricky:
  • The lack of a well-defined 0 probability in log space prohibits using binary states as inputs to the model
  • Calculating the marginal probabilities required in the update step still means moving between log and normal probabilities which does not aid in reducing numerical under/overflow

Analysis/Plotting

Plotting the result of the model as a time-space diagram like a regular CA requires aggregating the joint probability distribution of each cell in some manner, but I also looked to choose statistics that might reveal underlying dynamics of the probabilistic CA:

• The probability distribution for each cell are easily recovered as the marginals of the joint distribution

  \[P\left(s_{i}^{t}\right)=\sum_{s_{i+1}^{t}}P\left(s_{i}^{t}, s_{i+1}^{t}\right)\]

• The mutual information of the joint probabilities also seems like it should be a useful, giving something like mutual dependence between neighbouring cells. Here defined for a single cell as

  \[I_{i}^{t} = \sum_{s_{i}^{t}, s_{i+1}^{t}}P\left(s_{i}^{t}, s_{i+1}^{t}\right) \text{log}\left(\frac{P\left(s_{i}^{t}, s_{i+1}^{t}\right)}{P\left(s_{i}^{t}\right)P\left(s_{i+1}^{t}\right)}\right)\]

An additional issue is that in many cases the relative difference (of these statistics) between cells decreases over time, meaning plots can fail to adequately show features contained in resulting arrays. Applying min-max scaling across rows of the array is one approach used to tackle this issue, though care should be taken in how this is interpreted, given that small relative differences could also be a result of numerical precision.

Results

Given the large number of potential parametrizations of the model (mixing probabilistic update rules, and initial states) I looked to focus on two simple cases:

• Standard update rules (i.e. deterministic) with a randomly chosen discrete initial state containing a single cell in a mixed state. Used to investigate how the uncertainty from a single cell propagates through the state over time.

• Update rules with the same perturbation applied to each mapping, and a randomly chosen discrete initial state. For example if the undeterred rule maps $(0,0,0) \rightarrow 0$ then the perturbed probabilities are $P(0\vert 0,0,0)=1-\epsilon$ and $P(1\vert 0,0,0)=\epsilon$ (applied to all the permutations). This could be thought of as a probability of error when applying the update rule, we can then consider how robust rules and patterns are to these errors.

Below are some examples of space-time diagrams from these cases (though far exhaustive given the possible combinations of parameters). In each case the left image is the equivalent regular CA evolution (i.e. no perturbation in the update rule or initial state). The right hand side has been coloured to emphasise values between 0-1.

Perturbed Initial State

Perturbed Update Rule

Conclusion

In terms of a model, I feel this turned out quite well. Extending a regular CA into one that supports mixed/probabilistic cell states turned out to be quite a neat algorithm, and the resulting implementation relatively speedy. It’d be nice to have a robust way of working in log-probability space, but the current model seems relatively robust in most cases. Currently, the model-runner only supports binary states, but the extension to larger state spaces should be relatively straight forward.

Unfortunately the results are mostly qualitative, there seem to be some nice features revealed that point towards interesting dynamics and information propagation between cells. One could maybe make some statements about robustness of patterns, or the propagation of uncertainty relate to the classes of behaviours (described earlier in this post) but this is in need of further analysis.

One thing that might be nice to explore is the transition between the behaviour of update rules. Being able to apply probabilistic update rules means one could explore the continuous space of update rules inside the $\vert S\vert$-dimensional interval (of which the regular discrete CA rules form the corners).

At a minimum though, some of the images would make for cool album covers!

All the code needed to run the model and produce the plots included in this post is available here so please go ahead and try it out, it’d be interesting to see what other people come up with!