Optimise multi colony

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
@@ -21,13 +24,21 @@ 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
     -------
     chex.Array
         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[pos[:, 0], pos[:, 1]].get()
+        is_nest = (is_nest - 1) == colony_idxs
+
     rewards = jnp.where(is_nest, d_carrying, 0.0)
     return rewards

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

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

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,

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[x[:, 0], x[:, 1]].get() == (i + 1)
 
-    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((n_colonies, *dims), 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