Cleanup

README.md

@@ -34,11 +34,15 @@ pip install thants
 
 #### Single Colony
 
+The single colony environment follows the [Jumanji](https://github.com/instadeepai/jumanji)
+environment API, with actions provided as an array of individual
+actions:
+
 ```python
-from thants.envs import ThantsMonoColony
+from thants.envs import ThantsMono
 import jax
 
-env = ThantsMonoColony(dims=(50, 50))
+env = ThantsMono(dims=(50, 50))
 key = jax.random.PRNGKey(101)
 state, obs = env.reset(key)
 state_history = [state]
@@ -57,15 +61,15 @@ env.animate(state_history, 100, "mono_colony.gif")
 #### Multi-Colony
 
 In the multi-colony case each colony is treated independently (and can be
-different sizes), so actions and observations are list/tuples of arrays/structs.
+different sizes), so actions, observations, timesteps are list/tuples of
+arrays/structs:
 
 ```python
-from thants.envs import ThantsMultiColony
+from thants.envs import Thants
 import jax
 import jax.numpy as jnp
 
-
-env = ThantsMultiColony((50, 100))
+env = Thants((50, 100))
 key = jax.random.PRNGKey(101)
 state, obs = env.reset(key)
 state_history = [state]
@@ -83,6 +87,9 @@ for _ in range(50):
 env.animate(state_history, 100, "multi_colony.gif")
 ```
 
+Preset simple environments can be imported from `thants.envs.ThantsDual` and
+`thants.envs.ThantsQuad` with 2 and 4 colonies respectively.
+
 ## Environment
 
 <div align="center">

src/thants/envs/__init__.py

@@ -1,2 +1,3 @@
-from thants.envs.mono import ThantsMonoColony
-from thants.envs.multi import ThantsMultiColony
+from thants.envs.mono import ThantsMono
+from thants.envs.multi import Thants
+from thants.envs.presets import ThantsDual, ThantsQuad

src/thants/envs/mono.py

@@ -7,7 +7,7 @@
 from jumanji.viewer import Viewer
 from matplotlib.animation import FuncAnimation
 
-from thants.envs.multi import ThantsMultiColony
+from thants.envs.multi import Thants
 from thants.generators.colonies.mono import (
     BasicColonyGenerator,
     ColonyGenerator,
@@ -20,7 +20,7 @@
 from thants.types import Observations, State
 
 
-class ThantsMonoColony(Environment):
+class ThantsMono(Environment):
     """Environment with a single colony"""
 
     def __init__(
@@ -78,7 +78,7 @@ def __init__(
         """
         colony_generator = colony_generator or BasicColonyGenerator(25, 2, (5, 5))
         colony_generator = SingleColonyWrapper(colony_generator)
-        self.env = ThantsMultiColony(
+        self.env = Thants(
             dims=dims,
             colonies_generator=colony_generator,
             food_generator=food_generator,
@@ -141,6 +141,10 @@ def step(
         state, timestep = self.env.step(state, [actions])
         return state, timestep[0]
 
+    @cached_property
+    def dims(self) -> tuple[int, int]:
+        return self.env.dims
+
     @cached_property
     def num_agents(self) -> int:
         return self.env.num_agents[0]

src/thants/envs/multi.py

@@ -11,8 +11,8 @@
 
 from thants.actions import derive_actions
 from thants.generators.colonies.multi import (
-    BasicColoniesGenerator,
     ColoniesGenerator,
+    DualBasicColoniesGenerator,
 )
 from thants.generators.food import BasicFoodGenerator, FoodGenerator
 from thants.generators.terrain import OpenTerrainGenerator, TerrainGenerator
@@ -31,7 +31,7 @@
 from thants.viewer import ThantsMultiColonyViewer
 
 
-class ThantsMultiColony(Environment):
+class Thants(Environment):
     """
     Thants environment with multiple colonies
     """
@@ -95,7 +95,7 @@ def __init__(
         self.deposit_food_amount = deposit_food_amount
         self.signal_deposit_amount = signal_deposit_amount
         self.max_steps = max_steps
-        self._colonies_generator = colonies_generator or BasicColoniesGenerator(
+        self._colonies_generator = colonies_generator or DualBasicColoniesGenerator(
             [25, 25], 2, (5, 5)
         )
         self._food_generator = food_generator or BasicFoodGenerator((5, 5), 100, 1.0)

src/thants/envs/presets.py

@@ -0,0 +1,132 @@
+from thants.envs.multi import Thants
+from thants.generators.colonies.multi import (
+    DualBasicColoniesGenerator,
+    QuadBasicColoniesGenerator,
+)
+from thants.generators.food import BasicFoodGenerator
+
+
+class ThantsDual(Thants):
+    """
+    Environment with two evenly sized and spaced rectangular colonies
+    """
+
+    def __init__(
+        self,
+        dims: tuple[int, int] = (50, 100),
+        n_agents: int = 36,
+        n_signals: int = 2,
+        nest_dims: tuple[int, int] = (5, 5),
+        food_drop_dims: tuple[int, int] = (5, 5),
+        food_drop_interval: int = 50,
+        max_steps: int = 10_000,
+        carry_capacity: float = 1.0,
+        take_food_amount: float = 0.1,
+        deposit_food_amount: float = 0.1,
+        signal_deposit_amount: float = 0.1,
+    ) -> None:
+        """
+        Initialise the environment
+
+        Parameters
+        ----------
+        dims
+            Env dimensions
+        n_agents
+            Number of agents in each colony
+        n_signals
+            Number of signal channels
+        nest_dims
+            Rectangular dimensions of each colonies nest
+        food_drop_dims
+            Rectangular dimensions of food deposits
+        food_drop_interval
+            Interval between new randomly placed food deposits
+        max_steps
+            Maximum environment steps
+        carry_capacity
+            Ant food carrying capacity
+        take_food_amount
+            Max food that can be picked up by an ant in a single step
+        deposit_food_amount
+            Max food that can be dropped by an ant in a single step
+        signal_deposit_amount
+            Amount of signal deposited in a single step
+        """
+        food_generator = BasicFoodGenerator(food_drop_dims, food_drop_interval)
+        super().__init__(
+            dims=dims,
+            colonies_generator=DualBasicColoniesGenerator(
+                n_agents=(n_agents, n_agents), n_signals=n_signals, nest_dims=nest_dims
+            ),
+            food_generator=food_generator,
+            max_steps=max_steps,
+            carry_capacity=carry_capacity,
+            take_food_amount=take_food_amount,
+            deposit_food_amount=deposit_food_amount,
+            signal_deposit_amount=signal_deposit_amount,
+        )
+
+
+class ThantsQuad(Thants):
+    """
+    Environment with four evenly sized and spaced rectangular colonies
+    """
+
+    def __init__(
+        self,
+        dims: tuple[int, int] = (100, 100),
+        n_agents: int = 36,
+        n_signals: int = 2,
+        nest_dims: tuple[int, int] = (5, 5),
+        food_drop_dims: tuple[int, int] = (5, 5),
+        food_drop_interval: int = 100,
+        max_steps: int = 10_000,
+        carry_capacity: float = 1.0,
+        take_food_amount: float = 0.1,
+        deposit_food_amount: float = 0.1,
+        signal_deposit_amount: float = 0.1,
+    ) -> None:
+        """
+        Initialise the environment
+
+        Parameters
+        ----------
+        dims
+            Env dimensions
+        n_agents
+            Number of agents in each colony
+        n_signals
+            Number of signal channels
+        nest_dims
+            Rectangular dimensions of each colonies nest
+        food_drop_dims
+            Rectangular dimensions of food deposits
+        food_drop_interval
+            Interval between new randomly placed food deposits
+        max_steps
+            Maximum environment steps
+        carry_capacity
+            Ant food carrying capacity
+        take_food_amount
+            Max food that can be picked up by an ant in a single step
+        deposit_food_amount
+            Max food that can be dropped by an ant in a single step
+        signal_deposit_amount
+            Amount of signal deposited in a single step
+        """
+        food_generator = BasicFoodGenerator(food_drop_dims, food_drop_interval)
+        super().__init__(
+            dims=dims,
+            colonies_generator=QuadBasicColoniesGenerator(
+                n_agents=(n_agents, n_agents, n_agents, n_agents),
+                n_signals=n_signals,
+                nest_dims=nest_dims,
+            ),
+            food_generator=food_generator,
+            max_steps=max_steps,
+            carry_capacity=carry_capacity,
+            take_food_amount=take_food_amount,
+            deposit_food_amount=deposit_food_amount,
+            signal_deposit_amount=signal_deposit_amount,
+        )

src/thants/generators/colonies/mono.py

@@ -2,7 +2,7 @@
 
 import chex
 
-from thants.generators.colonies.utils import init_colony
+from thants.generators.colonies.utils import BBox, init_colony
 from thants.types import Colony
 
 
@@ -94,6 +94,5 @@ def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> Colony:
         Colony
             Colony state containing ants, signals, and nest states
         """
-        return init_colony(
-            dims, (0, 0), dims, self.nest_dims, self.n_agents, self.n_signals
-        )
+        bounds = BBox(x0=(0, 0), x1=dims)
+        return init_colony(dims, bounds, self.nest_dims, self.n_agents, self.n_signals)

src/thants/generators/colonies/multi.py

@@ -4,7 +4,7 @@
 import chex
 
 from thants.generators.colonies.mono import ColonyGenerator
-from thants.generators.colonies.utils import init_colony
+from thants.generators.colonies.utils import BBox, init_colonies
 from thants.types import Colony
 
 
@@ -82,9 +82,9 @@ def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> Sequence[Colony]
         return [self.generator(dims, key)]
 
 
-class BasicColoniesGenerator(ColoniesGenerator):
+class DualBasicColoniesGenerator(ColoniesGenerator):
     """
-    Basic generator that create 2 evenly spaced colonies
+    Basic generator that create 2 evenly spaced rectangular colonies
     """
 
     def __init__(
@@ -122,21 +122,65 @@ def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> Sequence[Colony]
             List of initialised colonies
         """
         mid = dims[1] // 2
-        return [
-            init_colony(
-                dims,
-                (0, 0),
-                (dims[0], mid),
-                self.nest_dims,
-                self.n_agents[0],
-                self.n_signals,
-            ),
-            init_colony(
-                dims,
-                (0, mid),
-                (dims[0], dims[1]),
-                self.nest_dims,
-                self.n_agents[1],
-                self.n_signals,
-            ),
+        bounds = [
+            BBox(x0=(0, 0), x1=(dims[0], mid)),
+            BBox(x0=(0, mid), x1=dims),
         ]
+        return init_colonies(
+            dims, self.nest_dims, self.n_agents, self.n_signals, bounds
+        )
+
+
+class QuadBasicColoniesGenerator(ColoniesGenerator):
+    """
+    Basic generator that create 4 evenly spaced rectangular colonies
+    """
+
+    def __init__(
+        self,
+        n_agents: tuple[int, int, int, int],
+        n_signals: int,
+        nest_dims: tuple[int, int],
+    ) -> None:
+        """
+        Initialise a basic generator
+
+        Parameters
+        ----------
+        n_agents
+            Number of agents in each colony
+        n_signals
+            Number of colony signal-channels
+        nest_dims
+            Rectangular nest dimensions
+        """
+        self.nest_dims = nest_dims
+        super().__init__(n_agents, n_signals)
+
+    def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> Sequence[Colony]:
+        """
+        Initialise the pair of colonies
+
+        Parameters
+        ----------
+        dims
+            Dimensions of the environment
+        key
+            JAX random key
+
+        Returns
+        -------
+        Sequence[Colony]
+            List of initialised colonies
+        """
+
+        mid = (dims[0] // 2, dims[1] // 2)
+        bounds = [
+            BBox(x0=(0, 0), x1=mid),
+            BBox(x0=mid, x1=dims),
+            BBox(x0=(mid[0], 0), x1=(dims[0], mid[1])),
+            BBox(x0=(0, mid[1]), x1=(mid[0], dims[1])),
+        ]
+        return init_colonies(
+            dims, self.nest_dims, self.n_agents, self.n_signals, bounds
+        )