From b9c80cbd4360d0bd3d7697b871d11fdea2f078bc Mon Sep 17 00:00:00 2001 From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com> Date: Thu, 16 Oct 2025 23:26:43 +0100 Subject: [PATCH 1/5] Cleanup colony generators --- src/thants/envs/multi.py | 4 +- src/thants/generators/colonies/mono.py | 7 +- src/thants/generators/colonies/multi.py | 84 ++++++++++++++----- src/thants/generators/colonies/utils.py | 63 ++++++++++++-- tests/test_envs/test_multi.py | 4 +- .../test_multi_colony_generators.py | 4 +- 6 files changed, 128 insertions(+), 38 deletions(-) diff --git a/src/thants/envs/multi.py b/src/thants/envs/multi.py index 3908e0a..b806dc6 100644 --- a/src/thants/envs/multi.py +++ b/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 @@ -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) diff --git a/src/thants/generators/colonies/mono.py b/src/thants/generators/colonies/mono.py index 316f04c..dbf0198 100644 --- a/src/thants/generators/colonies/mono.py +++ b/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) diff --git a/src/thants/generators/colonies/multi.py b/src/thants/generators/colonies/multi.py index b9548b4..2169da4 100644 --- a/src/thants/generators/colonies/multi.py +++ b/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 + ) diff --git a/src/thants/generators/colonies/utils.py b/src/thants/generators/colonies/utils.py index 4c39be4..7593858 100644 --- a/src/thants/generators/colonies/utils.py +++ b/src/thants/generators/colonies/utils.py @@ -1,4 +1,6 @@ import math +from dataclasses import dataclass +from typing import Sequence import chex import jax.numpy as jnp @@ -25,10 +27,17 @@ def get_rectangular_indices(rec_dims: tuple[int, int]) -> chex.Array: return idxs +@dataclass +class BBox: + """Rectangular region bounding box""" + + x0: tuple[int, int] + x1: tuple[int, int] + + def init_colony( dims: tuple[int, int], - x0: tuple[int, int], - x1: tuple[int, int], + bounds: BBox, nest_dims: tuple[int, int], n_agents: int, n_signals: int, @@ -40,10 +49,8 @@ def init_colony( ---------- dims Environment dimensions - x0 - Ids of the origin of the rectangular region - x1 - Dimensions of the rectangular region + bounds + Bounding box of the rectangular region nest_dims Rectangular nest dimensions n_agents @@ -56,8 +63,8 @@ def init_colony( Colony Initialised colony """ - x0 = jnp.array(x0) - x1 = jnp.array(x1) + x0 = jnp.array(bounds.x0) + x1 = jnp.array(bounds.x1) dims = jnp.array(dims) centre = (x0 + ((x1 - x0) // 2))[jnp.newaxis] d = math.ceil(math.sqrt(n_agents)) @@ -79,3 +86,43 @@ def init_colony( signals = jnp.zeros((n_signals, *dims)) return Colony(ants=ants, signals=signals, nest=nest) + + +def init_colonies( + env_dims: tuple[int, int], + nest_dims: tuple[int, int], + n_agents: Sequence[BBox], + n_signals: int, + bounds: Sequence[BBox], +) -> Sequence[Colony]: + """ + Initialise multiple rectangular colonies + + Parameters + ---------- + env_dims + Environment dimensions + nest_dims + Rectangular nest dimensions + n_agents + Number of agents in each colony + n_signals + Number of signal channels + bounds + Bounding boxes of each colony + + Returns + ------- + Sequence[Colony] + List of initialised colonies + """ + return [ + init_colony( + env_dims, + b, + nest_dims, + n, + n_signals, + ) + for b, n in zip(bounds, n_agents) + ] diff --git a/tests/test_envs/test_multi.py b/tests/test_envs/test_multi.py index e495801..ecf40ec 100644 --- a/tests/test_envs/test_multi.py +++ b/tests/test_envs/test_multi.py @@ -3,7 +3,7 @@ import pytest from thants.envs.multi import ThantsMultiColony -from thants.generators.colonies.multi import BasicColoniesGenerator +from thants.generators.colonies.multi import DualBasicColoniesGenerator from thants.generators.food import BasicFoodGenerator from thants.types import Observations, State @@ -11,7 +11,7 @@ @pytest.fixture def env() -> ThantsMultiColony: dims = (50, 100) - colony_generator = BasicColoniesGenerator((64, 36), 2, (5, 5)) + colony_generator = DualBasicColoniesGenerator((64, 36), 2, (5, 5)) food_generator = BasicFoodGenerator( (2, 2), 50, diff --git a/tests/test_generators/test_multi_colony_generators.py b/tests/test_generators/test_multi_colony_generators.py index 7849d43..0d38e12 100644 --- a/tests/test_generators/test_multi_colony_generators.py +++ b/tests/test_generators/test_multi_colony_generators.py @@ -1,12 +1,12 @@ import jax.numpy as jnp -from thants.generators.colonies.multi import BasicColoniesGenerator +from thants.generators.colonies.multi import DualBasicColoniesGenerator from thants.types import Colony def test_colony_generator(key): dims = (50, 100) - generator = BasicColoniesGenerator([25, 25], 2, (5, 5)) + generator = DualBasicColoniesGenerator([25, 25], 2, (5, 5)) colonies = generator(dims, key) assert isinstance(colonies, list) assert len(colonies) == 2 From 28f37bc679ae68734dd1ca168eb3089ea2f3f7d8 Mon Sep 17 00:00:00 2001 From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com> Date: Fri, 17 Oct 2025 00:12:36 +0100 Subject: [PATCH 2/5] Add preset environments --- README.md | 19 +++-- src/thants/envs/__init__.py | 135 +++++++++++++++++++++++++++++++++- src/thants/envs/mono.py | 6 +- src/thants/envs/multi.py | 2 +- tests/test_envs/test_mono.py | 10 +-- tests/test_envs/test_multi.py | 8 +- 6 files changed, 159 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 1540ab6..c95cea2 100644 --- a/README.md +++ b/README.md @@ -34,11 +34,15 @@ pip install thants #### Single Colony +The single colony environment follow 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
diff --git a/src/thants/envs/__init__.py b/src/thants/envs/__init__.py index 1626d9a..4305f73 100644 --- a/src/thants/envs/__init__.py +++ b/src/thants/envs/__init__.py @@ -1,2 +1,133 @@ -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.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, + ) diff --git a/src/thants/envs/mono.py b/src/thants/envs/mono.py index f8caa13..84c4da9 100644 --- a/src/thants/envs/mono.py +++ b/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, diff --git a/src/thants/envs/multi.py b/src/thants/envs/multi.py index b806dc6..a893d92 100644 --- a/src/thants/envs/multi.py +++ b/src/thants/envs/multi.py @@ -31,7 +31,7 @@ from thants.viewer import ThantsMultiColonyViewer -class ThantsMultiColony(Environment): +class Thants(Environment): """ Thants environment with multiple colonies """ diff --git a/tests/test_envs/test_mono.py b/tests/test_envs/test_mono.py index b493fb9..1afca32 100644 --- a/tests/test_envs/test_mono.py +++ b/tests/test_envs/test_mono.py @@ -6,26 +6,26 @@ check_env_specs_does_not_smoke, ) -from thants.envs.mono import ThantsMonoColony +from thants.envs.mono import ThantsMono from thants.generators.colonies.mono import BasicColonyGenerator from thants.generators.food import BasicFoodGenerator from thants.types import Observations @pytest.fixture -def env() -> ThantsMonoColony: +def env() -> ThantsMono: dims = (50, 50) colony_generator = BasicColonyGenerator(100, 2, (5, 5)) food_generator = BasicFoodGenerator( (2, 2), 50, ) - return ThantsMonoColony( + return ThantsMono( dims=dims, colony_generator=colony_generator, food_generator=food_generator ) -def test_env_does_not_smoke(env: ThantsMonoColony) -> None: +def test_env_does_not_smoke(env: ThantsMono) -> None: """Test that we can run an episode without any errors.""" env.max_steps = 20 @@ -35,6 +35,6 @@ def select_action(action_key: chex.PRNGKey, _state: Observations) -> chex.Array: check_env_does_not_smoke(env, select_action=select_action) -def test_env_specs_do_not_smoke(env: ThantsMonoColony) -> None: +def test_env_specs_do_not_smoke(env: ThantsMono) -> None: """Test that we can access specs without any errors.""" check_env_specs_does_not_smoke(env) diff --git a/tests/test_envs/test_multi.py b/tests/test_envs/test_multi.py index ecf40ec..b2bb72e 100644 --- a/tests/test_envs/test_multi.py +++ b/tests/test_envs/test_multi.py @@ -2,26 +2,26 @@ import jax.random import pytest -from thants.envs.multi import ThantsMultiColony +from thants.envs.multi import Thants from thants.generators.colonies.multi import DualBasicColoniesGenerator from thants.generators.food import BasicFoodGenerator from thants.types import Observations, State @pytest.fixture -def env() -> ThantsMultiColony: +def env() -> Thants: dims = (50, 100) colony_generator = DualBasicColoniesGenerator((64, 36), 2, (5, 5)) food_generator = BasicFoodGenerator( (2, 2), 50, ) - return ThantsMultiColony( + return Thants( dims=dims, colonies_generator=colony_generator, food_generator=food_generator ) -def test_env_does_not_smoke(key: chex.Array, env: ThantsMultiColony) -> None: +def test_env_does_not_smoke(key: chex.Array, env: Thants) -> None: """Test that we can run an episode without any errors.""" env.max_steps = 100 From d6b9b7be2ad492985aa050174b8df4ade76969b1 Mon Sep 17 00:00:00 2001 From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com> Date: Fri, 17 Oct 2025 20:03:35 +0100 Subject: [PATCH 3/5] Refactor and test presets --- src/thants/envs/__init__.py | 132 +------------------------------- src/thants/envs/presets.py | 132 ++++++++++++++++++++++++++++++++ tests/test_envs/test_presets.py | 47 ++++++++++++ 3 files changed, 180 insertions(+), 131 deletions(-) create mode 100644 src/thants/envs/presets.py create mode 100644 tests/test_envs/test_presets.py diff --git a/src/thants/envs/__init__.py b/src/thants/envs/__init__.py index 4305f73..7551fba 100644 --- a/src/thants/envs/__init__.py +++ b/src/thants/envs/__init__.py @@ -1,133 +1,3 @@ from thants.envs.mono import ThantsMono 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, - ) +from thants.envs.presets import ThantsDual, ThantsQuad diff --git a/src/thants/envs/presets.py b/src/thants/envs/presets.py new file mode 100644 index 0000000..2fd2f4b --- /dev/null +++ b/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, + ) diff --git a/tests/test_envs/test_presets.py b/tests/test_envs/test_presets.py new file mode 100644 index 0000000..219324e --- /dev/null +++ b/tests/test_envs/test_presets.py @@ -0,0 +1,47 @@ +from typing import Type + +import chex +import jax +import pytest + +from thants.envs import ThantsDual, ThantsQuad +from thants.types import Observations + + +@pytest.mark.parametrize( + "env_type, dims", + [ + (ThantsDual, (20, 10)), + (ThantsQuad, (20, 20)), + ], +) +def test_env_runs( + key: chex.PRNGKey, + env_type: Type[ThantsDual | ThantsQuad], + dims: tuple[int, int], +) -> None: + n_steps = 10 + n_agents = 9 + n_signals = 2 + + env = env_type(dims=dims, n_agents=n_agents, nest_dims=(3, 3), n_signals=n_signals) + + env.max_steps = 100 + state, _ = env.reset(key) + + def step(_state, _): + k = jax.random.split(_state.key, env.num_colonies) + actions = [jax.random.choice(_k, 9, (n,)) for _k, n in zip(k, env.num_agents)] + _state, timesteps = env.step(_state, actions) + return _state, [t.observation for t in timesteps] + + final_state, obs = jax.lax.scan(step, state, None, length=n_steps) + + assert isinstance(obs, list) + assert all([isinstance(x, Observations)] for x in obs) + + for o in obs: + assert o.food.shape == (n_steps, n_agents, 9) + assert o.signals.shape == (n_steps, n_agents, n_signals, 9) + assert o.terrain.shape == (n_steps, n_agents, 9) + assert o.ants.shape == (n_steps, n_agents, env.num_colonies, 9) From 310085310109f6b1fcb91ca47a17b8f2f79b6168 Mon Sep 17 00:00:00 2001 From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com> Date: Fri, 17 Oct 2025 22:43:44 +0100 Subject: [PATCH 4/5] Extend tests --- src/thants/envs/mono.py | 4 ++ tests/test_envs/test_mono.py | 33 ++++++++++++--- tests/test_envs/test_multi.py | 25 ++++++++--- tests/test_generators/test_food_generator.py | 18 +++++++- .../test_mono_colony_generators.py | 5 ++- .../test_multi_colony_generators.py | 23 ++++++---- tests/test_rewards.py | 2 +- tests/test_viewer.py | 42 +++++++++++++++++++ 8 files changed, 128 insertions(+), 24 deletions(-) create mode 100644 tests/test_viewer.py diff --git a/src/thants/envs/mono.py b/src/thants/envs/mono.py index 84c4da9..a66b7e7 100644 --- a/src/thants/envs/mono.py +++ b/src/thants/envs/mono.py @@ -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] diff --git a/tests/test_envs/test_mono.py b/tests/test_envs/test_mono.py index 1afca32..1716884 100644 --- a/tests/test_envs/test_mono.py +++ b/tests/test_envs/test_mono.py @@ -1,5 +1,6 @@ import chex import jax +import jax.numpy as jnp import pytest from jumanji.testing.env_not_smoke import ( check_env_does_not_smoke, @@ -9,16 +10,16 @@ from thants.envs.mono import ThantsMono from thants.generators.colonies.mono import BasicColonyGenerator from thants.generators.food import BasicFoodGenerator -from thants.types import Observations +from thants.types import Observations, State @pytest.fixture def env() -> ThantsMono: - dims = (50, 50) - colony_generator = BasicColonyGenerator(100, 2, (5, 5)) + dims = (20, 20) + colony_generator = BasicColonyGenerator(16, 2, (5, 5)) food_generator = BasicFoodGenerator( - (2, 2), - 50, + (5, 5), + 5, ) return ThantsMono( dims=dims, colony_generator=colony_generator, food_generator=food_generator @@ -38,3 +39,25 @@ def select_action(action_key: chex.PRNGKey, _state: Observations) -> chex.Array: def test_env_specs_do_not_smoke(env: ThantsMono) -> None: """Test that we can access specs without any errors.""" check_env_specs_does_not_smoke(env) + + +def state_checks(env: ThantsMono, state: State) -> None: + assert isinstance(state, State) + assert state.food.shape == env.dims + assert state.terrain.shape == env.dims + assert state.colonies.ants.pos.shape == (env.num_agents, 2) + assert state.colonies.nests.shape == env.dims + assert jnp.all( + jnp.logical_not(jnp.logical_and(state.food > 0.0, state.colonies.nests > 0)) + ) + + +def test_env_outputs(key: chex.Array, env: ThantsMono) -> None: + state, obs = env.reset(key) + + state_checks(env, state) + + actions = jnp.zeros((env.num_agents,), dtype=int) + state, obs = env.step(state, actions) + + state_checks(env, state) diff --git a/tests/test_envs/test_multi.py b/tests/test_envs/test_multi.py index b2bb72e..428f797 100644 --- a/tests/test_envs/test_multi.py +++ b/tests/test_envs/test_multi.py @@ -1,5 +1,6 @@ import chex -import jax.random +import jax +import jax.numpy as jnp import pytest from thants.envs.multi import Thants @@ -10,11 +11,11 @@ @pytest.fixture def env() -> Thants: - dims = (50, 100) - colony_generator = DualBasicColoniesGenerator((64, 36), 2, (5, 5)) + dims = (20, 40) + colony_generator = DualBasicColoniesGenerator((16, 9), 2, (5, 5)) food_generator = BasicFoodGenerator( - (2, 2), - 50, + (5, 5), + 5, ) return Thants( dims=dims, colonies_generator=colony_generator, food_generator=food_generator @@ -24,6 +25,7 @@ def env() -> Thants: def test_env_does_not_smoke(key: chex.Array, env: Thants) -> None: """Test that we can run an episode without any errors.""" env.max_steps = 100 + n_steps = 50 def step(_state: State, _: None) -> tuple[State, list[Observations]]: k1, k2 = jax.random.split(_state.key, 2) @@ -36,9 +38,20 @@ def step(_state: State, _: None) -> tuple[State, list[Observations]]: state, _ = env.reset(key) - state, obs = jax.lax.scan(step, state, None, 50) + state, obs = jax.lax.scan(step, state, None, n_steps) assert isinstance(state, State) + assert state.food.shape == env.dims + assert state.terrain.shape == env.dims + assert state.colonies.ants.pos.shape == (sum(env.num_agents), 2) + assert state.colonies.nests.shape == env.dims + assert jnp.all( + jnp.logical_not(jnp.logical_and(state.food > 0.0, state.colonies.nests > 0)) + ) + assert isinstance(obs, list) assert len(obs) == 2 assert all([isinstance(x, Observations) for x in obs]) + + for n, o in zip(env.num_agents, obs): + assert o.ants.shape == (n_steps, n, 2, 9) diff --git a/tests/test_generators/test_food_generator.py b/tests/test_generators/test_food_generator.py index 53bf27d..94756b2 100644 --- a/tests/test_generators/test_food_generator.py +++ b/tests/test_generators/test_food_generator.py @@ -1,13 +1,27 @@ +import chex +import jax import jax.numpy as jnp from thants.generators.food import BasicFoodGenerator -def test_food_generator(key) -> None: +def test_food_generator(key: chex.PRNGKey) -> None: env_dims = (5, 10) + k1, k2 = jax.random.split(key) + food_generator = BasicFoodGenerator((2, 2), 10, 0.5) - food = food_generator.init(env_dims, key) + food = food_generator.init(env_dims, k1) + + assert food.shape == env_dims + assert jnp.isclose(jnp.sum(food), 2.0) + + food = food_generator.update(k2, 4, food) assert food.shape == env_dims assert jnp.isclose(jnp.sum(food), 2.0) + + food = food_generator.update(k2, 9, food) + + assert food.shape == env_dims + assert jnp.isclose(jnp.sum(food), 4.0) diff --git a/tests/test_generators/test_mono_colony_generators.py b/tests/test_generators/test_mono_colony_generators.py index 291af36..26013b5 100644 --- a/tests/test_generators/test_mono_colony_generators.py +++ b/tests/test_generators/test_mono_colony_generators.py @@ -6,8 +6,9 @@ def test_basic_colony_generator(key) -> None: env_dims = (5, 10) n_agents = 8 + n_signals = 2 - colony_generator = BasicColonyGenerator(n_agents, 2, (2, 2)) + colony_generator = BasicColonyGenerator(n_agents, n_signals, (2, 2)) colony = colony_generator(env_dims, key) @@ -21,3 +22,5 @@ def test_basic_colony_generator(key) -> None: assert colony.nest.shape == env_dims assert jnp.isclose(jnp.sum(colony.nest), 4) + assert colony.ants.carrying.shape == (n_agents,) + assert colony.signals.shape == (n_signals, *env_dims) diff --git a/tests/test_generators/test_multi_colony_generators.py b/tests/test_generators/test_multi_colony_generators.py index 0d38e12..26e0f38 100644 --- a/tests/test_generators/test_multi_colony_generators.py +++ b/tests/test_generators/test_multi_colony_generators.py @@ -4,20 +4,25 @@ from thants.types import Colony -def test_colony_generator(key): - dims = (50, 100) - generator = DualBasicColoniesGenerator([25, 25], 2, (5, 5)) +def test_colony_generator(key) -> None: + dims = (25, 50) + n_agents = (25, 16) + generator = DualBasicColoniesGenerator(n_agents, 2, (5, 5)) colonies = generator(dims, key) assert isinstance(colonies, list) assert len(colonies) == 2 assert all([isinstance(c, Colony) for c in colonies]) - assert all([c.ants.pos.shape == (25, 2) for c in colonies]) + assert all( + [c.ants.pos.shape == (n, 2) for c, n in zip(colonies, generator.n_agents)] + ) occupation = jnp.zeros(dims, dtype=int) - pos_0 = colonies[0].ants.pos - occupation = occupation.at[pos_0[:, 0], pos_0[:, 1]].add(1) - pos_1 = colonies[1].ants.pos - occupation = occupation.at[pos_1[:, 0], pos_1[:, 1]].add(1) - assert jnp.sum(occupation) == 50 + for n, colony in zip(generator.n_agents, colonies): + pos = colony.ants.pos + assert pos.shape == (n, 2) + assert jnp.sum(colony.nest) == 25 + occupation = occupation.at[pos[:, 0], pos[:, 1]].add(1) + + assert jnp.sum(occupation) == sum(n_agents) assert jnp.max(occupation) == 1 diff --git a/tests/test_rewards.py b/tests/test_rewards.py index 8546cba..a98a21a 100644 --- a/tests/test_rewards.py +++ b/tests/test_rewards.py @@ -5,7 +5,7 @@ from thants.types import Ants, Colony, State -def test_delivered_food_rewards(): +def test_delivered_food_rewards() -> None: dims = (3, 1) nest = jnp.ones(dims, dtype=bool).at[0, 0].set(False) pos = jnp.array([[0, 0], [1, 0], [2, 0]]) diff --git a/tests/test_viewer.py b/tests/test_viewer.py new file mode 100644 index 0000000..835e6c6 --- /dev/null +++ b/tests/test_viewer.py @@ -0,0 +1,42 @@ +import chex +import jax.numpy as jnp +import pytest +from matplotlib.animation import FuncAnimation +from matplotlib.pyplot import Figure + +from thants.envs.multi import Thants +from thants.generators.colonies.multi import DualBasicColoniesGenerator +from thants.generators.food import BasicFoodGenerator + + +@pytest.fixture +def env() -> Thants: + dims = (20, 40) + colony_generator = DualBasicColoniesGenerator((16, 9), 2, (5, 5)) + food_generator = BasicFoodGenerator( + (5, 5), + 5, + ) + return Thants( + dims=dims, colonies_generator=colony_generator, food_generator=food_generator + ) + + +def test_render(monkeypatch, key: chex.PRNGKey, env: Thants) -> None: + monkeypatch.setattr(Figure, "show", lambda _: None) + state, _ = env.reset(key) + env.render(state) + + +def test_animation(key: chex.PRNGKey, env: Thants) -> None: + state, _ = env.reset(key) + + states = [state] + + for _ in range(2): + actions = [jnp.zeros((n,), dtype=int) for n in env.num_agents] + state, _ = env.step(state, actions) + states.append(state) + + animation = env.animate(states) + assert isinstance(animation, FuncAnimation) From c154c7f55a7623f3c96a0fd8b6800553a8035954 Mon Sep 17 00:00:00 2001 From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com> Date: Fri, 17 Oct 2025 22:48:28 +0100 Subject: [PATCH 5/5] Fix typo --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c95cea2..2d82af2 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ pip install thants #### Single Colony -The single colony environment follow the [Jumanji](https://github.com/instadeepai/jumanji) +The single colony environment follows the [Jumanji](https://github.com/instadeepai/jumanji) environment API, with actions provided as an array of individual actions: