Jax for scaling ALife tutorial: Biomaker CA

Ettore Randazzo* and Bert Chan**

*Google, Paradigms of Intelligence Team

**Google DeepMind

The environment

That is the core…

But if you want to see how to add more complex properties such as gradients of colors, here is the full implementation

key, env and programs are NOT static.

Highlights of the logic

• All materials operations (air and earth) are executed for the entire grid, and then masked if the material is wrong.
• This is because in jax you would anyway execute all conditions when vmapping.
• Agent ops are more expensive, so instead we have a fixed budget of number of agents to run and we sparsely compute the updates.
• We make sure that in our simulations the budget is always > #alive agents.
• Exclusive operations (acting on somebody else) can cause conflicts, so they are atomically executed:
• You propose an action, and we execute a subset of non-conflicting actions at each step.

The code is quite complex, but it is jittable. If you are brave, here it is.

Implementing exclusive ops

Implementing exclusive ops

​

Implementing exclusive ops

Implementation of AIR and EARTH ops

Now let’s implement some new materials

• Disable earth nutrient generation from IMMUTABLE materials.
• Create WATER: it falls, and if it touches EARTH, it gives it earth nutrients.
• Create a WATER_GENERATOR: it creates WATER under it with a certain chance.

Extra material

Cut for time…

This is a variable input. As long as the size is the same, any function taking this as input won’t recompile.

This is a static input.

If you change it, the function recompiles.

Example on how to create environments and configs

A jittable function.

It detects where to place the seed.

The position is dynamic, so at-set doesn’t compile.

It needs to know the exact size of the update

Agent logic

You just think about what to do yourself.

Then, we vmap it twice to create the logic for a grid.

We can also (and do) optimize it for sparse computations.

Example: performing meta-evolution

See the colab