Add unique colors to plants, and shade to aging plants.

Add the option (recommended but not default) to have self-balancing environments where soil never takes over the who environment or disappears.
Update two notebooks to show how to add self-balancing soils to the simulations.

PiperOrigin-RevId: 584826124

self_organising_systems/biomakerca/env_logic.py

@@ -1429,3 +1429,85 @@ def env_increase_age(env: Environment, etd: EnvTypeDef) -> Environment:
   new_state_grid = env.state_grid.at[:,:, evm.AGE_IDX].set(
       env.state_grid[:,:, evm.AGE_IDX] + is_aging_m)
   return evm.update_env_state_grid(env, new_state_grid)
+
+### Balancing the soil
+
+def balance_soil(key: KeyType, env: Environment, config: EnvConfig):
+  """Balance the earth/air proportion for each vertical slice.
+
+  If any vertical slice has a proportion of Earth to air past the
+  config.soil_unbalance_limit (a percentage of the maximum height), then we find
+  the boundary of (earth|immovable) and (air|sun), and replace one of them with:
+  1) air, for high altitudes, or 2) earth, for low altitudes.
+  This is approximated by detecting this boundary and checking its altitude.
+  That is, we do not actually count the amounts of earth and air cells.
+
+  upd_perc is used to make this effect gradual.
+
+  Note that if agent cells are at the boundary, nothing happens. So, if plants
+  actually survive past the boundary, as long as they don't die, soil is not
+  balanced.
+  """
+  type_grid = env.type_grid
+  etd = config.etd
+
+  h = type_grid.shape[0]
+  unbalance_limit = config.soil_unbalance_limit
+  # This is currently hard coded. Feel free to change this if needed and
+  # consider sending a pull request.
+  upd_perc = 0.05
+  earth_min_r = jp.array(h * unbalance_limit, dtype=jp.int32)
+  air_max_r = jp.array(h * (1 -unbalance_limit), dtype=jp.int32)
+
+  # create a mask that checks: 'you are earth|immovable and above is air|sun'.
+  # the highest has index 0 (it is inverted).
+  # the index starts from 1 (not 0) since the first row can never be fine.
+  mask = (((type_grid[1:] == etd.types.EARTH) |
+           (type_grid[1:] == etd.types.IMMOVABLE)) &
+          ((type_grid[:-1] == etd.types.AIR) |
+           (type_grid[:-1] == etd.types.SUN)))
+  # only get the highest value, if it exists.
+  # note that these indexes now would return the position of the 'high-end' of
+  # the seed. That is, the idx is where the top cell (leaf) would be, and idx+1
+  # would be where the bottom (root) cell would be.
+  best_idx_per_column = (mask * jp.arange(mask.shape[0]+1, 1, -1)[:, None]
+                         ).argmax(axis=0)
+  # We need to make sure that it is a valid position, otherwise a value of 0
+  # would be confused.
+  column_m = mask[best_idx_per_column, jp.arange(mask.shape[1])]
+
+  # if best idx is low, create air where earth would be.
+  target_for_air = (best_idx_per_column+1)
+  ku, key = jr.split(key)
+  make_air_mask = (column_m & (target_for_air < earth_min_r) &
+                   (jr.uniform(ku, target_for_air.shape) < upd_perc))
+
+  column_idx = jp.arange(type_grid.shape[1])
+  type_grid = type_grid.at[target_for_air, column_idx].set(
+      etd.types.AIR * make_air_mask +
+      type_grid[target_for_air, column_idx] * (1 - make_air_mask))
+  env = evm.update_env_type_grid(env, type_grid)
+  # We need to reset the state too.
+  state_grid = env.state_grid
+  old_state_slice = state_grid[target_for_air, column_idx]
+  state_grid = state_grid.at[target_for_air, column_idx].set(
+      jp.zeros_like(old_state_slice) * make_air_mask[..., None] +
+      old_state_slice * (1 - make_air_mask[..., None]))
+  env = evm.update_env_state_grid(env, state_grid)
+
+  # if best idx is high, create earth where air would be.
+  ku, key = jr.split(key)
+  make_earth_mask = (column_m & (best_idx_per_column > air_max_r) &
+                     (jr.uniform(ku, target_for_air.shape) < upd_perc))
+  type_grid = type_grid.at[best_idx_per_column, column_idx].set(
+      etd.types.EARTH * make_earth_mask +
+      type_grid[best_idx_per_column, column_idx] * (1 - make_earth_mask))
+  env = evm.update_env_type_grid(env, type_grid)
+  # We need to reset the state too.
+  old_state_slice = state_grid[best_idx_per_column, column_idx]
+  state_grid = state_grid.at[best_idx_per_column, column_idx].set(
+      jp.zeros_like(old_state_slice) * make_earth_mask[..., None] +
+      old_state_slice * (1 - make_earth_mask[..., None]))
+  env = evm.update_env_state_grid(env, state_grid)
+
+  return env

self_organising_systems/biomakerca/environments.py

@@ -423,6 +423,15 @@ class EnvConfig:
       increasing number of materials until they would lose 100% of them at 
       max_lifetime. You can essentially disable this feature by setting an 
       enormous value for this attribute.
+    soil_unbalance_limit: if this value is > 0, the environment will try to
+      balance the earth/air proportion for each vertical slice as described in
+      env_logic.balance_soil. This is useful for making sure that the 
+      environment doesn't degenerate into losing either all earth or air.
+      Defaults to 0 to not break older experiments, but the recommended value is
+      1/3.
+      Note that we could add one more parameter here: the update probability.
+      But since it is almost entirely irrelevant, we keep a default value of
+      0.05. Let us know if there are reasons to vary this.
   """
 
   def __init__(self,
@@ -439,7 +448,8 @@ def __init__(self,
                n_reproduce_per_step=2,
                nutrient_cap=DEFAULT_NUTRIENT_CAP,
                material_nutrient_cap=DEFAULT_MATERIAL_NUTRIENT_CAP,
-               max_lifetime=int(1e6)):
+               max_lifetime=int(1e6),
+               soil_unbalance_limit=0):
     self.agent_state_size = agent_state_size
     self.etd = etd
     self.env_state_size = 4 + self.agent_state_size
@@ -455,6 +465,7 @@ def __init__(self,
     self.nutrient_cap = nutrient_cap
     self.material_nutrient_cap = material_nutrient_cap
     self.max_lifetime = max_lifetime
+    self.soil_unbalance_limit = soil_unbalance_limit
 
   def __str__(self):
     return stringify_class(self)
@@ -758,15 +769,48 @@ def slice_environment_from_center(env, new_w):
 
 ### Visualization of environments
 
-@partial(jit, static_argnames=["config"])
-def grab_image_from_env(env, config):
+def hue_to_rgb(p, q, t):
+  t = jp.mod(t, 1.0)
+
+  # exclusive conditions
+  t_lt_1d6 = (t < 1/6)
+  done = t_lt_1d6
+  t_lt_1d2 = jp.logical_and(jp.logical_not(done), (t < 1/2))
+  done = jp.logical_or(done, t_lt_1d2)
+  t_lt_2d3 = jp.logical_and(jp.logical_not(done), (t < 2/3))
+  t_else = jp.logical_and(jp.logical_not(done), jp.logical_not(t_lt_2d3))
+  return (t_lt_1d6 * (p+(q-p)*6*t) +
+          t_lt_1d2 * q +
+          t_lt_2d3 * (p + (q - p) * (2/3 - t) * 6) +
+          t_else * p)
+
+
+def hsl_to_rgb(h,s,l):
+  """Return an array containing the converted RGB colors, in range [0,1].
+  Assumes h,s,l are in the range [0,1].
+  """
+  l_lt_05 = jp.array(l < 0.5).astype(jp.float32)
+  q = l_lt_05 * l * (1 + s) + (1. - l_lt_05) * (l + s - l * s)
+  print(q)
+  p = 2 * l - q
+  print(p)
+  return jp.stack(
+      [hue_to_rgb(p,q,h+1/3), hue_to_rgb(p,q,h), hue_to_rgb(p,q,h-1/3)], -1)
+
+
+@partial(jit, static_argnames=["config", "color_by_id"])
+def grab_image_from_env(env, config, color_by_id=True, id_color_intensity=0.15):
   """Create a visualization of the environment.
-  
+
   Resulting values are floats ranging from [0,1].
+
+  If color_by_id is True, we blend the agent cell colors with unique per id
+  colors with a mix of id_color_intensity.
   """
   etd = config.etd
   def map_cell(cell_type, state):
     env_c = etd.type_color_map[cell_type]
+
     # EARTH and AIR colors degrade by how little nutrients they have.
     is_earth_f = (cell_type == etd.types.EARTH).astype(jp.float32)
     is_air_f = (cell_type == etd.types.AIR).astype(jp.float32)
@@ -777,4 +821,30 @@ def map_cell(cell_type, state):
         state[EN_ST+AIR_NUTRIENT_RPOS]/
         config.material_nutrient_cap[AIR_NUTRIENT_RPOS])*0.7)
     return env_c
-  return vmap2(map_cell)(env.type_grid, env.state_grid)
+  env_c_grid = vmap2(map_cell)(env.type_grid, env.state_grid)
+
+  if color_by_id:
+    def add_id_colors(env_c, cell_type, agent_id):
+      is_agent_f = etd.is_agent_fn(cell_type).astype(jp.float32)
+      # Agents are slightly colored towards a random hue based on the agent id.
+      # Just using two prime numbers for breakign cyclical coloring.
+      agent_hue = jp.mod(agent_id * 41 / 137, 1.)
+      agent_c = hsl_to_rgb(agent_hue, 0.5, 0.5)
+      env_c = env_c * (1. - is_agent_f) + is_agent_f * (
+          env_c * (1. - id_color_intensity) + agent_c * id_color_intensity)
+      return env_c
+    env_c_grid = vmap2(add_id_colors)(
+        env_c_grid, env.type_grid, env.agent_id_grid)
+
+  # Then degrade colors by how old agents are.
+  def decay_by_age(env_c, cell_type, state):
+    # Agents colors degrade by how old they are.
+    is_agent_f = etd.is_agent_fn(cell_type).astype(jp.float32)
+    age_perc = state[AGE_IDX] / config.max_lifetime
+    env_c = env_c * (1. - is_agent_f) + env_c * is_agent_f * (
+        0.3 + (1 - age_perc) * 0.7)
+    return env_c
+  env_c_grid = vmap2(decay_by_age)(
+      env_c_grid, env.type_grid, env.state_grid)
+
+  return env_c_grid

self_organising_systems/biomakerca/examples/notebooks/eruption_example.ipynb

@@ -10,6 +10,8 @@
         "\n",
         "This colab shows how to run experiments on the Eruption configuration. Eruption contains LAVA and FIRE as extra materials. See the code base for examples on how to create your own materials.\n",
         "\n",
+        "Set soil_unbalance_limit to 1/3 for a more optimal set of experiments. The value of 0 reproduces experiments in the [original YouTube video](https://youtu.be/e8Gl0Ns4XiM).\n",
+        "\n",
         "Copyright 2023 Google LLC\n",
         "\n",
         "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
@@ -124,7 +126,10 @@
       },
       "outputs": [],
       "source": [
+        "soil_unbalance_limit = 0 #@param [0, \"1/3\"] {type:\"raw\"}\n",
+        "\n",
         "config = get_eruption_config()\n",
+        "config.soil_unbalance_limit = soil_unbalance_limit\n",
         "etd = config.etd\n",
         "# Create the exclusive fs\n",
         "excl_fs = make_eruption_excl_fs(etd)\n",
@@ -410,18 +415,80 @@
         "agent_logic = BasicAgentLogic(config, minimal_net=agent_model==\"minimal\")\n",
         "\n",
         "mutator_type = \"randomly_adaptive\" #@param ['basic', 'randomly_adaptive']\n",
-        "sd = 1e-3\n",
+        "sd = 1e-3 #@param ['1e-3', '1e-4'] {type:\"raw\"}\n",
         "mutator = (BasicMutator(sd=sd, change_perc=0.2) if mutator_type == \"basic\"\n",
         "           else RandomlyAdaptiveMutator(init_sd=sd, change_perc=0.2))"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FpWCdi2GHeoS"
+      },
+      "source": [
+        "### Testing init on a regular height (72)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JawG5qGMHeoe"
+      },
+      "outputs": [],
+      "source": [
+        "# first just initialize the environment to a certain state, without making videos.\n",
+        "key = jr.PRNGKey(47)\n",
+        "\n",
+        "h = 72\n",
+        "\n",
+        "n_steps = 50000\n",
+        "\n",
+        "use_tiny_dna = False\n",
+        "\n",
+        "if use_tiny_dna:\n",
+        "  program = updated_tiny_dna\n",
+        "else:\n",
+        "  ku, key = jr.split(key)\n",
+        "  program = agent_logic.initialize(ku)\n",
+        "\n",
+        "ku, key = jr.split(key)\n",
+        "program = mutator.initialize(ku, program)\n",
+        "\n",
+        "# 128 is TOO SMALL!\n",
+        "N_MAX_PROGRAMS = 256\n",
+        "\n",
+        "programs = jp.repeat(program[None, :], N_MAX_PROGRAMS, axis=0)\n",
+        "\n",
+        "env = create_eruption_env(h, config)\n",
+        "ku, key = jr.split(key)\n",
+        "programs, env = run_eruption_env(\n",
+        "    ku, config, programs, env, agent_logic, mutator, n_steps, zoom_sz=6)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "jq-haImhJf0l"
+      },
+      "outputs": [],
+      "source": [
+        "# continue...\n",
+        "n_steps = 50000\n",
+        "ku, key = jr.split(key)\n",
+        "programs, env = run_eruption_env(\n",
+        "    ku, config, programs, env, agent_logic, mutator, n_steps, zoom_sz=6,\n",
+        "    steps_per_frame=64, when_to_double_speed=[])"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "WiFLI7fAPNWF"
       },
       "source": [
-        "### Testing on a small height (36)"
+        "### Testing tiny_dna on a small height (36)"
       ]
     },
     {
@@ -462,13 +529,27 @@
         "    ku, config, programs, env, agent_logic, mutator, n_steps, zoom_sz=12)\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0v-T-EzGF2wh"
+      },
+      "outputs": [],
+      "source": [
+        "# If the video is too big, and trying to run it crashes the runtime, perform a\n",
+        "# manual download.\n",
+        "from colabtools import fileedit\n",
+        "fileedit.download_file('video.mp4', ephemeral=True)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "4mtJeIeUPR2z"
       },
       "source": [
-        "### Testing on a regular height (72)"
+        "### Testing tiny_dna on a regular height (72)"
       ]
     },
     {
@@ -485,7 +566,7 @@
         "\n",
         "h = 72\n",
         "\n",
-        "n_steps = 100000\n",
+        "n_steps = 50000\n",
         "\n",
         "use_tiny_dna = True\n",
         "\n",
@@ -509,6 +590,22 @@
         "    ku, config, programs, env, agent_logic, mutator, n_steps, zoom_sz=6)\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "nfUj1iWwNmIE"
+      },
+      "outputs": [],
+      "source": [
+        "# continue...\n",
+        "n_steps = 50000\n",
+        "ku, key = jr.split(key)\n",
+        "programs, env = run_eruption_env(\n",
+        "    ku, config, programs, env, agent_logic, mutator, n_steps, zoom_sz=6,\n",
+        "    steps_per_frame=64, when_to_double_speed=[])"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb

@@ -109,7 +109,9 @@
         "id": "aRRMQ1aNhqv6"
       },
       "source": [
-        "## Select the configuration, the agent logic and the mutator"
+        "## Select the configuration, the agent logic and the mutator\n",
+        "\n",
+        "Set soil_unbalance_limit to 0 to reproduce the original environment. Set it to 1/3 for having self-balancing environments (recommended)."
       ]
     },
     {
@@ -122,8 +124,11 @@
       "source": [
         "ec_id = \"pestilence\" #@param ['persistence', 'pestilence', 'collaboration', 'sideways']\n",
         "env_width_type = \"landscape\" #@param ['wide', 'landscape', 'square', 'petri']\n",
+        "soil_unbalance_limit = 0 #@param [0, \"1/3\"] {type:\"raw\"}\n",
+        "\n",
         "env_and_config = evm.get_env_and_config(ec_id, width_type=env_width_type)\n",
         "st_env, config = env_and_config\n",
+        "config.soil_unbalance_limit = soil_unbalance_limit\n",
         "\n",
         "agent_model = \"extended\" #@param ['minimal', 'extended']\n",
         "agent_logic = BasicAgentLogic(config, minimal_net=agent_model==\"minimal\")\n",