From bde48f6595fe2971ee0cd1be7c69531ca59d61ae Mon Sep 17 00:00:00 2001
From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com>
Date: Tue, 14 Oct 2025 12:35:03 +0100
Subject: [PATCH 1/2] Consolidate multi-colonies into single arrays

---
 src/thants/common/rewards.py          | 10 +++-
 src/thants/common/signals.py          |  8 +--
 src/thants/mono/env.py                |  4 +-
 src/thants/multi/env.py               | 76 ++++++++++++------------
 src/thants/multi/observations.py      | 83 ++++++++++++++++-----------
 src/thants/multi/rewards.py           | 40 ++++++++-----
 src/thants/multi/steps.py             | 42 ++++++++++++--
 src/thants/multi/types.py             | 12 +++-
 tests/test_multi/test_observations.py |  5 +-
 tests/test_multi/test_rewards.py      | 55 ++++++++++--------
 tests/test_multi/test_steps.py        |  3 +-
 11 files changed, 212 insertions(+), 126 deletions(-)

diff --git a/src/thants/common/rewards.py b/src/thants/common/rewards.py
index 2961fef..2534143 100644
--- a/src/thants/common/rewards.py
+++ b/src/thants/common/rewards.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 import chex
 import jax.numpy as jnp
 
@@ -7,6 +9,7 @@ def delivered_food(
     pos: chex.Array,
     carrying_before: chex.Array,
     carrying_after: chex.Array,
+    colony_idxs: Optional[chex.Array] = None,
 ) -> chex.Array:
     """
     Calculate food deposited by individual ant on a nest
@@ -28,6 +31,11 @@ def delivered_food(
         Array represented deposited food by each agent
     """
     d_carrying = carrying_before - carrying_after
-    is_nest = nest.at[pos[:, 0], pos[:, 1]].get()
+
+    if colony_idxs is None:
+        is_nest = nest.at[pos[:, 0], pos[:, 1]].get()
+    else:
+        is_nest = nest.at[colony_idxs, pos[:, 0], pos[:, 1]].get()
+
     rewards = jnp.where(is_nest, d_carrying, 0.0)
     return rewards
diff --git a/src/thants/common/signals.py b/src/thants/common/signals.py
index 21f19ee..a2165e5 100644
--- a/src/thants/common/signals.py
+++ b/src/thants/common/signals.py
@@ -74,10 +74,10 @@ def __call__(self, key: chex.PRNGKey, signals: chex.Array) -> chex.Array:
 
         signals = (
             (1.0 - self.dissipation_rate) * signals
-            + jnp.roll(dissipate, shift=-1, axis=1)
-            + jnp.roll(dissipate, shift=1, axis=1)
-            + jnp.roll(dissipate, shift=-1, axis=2)
-            + jnp.roll(dissipate, shift=1, axis=2)
+            + jnp.roll(dissipate, shift=-1, axis=-1)
+            + jnp.roll(dissipate, shift=1, axis=-1)
+            + jnp.roll(dissipate, shift=-1, axis=-2)
+            + jnp.roll(dissipate, shift=1, axis=-2)
         )
 
         return signals
diff --git a/src/thants/mono/env.py b/src/thants/mono/env.py
index 72cc884..8026b7c 100644
--- a/src/thants/mono/env.py
+++ b/src/thants/mono/env.py
@@ -37,7 +37,7 @@ class ThantsMonoColony(Environment):
 
     def __init__(
         self,
-        dims: tuple[int, int],
+        dims: tuple[int, int] = (50, 50),
         colony_generator: Optional[ColonyGenerator] = None,
         food_generator: Optional[FoodGenerator] = None,
         terrain_generator: Optional[TerrainGenerator] = None,
@@ -56,7 +56,7 @@ def __init__(
         Parameters
         ----------
         dims
-            Environment grid dimensions
+            Environment grid dimensions, default is a 50x50 environment
         colony_generator
             Initial ant colony state generator, initialises ants and nest values.
             By default, initialises a `BasicColonyGenerator` with 25 ants, 2
diff --git a/src/thants/multi/env.py b/src/thants/multi/env.py
index e0093a5..b25403d 100644
--- a/src/thants/multi/env.py
+++ b/src/thants/multi/env.py
@@ -3,6 +3,7 @@
 
 import chex
 import jax
+import jax.numpy as jnp
 from jumanji import Environment, specs
 from jumanji.types import TimeStep, restart, termination, transition
 from jumanji.viewer import Viewer
@@ -20,16 +21,17 @@
     get_observation_spec,
     get_reward_spec,
 )
-from thants.common.steps import deposit_signals
-from thants.common.types import Ants, Colony, Observations
+from thants.common.steps import update_food
+from thants.common.types import Ants, Observations
+from thants.mono.steps import update_positions
 from thants.multi.colonies_generator import (
     BasicColoniesGenerator,
     ColoniesGenerator,
 )
 from thants.multi.observations import observations_from_state
 from thants.multi.rewards import DeliveredFoodRewards, RewardFn
-from thants.multi.steps import clear_nest, update_food, update_positions
-from thants.multi.types import State
+from thants.multi.steps import clear_nest, deposit_signals, merge_colonies
+from thants.multi.types import Colonies, State
 from thants.multi.viewer import ThantsMultiColonyViewer
 
 
@@ -59,7 +61,7 @@ def __init__(
         Parameters
         ----------
         dims
-            Environment grid dimensions
+            Environment grid dimensions, default is a 50x100 environment
         colonies_generator
             Initial colonies state generator, initialises ants and nest states.
             By default, initialises a `BasicColoniesGenerator` with 2 colonies,
@@ -125,6 +127,7 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, list[TimeStep[Observations]]]
         """
         key, colony_key, food_key, terrain_key = jax.random.split(key, num=4)
         colonies = self._colonies_generator(self.dims, colony_key)
+        colonies = merge_colonies(colonies)
         food = self._food_generator.init(self.dims, food_key)
         terrain = self._terrain_generator(self.dims, terrain_key)
         state = State(
@@ -134,7 +137,7 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, list[TimeStep[Observations]]]
             food=food,
             terrain=terrain,
         )
-        observations = observations_from_state(state)
+        observations = observations_from_state(self.num_agents, state)
         time_steps = [
             restart(observation=obs, shape=(n,))
             for obs, n in zip(observations, self._colonies_generator.n_agents)
@@ -169,57 +172,50 @@ def step(
             Tuple containing new state and list of TimeSteps for each colony
         """
         key, food_key, signals_key = jax.random.split(state.key, num=3)
+        actions = jnp.concatenate(actions, axis=0)
         # Unwrap actions
-        actions = [
-            derive_actions(
-                a,
-                take_food_amount=self.take_food_amount,
-                deposit_food_amount=self.deposit_food_amount,
-                signal_deposit_amount=self.signal_deposit_amount,
-            )
-            for a in actions
-        ]
+        actions = derive_actions(
+            actions,
+            take_food_amount=self.take_food_amount,
+            deposit_food_amount=self.deposit_food_amount,
+            signal_deposit_amount=self.signal_deposit_amount,
+        )
 
         # Apply movements
         new_pos = update_positions(
             self.dims,
-            [c.ants.pos for c in state.colonies],
+            state.colonies.ants.pos,
             state.terrain,
-            [a.movements for a in actions],
+            actions.movements,
         )
         # Pick up and drop-off food for each colony
         new_food, new_carrying = update_food(
             state.food,
             new_pos,
-            [a.take_food for a in actions],
-            [a.deposit_food for a in actions],
-            [c.ants.carrying for c in state.colonies],
+            actions.take_food,
+            actions.deposit_food,
+            state.colonies.ants.carrying,
             self.carry_capacity,
         )
         # Drop any new food
         new_food = self._food_generator.update(food_key, state.step, new_food)
         # Propagate / disperse signals
-        new_signals = [
-            self._signal_dynamics(signals_key, c.signals) for c in state.colonies
-        ]
+        new_signals = self._signal_dynamics(signals_key, state.colonies.signals)
         # Deposit signals
-        new_signals = [
-            deposit_signals(signals, pos, a.deposit_signals)
-            for signals, pos, a in zip(new_signals, new_pos, actions)
-        ]
+        new_signals = deposit_signals(
+            new_signals, new_pos, state.colonies.colony_idx, actions.deposit_signals
+        )
         # Clear food dropped on nests
-        new_food = clear_nest([c.nest for c in state.colonies], new_food)
+        new_food = clear_nest(state.colonies.nests, new_food)
         # Gather updated state
-        colonies = [
-            Colony(
-                ants=Ants(pos=pos, health=c.ants.health, carrying=carrying),
-                signals=signals,
-                nest=c.nest,
-            )
-            for c, pos, signals, carrying in zip(
-                state.colonies, new_pos, new_signals, new_carrying
-            )
-        ]
+        colonies = Colonies(
+            ants=Ants(
+                pos=new_pos, health=state.colonies.ants.health, carrying=new_carrying
+            ),
+            colony_idx=state.colonies.colony_idx,
+            signals=new_signals,
+            nests=state.colonies.nests,
+        )
         new_state = State(
             step=state.step + 1,
             key=key,
@@ -228,9 +224,9 @@ def step(
             terrain=state.terrain,
         )
         # Rewards
-        rewards = self._reward_fn(old_state=state, new_state=new_state)
+        rewards = self._reward_fn(self.num_agents, old_state=state, new_state=new_state)
         # Observations
-        observations = observations_from_state(new_state)
+        observations = observations_from_state(self.num_agents, new_state)
         timestep = [
             jax.lax.cond(
                 state.step >= self.max_steps,
diff --git a/src/thants/multi/observations.py b/src/thants/multi/observations.py
index 59b4d6c..82a0f02 100644
--- a/src/thants/multi/observations.py
+++ b/src/thants/multi/observations.py
@@ -1,17 +1,22 @@
 import chex
 import jax
 import jax.numpy as jnp
+import numpy as np
 
-from thants.common.types import Colony, Observations
+from thants.common.types import Observations
 from thants.multi.types import State
 
 
-def observations_from_state(state: State) -> list[Observations]:
+def observations_from_state(
+    colony_sizes: list[int], state: State
+) -> list[Observations]:
     """
     Generate individual agent observations from state for each colony
 
     Parameters
     ----------
+    colony_sizes
+        List of colony sizes
     state
         Environment state
 
@@ -27,7 +32,7 @@ def observations_from_state(state: State) -> list[Observations]:
         - Food carried by each agent
         - Local environment terrain
     """
-    n_colonies = len(state.colonies)
+    n_colonies = len(colony_sizes)
     dims = state.food.shape
     dims_arr = jnp.array([dims])
     idxs = jnp.indices((3, 3))
@@ -39,37 +44,49 @@ def get_ant_view(arr: chex.Array, i: chex.Array, x: chex.Array) -> chex.Array:
     def get_view(arr: chex.Array, x: chex.Array) -> chex.Array:
         return arr.at[x[:, 0], x[:, 1]].get()
 
-    def get_signals(arr: chex.Array, x: chex.Array) -> chex.Array:
-        return arr.at[:, x[:, 0], x[:, 1]].get()
+    def get_signals(i: int, arr: chex.Array, x: chex.Array) -> chex.Array:
+        return arr.at[i, :, x[:, 0], x[:, 1]].get()
 
-    occupation = jnp.zeros(dims, dtype=float)
-    occupations = [
-        occupation.at[c.ants.pos[:, 0], c.ants.pos[:, 1]].set(1.0)
-        for c in state.colonies
-    ]
-    occupation = jnp.stack(occupations, axis=0)
+    def get_nest(i: int, arr: chex.Array, x: chex.Array) -> chex.Array:
+        return arr.at[i, x[:, 0], x[:, 1]].get()
 
-    def get_observation(i, colony: Colony) -> Observations:
-        view_idxs = jax.vmap(lambda x: (idxs + x) % dims_arr)(colony.ants.pos)
-        a_idxs = (i + jnp.arange(n_colonies)) % n_colonies
-        ants = jax.vmap(get_ant_view, in_axes=(None, None, 0))(
-            occupation, a_idxs, view_idxs
-        )
-        food = jax.vmap(get_view, in_axes=(None, 0))(state.food, view_idxs)
-        signals = jax.vmap(get_signals, in_axes=(None, 0))(colony.signals, view_idxs)
-        nest = jax.vmap(get_view, in_axes=(None, 0))(colony.nest, view_idxs).astype(
-            float
-        )
-        terrain = jax.vmap(get_view, in_axes=(None, 0))(
-            state.terrain, view_idxs
-        ).astype(float)
-        return Observations(
-            ants=ants,
-            food=food,
-            signals=signals,
-            nest=nest,
-            carrying=colony.ants.carrying,
-            terrain=terrain,
+    occupation = jnp.zeros(state.colonies.nests.shape, dtype=float)
+    occupation = occupation.at[
+        state.colonies.colony_idx,
+        state.colonies.ants.pos[:, 0],
+        state.colonies.ants.pos[:, 1],
+    ].set(1.0)
+
+    view_idxs = jax.vmap(lambda x: (idxs + x) % dims_arr)(state.colonies.ants.pos)
+    a_idxs = jax.vmap(lambda i: (i + jnp.arange(n_colonies)) % n_colonies)(
+        state.colonies.colony_idx
+    )
+    ants = jax.vmap(get_ant_view, in_axes=(None, 0, 0))(occupation, a_idxs, view_idxs)
+    food = jax.vmap(get_view, in_axes=(None, 0))(state.food, view_idxs)
+    signals = jax.vmap(get_signals, in_axes=(0, None, 0))(
+        state.colonies.colony_idx, state.colonies.signals, view_idxs
+    )
+    nest = jax.vmap(get_nest, in_axes=(0, None, 0))(
+        state.colonies.colony_idx, state.colonies.nests, view_idxs
+    )
+    terrain = jax.vmap(get_view, in_axes=(None, 0))(state.terrain, view_idxs).astype(
+        float
+    )
+
+    boundaries = [0] + colony_sizes
+    boundaries = np.array(boundaries)
+    boundaries = np.cumsum(boundaries)
+
+    observations = [
+        Observations(
+            ants=ants[a:b],
+            food=food[a:b],
+            signals=signals[a:, b],
+            nest=nest[a:b],
+            terrain=terrain[a:b],
+            carrying=state.colonies.ants.carrying[a:b],
         )
+        for a, b in zip(boundaries[:-1], boundaries[1:])
+    ]
 
-    return [get_observation(i, colony) for i, colony in enumerate(state.colonies)]
+    return observations
diff --git a/src/thants/multi/rewards.py b/src/thants/multi/rewards.py
index 5d7424b..f212ec6 100644
--- a/src/thants/multi/rewards.py
+++ b/src/thants/multi/rewards.py
@@ -2,14 +2,24 @@
 
 import chex
 import jax.numpy as jnp
+import numpy as np
 
 from thants.common.rewards import delivered_food
 from thants.multi.types import State
 
 
+def _get_boundaries(colony_sizes: list[int]) -> np.typing.NDArray:
+    boundaries = [0] + colony_sizes
+    boundaries = np.array(boundaries)
+    boundaries = np.cumsum(boundaries)
+    return boundaries
+
+
 class RewardFn(abc.ABC):
     @abc.abstractmethod
-    def __call__(self, old_state: State, new_state: State) -> list[chex.Array]:
+    def __call__(
+        self, colony_sizes: list[int], old_state: State, new_state: State
+    ) -> list[chex.Array]:
         """
         Generate individual ant rewards for each colony
 
@@ -28,7 +38,9 @@ def __call__(self, old_state: State, new_state: State) -> list[chex.Array]:
 
 
 class NullRewardFn(RewardFn):
-    def __call__(self, old_state: State, new_state: State) -> list[chex.Array]:
+    def __call__(
+        self, colony_sizes: list[int], old_state: State, new_state: State
+    ) -> list[chex.Array]:
         """
         Assigns 0 reward to all agents
 
@@ -44,11 +56,13 @@ def __call__(self, old_state: State, new_state: State) -> list[chex.Array]:
         list[chex.Array]
             List of rewards arrays per colony
         """
-        return [jnp.zeros(c.ants.pos.shape[:1]) for c in new_state.colonies]
+        return [jnp.zeros(n) for n in colony_sizes]
 
 
 class DeliveredFoodRewards(RewardFn):
-    def __call__(self, old_state: State, new_state: State) -> list[chex.Array]:
+    def __call__(
+        self, colony_sizes: list[int], old_state: State, new_state: State
+    ) -> list[chex.Array]:
         """
         Assigns rewards for ants depositing food on their own nest
 
@@ -64,12 +78,12 @@ def __call__(self, old_state: State, new_state: State) -> list[chex.Array]:
         list[chex.Array]
             List of rewards arrays per colony
         """
-        return [
-            delivered_food(
-                new_colony.nest,
-                new_colony.ants.pos,
-                old_colony.ants.carrying,
-                new_colony.ants.carrying,
-            )
-            for old_colony, new_colony in zip(old_state.colonies, new_state.colonies)
-        ]
+        boundaries = _get_boundaries(colony_sizes)
+        rewards = delivered_food(
+            new_state.colonies.nests,
+            new_state.colonies.ants.pos,
+            old_state.colonies.ants.carrying,
+            new_state.colonies.ants.carrying,
+            colony_idxs=new_state.colonies.colony_idx,
+        )
+        return [rewards[a:b] for a, b in zip(boundaries[:-1], boundaries[1:])]
diff --git a/src/thants/multi/steps.py b/src/thants/multi/steps.py
index 7d4488f..3c7cec9 100644
--- a/src/thants/multi/steps.py
+++ b/src/thants/multi/steps.py
@@ -5,6 +5,8 @@
 import jax.numpy as jnp
 
 from thants.common.steps import update_food as _update_food
+from thants.common.types import Ants, Colony, SignalActions
+from thants.multi.types import Colonies
 
 
 def _move(
@@ -106,7 +108,7 @@ def update_food(
     return new_food, new_carrying
 
 
-def clear_nest(nests: list[chex.Array], food: chex.Array) -> chex.Array:
+def clear_nest(nests: chex.Array, food: chex.Array) -> chex.Array:
     """
     Clear food deposited on each colony nest
 
@@ -122,6 +124,38 @@ def clear_nest(nests: list[chex.Array], food: chex.Array) -> chex.Array:
     chex.Array
         Food state
     """
-    removed = jnp.stack([jnp.where(nest, food, 0.0) for nest in nests])
-    removed = jnp.sum(removed, axis=0)
-    return food - removed
+    nests = jnp.any(nests, axis=0)
+    return jnp.where(nests, 0.0, food)
+
+
+def merge_colonies(colonies: list[Colony]) -> Colonies:
+    ant_pos = jnp.concatenate([c.ants.pos for c in colonies], axis=0)
+    ant_carrying = jnp.concatenate([c.ants.carrying for c in colonies], axis=0)
+    ant_health = jnp.concatenate([c.ants.health for c in colonies], axis=0)
+    colony_idx = jnp.concatenate(
+        [jnp.full(c.ants.carrying.shape, i, dtype=int) for i, c in enumerate(colonies)]
+    )
+    signals = jnp.stack([c.signals for c in colonies], axis=0)
+    nests = jnp.stack([c.nest for c in colonies], axis=0)
+
+    return Colonies(
+        ants=Ants(
+            pos=ant_pos,
+            carrying=ant_carrying,
+            health=ant_health,
+        ),
+        colony_idx=colony_idx,
+        signals=signals,
+        nests=nests,
+    )
+
+
+def deposit_signals(
+    signals: chex.Array,
+    pos: chex.Array,
+    colony_idx: chex.Array,
+    deposits: SignalActions,
+) -> chex.Array:
+    return signals.at[colony_idx, deposits.channel, pos[:, 0], pos[:, 1]].add(
+        deposits.amount
+    )
diff --git a/src/thants/multi/types.py b/src/thants/multi/types.py
index 2a8bd3e..ec978e8 100644
--- a/src/thants/multi/types.py
+++ b/src/thants/multi/types.py
@@ -2,7 +2,7 @@
 
 import chex
 
-from thants.common.types import Colony
+from thants.common.types import Ants
 
 if TYPE_CHECKING:
     from dataclasses import dataclass
@@ -10,6 +10,14 @@
     from chex import dataclass
 
 
+@dataclass
+class Colonies:
+    ants: Ants
+    colony_idx: chex.Array  # [n-ants,]
+    signals: chex.Array  # [n-channels, *env-size]
+    nests: chex.Array  # [*env-size]
+
+
 @dataclass
 class State:
     """
@@ -24,6 +32,6 @@ class State:
 
     step: int
     key: chex.PRNGKey
-    colonies: list[Colony]
+    colonies: Colonies
     food: chex.Array  # [*env-size]
     terrain: chex.Array  # [*env-size]
diff --git a/tests/test_multi/test_observations.py b/tests/test_multi/test_observations.py
index 7a64670..c7c82b9 100644
--- a/tests/test_multi/test_observations.py
+++ b/tests/test_multi/test_observations.py
@@ -3,6 +3,7 @@
 
 from thants.common.types import Ants, Colony
 from thants.multi.observations import observations_from_state
+from thants.multi.steps import merge_colonies
 from thants.multi.types import State
 
 
@@ -28,6 +29,8 @@ def test_colony_observations(key: chex.PRNGKey) -> None:
         ),
     ]
 
+    colonies = merge_colonies(colonies)
+
     state = State(
         step=0,
         key=key,
@@ -36,7 +39,7 @@ def test_colony_observations(key: chex.PRNGKey) -> None:
         terrain=jnp.ones(dims, dtype=bool),
     )
 
-    observations = observations_from_state(state)
+    observations = observations_from_state([2, 1], state)
 
     assert isinstance(observations, list)
     assert len(observations) == 2
diff --git a/tests/test_multi/test_rewards.py b/tests/test_multi/test_rewards.py
index 6ce79ff..c2f9f18 100644
--- a/tests/test_multi/test_rewards.py
+++ b/tests/test_multi/test_rewards.py
@@ -2,6 +2,7 @@
 
 from thants.common.types import Ants, Colony
 from thants.multi.rewards import DeliveredFoodRewards
+from thants.multi.steps import merge_colonies
 from thants.multi.types import State
 
 
@@ -25,18 +26,20 @@ def test_food_reward_function(key):
     s0 = State(
         step=0,
         key=key,
-        colonies=[
-            Colony(
-                ants=Ants(pos=pos_a, carrying=carrying_before, health=health),
-                nest=nest_a,
-                signals=signals,
-            ),
-            Colony(
-                ants=Ants(pos=pos_b, carrying=carrying_before, health=health),
-                nest=nest_b,
-                signals=signals,
-            ),
-        ],
+        colonies=merge_colonies(
+            [
+                Colony(
+                    ants=Ants(pos=pos_a, carrying=carrying_before, health=health),
+                    nest=nest_a,
+                    signals=signals,
+                ),
+                Colony(
+                    ants=Ants(pos=pos_b, carrying=carrying_before, health=health),
+                    nest=nest_b,
+                    signals=signals,
+                ),
+            ]
+        ),
         food=food,
         terrain=terrain,
     )
@@ -44,25 +47,27 @@ def test_food_reward_function(key):
     s1 = State(
         step=0,
         key=key,
-        colonies=[
-            Colony(
-                ants=Ants(pos=pos_a, carrying=carrying_after, health=health),
-                nest=nest_a,
-                signals=signals,
-            ),
-            Colony(
-                ants=Ants(pos=pos_b, carrying=carrying_after, health=health),
-                nest=nest_b,
-                signals=signals,
-            ),
-        ],
+        colonies=merge_colonies(
+            [
+                Colony(
+                    ants=Ants(pos=pos_a, carrying=carrying_after, health=health),
+                    nest=nest_a,
+                    signals=signals,
+                ),
+                Colony(
+                    ants=Ants(pos=pos_b, carrying=carrying_after, health=health),
+                    nest=nest_b,
+                    signals=signals,
+                ),
+            ]
+        ),
         food=food,
         terrain=terrain,
     )
 
     reward_fn = DeliveredFoodRewards()
 
-    rewards = reward_fn(s0, s1)
+    rewards = reward_fn([3, 3], s0, s1)
 
     assert isinstance(rewards, list)
     assert len(rewards) == 2
diff --git a/tests/test_multi/test_steps.py b/tests/test_multi/test_steps.py
index 503c8c6..080478c 100644
--- a/tests/test_multi/test_steps.py
+++ b/tests/test_multi/test_steps.py
@@ -67,10 +67,11 @@ def test_clear_food() -> None:
 
     nest_a = jnp.zeros(dims, dtype=bool).at[0].set(True)
     nest_b = jnp.zeros(dims, dtype=bool).at[2].set(True)
+    nests = jnp.stack([nest_a, nest_b], axis=0)
 
     food = jnp.ones(dims)
 
-    new_food = clear_nest([nest_a, nest_b], food)
+    new_food = clear_nest(nests, food)
     expected = jnp.array([[0.0], [1.0], [0.0]])
 
     assert jnp.allclose(new_food, expected)

From 5665813506e6ac373523a1d65705f68889da2b02 Mon Sep 17 00:00:00 2001
From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com>
Date: Tue, 14 Oct 2025 23:37:04 +0100
Subject: [PATCH 2/2] Better nest state and update viewer

---
 src/thants/common/rewards.py                  |  5 +-
 src/thants/multi/observations.py              |  4 +-
 src/thants/multi/rewards.py                   |  6 +++
 src/thants/multi/steps.py                     | 47 +++++++++++++++--
 src/thants/multi/types.py                     | 11 +++-
 src/thants/multi/viewer.py                    | 15 +++---
 tests/test_common/test_signals.py             | 26 +++++++++-
 tests/{test_basic => test_mono}/__init__.py   |  0
 .../test_colony_generator.py                  |  0
 tests/{test_basic => test_mono}/test_env.py   |  0
 .../test_observations.py                      |  0
 tests/{test_basic => test_mono}/test_steps.py |  0
 tests/test_multi/test_steps.py                | 51 +++++++++++++++++--
 13 files changed, 141 insertions(+), 24 deletions(-)
 rename tests/{test_basic => test_mono}/__init__.py (100%)
 rename tests/{test_basic => test_mono}/test_colony_generator.py (100%)
 rename tests/{test_basic => test_mono}/test_env.py (100%)
 rename tests/{test_basic => test_mono}/test_observations.py (100%)
 rename tests/{test_basic => test_mono}/test_steps.py (100%)

diff --git a/src/thants/common/rewards.py b/src/thants/common/rewards.py
index 2534143..7f586a6 100644
--- a/src/thants/common/rewards.py
+++ b/src/thants/common/rewards.py
@@ -24,6 +24,8 @@ def delivered_food(
         Food carried by ants at the start of the step
     carrying_after
         Food carried at the end of the step
+    colony_idxs
+        Colony indices of individual ants
 
     Returns
     -------
@@ -35,7 +37,8 @@ def delivered_food(
     if colony_idxs is None:
         is_nest = nest.at[pos[:, 0], pos[:, 1]].get()
     else:
-        is_nest = nest.at[colony_idxs, pos[:, 0], pos[:, 1]].get()
+        is_nest = nest.at[pos[:, 0], pos[:, 1]].get()
+        is_nest = (is_nest - 1) == colony_idxs
 
     rewards = jnp.where(is_nest, d_carrying, 0.0)
     return rewards
diff --git a/src/thants/multi/observations.py b/src/thants/multi/observations.py
index 82a0f02..7bc700b 100644
--- a/src/thants/multi/observations.py
+++ b/src/thants/multi/observations.py
@@ -48,9 +48,9 @@ def get_signals(i: int, arr: chex.Array, x: chex.Array) -> chex.Array:
         return arr.at[i, :, x[:, 0], x[:, 1]].get()
 
     def get_nest(i: int, arr: chex.Array, x: chex.Array) -> chex.Array:
-        return arr.at[i, x[:, 0], x[:, 1]].get()
+        return arr.at[x[:, 0], x[:, 1]].get() == (i + 1)
 
-    occupation = jnp.zeros(state.colonies.nests.shape, dtype=float)
+    occupation = jnp.zeros((n_colonies, *dims), dtype=float)
     occupation = occupation.at[
         state.colonies.colony_idx,
         state.colonies.ants.pos[:, 0],
diff --git a/src/thants/multi/rewards.py b/src/thants/multi/rewards.py
index f212ec6..18191f8 100644
--- a/src/thants/multi/rewards.py
+++ b/src/thants/multi/rewards.py
@@ -25,6 +25,8 @@ def __call__(
 
         Parameters
         ----------
+        colony_sizes
+            List of colony sizes (i.e. number of ants in each colony)
         old_state
             State at the start of the step
         new_state
@@ -46,6 +48,8 @@ def __call__(
 
         Parameters
         ----------
+        colony_sizes
+            List of colony sizes (i.e. number of ants in each colony)
         old_state
             State at the start of the step
         new_state
@@ -68,6 +72,8 @@ def __call__(
 
         Parameters
         ----------
+        colony_sizes
+            List of colony sizes (i.e. number of ants in each colony)
         old_state
             State at the start of the step
         new_state
diff --git a/src/thants/multi/steps.py b/src/thants/multi/steps.py
index 3c7cec9..a2b38c7 100644
--- a/src/thants/multi/steps.py
+++ b/src/thants/multi/steps.py
@@ -124,20 +124,38 @@ def clear_nest(nests: chex.Array, food: chex.Array) -> chex.Array:
     chex.Array
         Food state
     """
-    nests = jnp.any(nests, axis=0)
-    return jnp.where(nests, 0.0, food)
+    return jnp.where(nests > 0, 0.0, food)
 
 
 def merge_colonies(colonies: list[Colony]) -> Colonies:
+    """
+    Merge a list of colonies into a single state
+
+    Parameters
+    ----------
+    colonies
+        List of individual colonies
+
+    Returns
+    -------
+    Colonies
+        Joint state of the colonies, with added index appointing
+        each ant to a colony
+    """
     ant_pos = jnp.concatenate([c.ants.pos for c in colonies], axis=0)
     ant_carrying = jnp.concatenate([c.ants.carrying for c in colonies], axis=0)
     ant_health = jnp.concatenate([c.ants.health for c in colonies], axis=0)
     colony_idx = jnp.concatenate(
-        [jnp.full(c.ants.carrying.shape, i, dtype=int) for i, c in enumerate(colonies)]
+        [
+            jnp.full(colony.ants.carrying.shape, i, dtype=int)
+            for i, colony in enumerate(colonies)
+        ]
     )
     signals = jnp.stack([c.signals for c in colonies], axis=0)
-    nests = jnp.stack([c.nest for c in colonies], axis=0)
-
+    nests = jnp.stack(
+        [colony.nest.astype(int) * (i + 1) for i, colony in enumerate(colonies)], axis=0
+    )
+    nests = jnp.max(nests, axis=0)
     return Colonies(
         ants=Ants(
             pos=ant_pos,
@@ -156,6 +174,25 @@ def deposit_signals(
     colony_idx: chex.Array,
     deposits: SignalActions,
 ) -> chex.Array:
+    """
+    Deposit signals for the relevant colony and channel
+
+    Parameters
+    ----------
+    signals
+        Current signal states
+    pos
+        Ant positions
+    colony_idx
+        Colony index of each ant
+    deposits
+        Signal deposit amounts
+
+    Returns
+    -------
+    chex.Array
+        Update signal states
+    """
     return signals.at[colony_idx, deposits.channel, pos[:, 0], pos[:, 1]].add(
         deposits.amount
     )
diff --git a/src/thants/multi/types.py b/src/thants/multi/types.py
index ec978e8..a401216 100644
--- a/src/thants/multi/types.py
+++ b/src/thants/multi/types.py
@@ -12,9 +12,18 @@
 
 @dataclass
 class Colonies:
+    """
+    Joint state of multiple colonies
+
+    ants: State of all ants
+    colony_idx: Idx assigning each ant to a colony
+    signals: Signals states for each colony
+    nests: Nest of indices indicating if a cell is a nest of a colony
+    """
+
     ants: Ants
     colony_idx: chex.Array  # [n-ants,]
-    signals: chex.Array  # [n-channels, *env-size]
+    signals: chex.Array  # [n-colonies, n-channels, *env-size]
     nests: chex.Array  # [*env-size]
 
 
diff --git a/src/thants/multi/viewer.py b/src/thants/multi/viewer.py
index 41b6e00..e3c88a5 100644
--- a/src/thants/multi/viewer.py
+++ b/src/thants/multi/viewer.py
@@ -20,11 +20,9 @@ def _draw_env(
     terrain = state.terrain.astype(int)
     terrain = colors.terrain.at[terrain].get()
     nest_colors = jnp.clip(colors.ants + 0.1, 0.0, 1.0)
-    nests = [
-        colony.nest[:, :, jnp.newaxis] * nest_colors[i, jnp.newaxis]
-        for i, colony in enumerate(state.colonies)
-    ]
-    nests = jnp.sum(jnp.stack(nests, axis=0), axis=0)
+    empty_color = jnp.zeros((1, 4))
+    nest_colors = jnp.concatenate([empty_color, nest_colors], axis=0)
+    nests = nest_colors[state.colonies.nests]
     food = jnp.full((*state.food.shape, 4), colors.food)
     return terrain, nests, food
 
@@ -33,10 +31,9 @@ def _draw_env(
 def _draw_ants(state: State, colors: ColorScheme) -> chex.Array:
     dims = state.food.shape
     ants = jnp.zeros((*dims, 4))
-
-    for i, colony in enumerate(state.colonies):
-        ants = ants.at[colony.ants.pos[:, 0], colony.ants.pos[:, 1]].set(colors.ants[i])
-
+    ants = ants.at[state.colonies.ants.pos[:, 0], state.colonies.ants.pos[:, 1]].set(
+        colors.ants[state.colonies.colony_idx]
+    )
     return ants
 
 
diff --git a/tests/test_common/test_signals.py b/tests/test_common/test_signals.py
index 0aed8db..dd2e97a 100644
--- a/tests/test_common/test_signals.py
+++ b/tests/test_common/test_signals.py
@@ -7,7 +7,7 @@
 from thants.common.types import SignalActions
 
 
-def test_total_signals_preserved(key) -> None:
+def test_total_signals_preserved_single(key) -> None:
     signals = jnp.zeros((2, 3, 3))
     signals = signals.at[0, 1, 1].set(1.0)
 
@@ -27,6 +27,30 @@ def step(s: chex.Array, _: None) -> tuple[chex.Array, chex.Array]:
     assert jnp.allclose(total_signals[:, 1], 0.0)
 
 
+def test_total_signals_preserved_multi(key) -> None:
+    signals = jnp.zeros((3, 2, 3, 3))
+    signals = signals.at[0, 0, 1, 1].set(1.0)
+    signals = signals.at[1, 0, 1, 1].set(1.0)
+
+    propagator = BasicSignalPropagator(0.0, 0.1)
+
+    def step(s: chex.Array, _: None) -> tuple[chex.Array, chex.Array]:
+        s = propagator(key, s)
+        return s, s
+
+    _, signal_ts = jax.lax.scan(step, signals, None, 20)
+
+    assert signal_ts.shape == (20, 3, 2, 3, 3)
+
+    total_signals = jnp.sum(signal_ts, axis=(3, 4))
+
+    assert jnp.allclose(total_signals[:, 0, 0], 1.0)
+    assert jnp.allclose(total_signals[:, 0, 1], 0.0)
+    assert jnp.allclose(total_signals[:, 1, 0], 1.0)
+    assert jnp.allclose(total_signals[:, 1, 1], 0.0)
+    assert jnp.allclose(total_signals[:, 2, :], 0.0)
+
+
 def test_deposit_signals() -> None:
     signals = jnp.zeros((2, 3, 3), dtype=float)
     pos = jnp.array([[1, 1]])
diff --git a/tests/test_basic/__init__.py b/tests/test_mono/__init__.py
similarity index 100%
rename from tests/test_basic/__init__.py
rename to tests/test_mono/__init__.py
diff --git a/tests/test_basic/test_colony_generator.py b/tests/test_mono/test_colony_generator.py
similarity index 100%
rename from tests/test_basic/test_colony_generator.py
rename to tests/test_mono/test_colony_generator.py
diff --git a/tests/test_basic/test_env.py b/tests/test_mono/test_env.py
similarity index 100%
rename from tests/test_basic/test_env.py
rename to tests/test_mono/test_env.py
diff --git a/tests/test_basic/test_observations.py b/tests/test_mono/test_observations.py
similarity index 100%
rename from tests/test_basic/test_observations.py
rename to tests/test_mono/test_observations.py
diff --git a/tests/test_basic/test_steps.py b/tests/test_mono/test_steps.py
similarity index 100%
rename from tests/test_basic/test_steps.py
rename to tests/test_mono/test_steps.py
diff --git a/tests/test_multi/test_steps.py b/tests/test_multi/test_steps.py
index 080478c..321f16e 100644
--- a/tests/test_multi/test_steps.py
+++ b/tests/test_multi/test_steps.py
@@ -1,6 +1,13 @@
 import jax.numpy as jnp
 
-from thants.multi.steps import clear_nest, update_food, update_positions
+from thants.common.types import Ants, Colony
+from thants.multi.steps import (
+    clear_nest,
+    merge_colonies,
+    update_food,
+    update_positions,
+)
+from thants.multi.types import Colonies
 
 
 def test_colony_movement() -> None:
@@ -65,13 +72,47 @@ def test_colony_food_update() -> None:
 def test_clear_food() -> None:
     dims = (3, 1)
 
-    nest_a = jnp.zeros(dims, dtype=bool).at[0].set(True)
-    nest_b = jnp.zeros(dims, dtype=bool).at[2].set(True)
-    nests = jnp.stack([nest_a, nest_b], axis=0)
-
+    nests = jnp.array([[1], [0], [2]])
     food = jnp.ones(dims)
 
     new_food = clear_nest(nests, food)
     expected = jnp.array([[0.0], [1.0], [0.0]])
 
     assert jnp.allclose(new_food, expected)
+
+
+def test_merge_colonies():
+    dims = (2, 2)
+
+    colony_a = Colony(
+        ants=Ants(
+            pos=jnp.array([[0, 0], [0, 1]]),
+            carrying=jnp.zeros((2,)),
+            health=jnp.ones((2,)),
+        ),
+        nest=jnp.zeros(dims, dtype=bool).at[0, 0].set(True),
+        signals=jnp.zeros((2, *dims)),
+    )
+    colony_b = Colony(
+        ants=Ants(
+            pos=jnp.array([[1, 1]]),
+            carrying=jnp.zeros((1,)),
+            health=jnp.ones((1,)),
+        ),
+        nest=jnp.zeros(dims, dtype=bool).at[1, 1].set(True),
+        signals=jnp.zeros((2, *dims)),
+    )
+
+    colonies = merge_colonies([colony_a, colony_b])
+
+    assert isinstance(colonies, Colonies)
+
+    expected_pos = jnp.array([[0, 0], [0, 1], [1, 1]])
+    assert jnp.array_equal(colonies.ants.pos, expected_pos)
+    assert colonies.ants.carrying.shape == (3,)
+    assert colonies.ants.health.shape == (3,)
+    expected_nests = jnp.array([[1, 0], [0, 2]])
+    assert jnp.array_equal(colonies.nests, expected_nests)
+    assert colonies.signals.shape == (2, 2, *dims)
+    expected_idxs = jnp.array([0, 0, 1])
+    assert jnp.array_equal(colonies.colony_idx, expected_idxs)