diff --git a/README.md b/README.md index 0933af2..1540ab6 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ pip install thants #### Single Colony ```python -from thants.mono import ThantsMonoColony +from thants.envs import ThantsMonoColony import jax env = ThantsMonoColony(dims=(50, 50)) @@ -57,10 +57,10 @@ 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 of arrays/structs. +different sizes), so actions and observations are list/tuples of arrays/structs. ```python -from thants.multi import ThantsMultiColony +from thants.envs import ThantsMultiColony import jax import jax.numpy as jnp @@ -101,23 +101,25 @@ in the same colony. #### Colonies -The state of each individual colony has the following components: +The state of the ant colonies is represented by a single struct: -- *Ants*: Individual ants in the colony have several components: +- *Ants*: Individual ants themselves have several components: - *Positions*: 2d indices of ant positions on the environment grid. - *Carrying*: The amount of food being carried by each ant. - *Health*: Ant health (currently unused). -- *Nest*: 2d array indicating if a cell is designated as a nest for this colony. -- *Signals*: 3d array of signal deposits at each cell. Signals have multiple channels - to facilitate communication between ants (i.e. the first dimension of the array - is the signal channel). +- *Colony Index*: The index of the colony each ant belongs to +- *Nests*: 2d array indicating the index of the colony each cell belongs to + (`0` in the case a cell is not the nest of any colony). +- *Signals*: 4d array of signal deposits at each cell for each colony. Signals have + multiple channels to facilitate communication between ants (i.e. the 2nd dimension + of the array is the signal channel). #### Environment -The state of the environment then consists of the colony/colonies and state shared +The state of the environment then consists of the colonies and state shared by all the colonies -- *Colony/Colonies*: Colony or list of colony states +- *Colonies*: Ant colonies state - *Food*: 2d array representing the amount of food deposited at each cell - *Terrain*: Array of flags indicating if a cell can be occupied by an ant. This allows obstacles to be placed on the environment. @@ -136,10 +138,10 @@ Each step of the environment performs the following update to the state: by the colony) The behaviour of the dynamics of signals can be customised by implementing the -`thants.common.signals.SignalPropagator` base class and passing it when initialising -the environment. +[`thants.signals.SignalPropagator`](src/thants/signals.py) base class and +passing it when initialising the environment. -#### Generators +#### State Generators The initialisation of the environment can be customised by implementing the relevant base class and passing them to the environment. @@ -165,18 +167,17 @@ individually made for the local neighbourhood of each ant, i.e. the 8 surroundin cells on the environment grid, and their own cell: - `ants`: Flag indicating if a cell in the neighbourhood is occupied by an ant, - for the single colony environment this has shape `[n-ants, 9]`, but in the multi-colony - case this is `[n-ants, n-colonies, 9]` where the second dimensions indicates the - individual colonies. + with the shape `[n-ants, n-colonies, 9]` where the second dimensions indicates the + individual colonies. The ants from the same colony will always be on the first row. - `signals`: Signal deposits in the neighbourhood (across all channels), signals are observed individually for each colony. - `food`: Food deposits within the neighbourhood. -- `nest`: Flag indicating if a neighbouring cell is designated as a nest. +- `nest`: Flag indicating if a neighbouring cell is designated as a nest for the ants colony. - `terrain`: Flag indicating if a neighbouring cell can be occupied. - `carrying`: Amount of food currently being carried by each ant. ### Rewards By default, rewards are granted to ants when they deposit food on their colonies nest. -Reward signals customised by implementing the respective base classes -`thants.mono.rewards.RewardFn` or `thants.multi.rewards.RewardFn`. +Reward signals can be customised by implementing the respective base classes +[`thants.rewards.RewardFn`](src/thants/rewards.py). diff --git a/src/thants/common/actions.py b/src/thants/actions.py similarity index 96% rename from src/thants/common/actions.py rename to src/thants/actions.py index b57a681..6b46baf 100644 --- a/src/thants/common/actions.py +++ b/src/thants/actions.py @@ -1,7 +1,7 @@ import chex import jax.numpy as jnp -from thants.common.types import Actions, SignalActions +from thants.types import Actions, SignalActions def derive_actions( diff --git a/src/thants/common/rewards.py b/src/thants/common/rewards.py deleted file mode 100644 index 7f586a6..0000000 --- a/src/thants/common/rewards.py +++ /dev/null @@ -1,44 +0,0 @@ -from typing import Optional - -import chex -import jax.numpy as jnp - - -def delivered_food( - nest: chex.Array, - 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 - - Parameters - ---------- - nest - Array of nest flags - pos - Array of ant positions - carrying_before - 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 - - 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 diff --git a/src/thants/common/steps.py b/src/thants/common/steps.py deleted file mode 100644 index b3066ed..0000000 --- a/src/thants/common/steps.py +++ /dev/null @@ -1,77 +0,0 @@ -""" -Environment update steps -""" -import chex -import jax.numpy as jnp - -from .types import SignalActions - - -def deposit_signals( - signals: chex.Array, pos: chex.Array, deposits: SignalActions -) -> chex.Array: - """ - Deposit signals at ant locations - - Parameters - ---------- - signals - Signal state - pos - Ant positions - deposits - Amount to deposit by each ant - - Returns - ------- - chex.Array - Updated signal state - """ - new_signals = signals.at[deposits.channel, pos[:, 0], pos[:, 1]].add( - deposits.amount - ) - return new_signals - - -def update_food( - food: chex.Array, - pos: chex.Array, - take: chex.Array, - deposit: chex.Array, - carrying: chex.Array, - capacity: float, -) -> tuple[chex.Array, chex.Array]: - """ - Update food piles due to ant actions - - Parameters - ---------- - food - Food deposit state array - pos - Ant positions - take - Food pick-up action amounts for individual ants - deposit - Food pick-up deposit amounts for individual ants - carrying - Current ant carrying amounts - capacity - Ant carrying capacity - - Returns - ------- - tuple[chex.Array, chex.Array] - Tuple containing - - - Update food deposit state - - Updated ant carrying amounts - """ - available_food = food.at[pos[:, 0], pos[:, 1]].get() - available_capacity = capacity - carrying - taken_food = jnp.minimum(jnp.minimum(available_food, take), available_capacity) - deposited = jnp.minimum(deposit, carrying) - new_food = food.at[pos[:, 0], pos[:, 1]].add(deposited) - new_food = new_food.at[pos[:, 0], pos[:, 1]].subtract(taken_food) - new_carrying = carrying + taken_food - deposited - return new_food, new_carrying diff --git a/src/thants/envs/__init__.py b/src/thants/envs/__init__.py new file mode 100644 index 0000000..1626d9a --- /dev/null +++ b/src/thants/envs/__init__.py @@ -0,0 +1,2 @@ +from thants.envs.mono import ThantsMonoColony +from thants.envs.multi import ThantsMultiColony diff --git a/src/thants/envs/mono.py b/src/thants/envs/mono.py new file mode 100644 index 0000000..f8caa13 --- /dev/null +++ b/src/thants/envs/mono.py @@ -0,0 +1,224 @@ +from functools import cached_property +from typing import Optional, Sequence + +import chex +from jumanji import Environment, specs +from jumanji.types import TimeStep +from jumanji.viewer import Viewer +from matplotlib.animation import FuncAnimation + +from thants.envs.multi import ThantsMultiColony +from thants.generators.colonies.mono import ( + BasicColonyGenerator, + ColonyGenerator, +) +from thants.generators.colonies.multi import SingleColonyWrapper +from thants.generators.food import FoodGenerator +from thants.generators.terrain import TerrainGenerator +from thants.rewards import RewardFn +from thants.signals import SignalPropagator +from thants.types import Observations, State + + +class ThantsMonoColony(Environment): + """Environment with a single colony""" + + def __init__( + self, + dims: tuple[int, int] = (50, 50), + colony_generator: Optional[ColonyGenerator] = None, + food_generator: Optional[FoodGenerator] = None, + terrain_generator: Optional[TerrainGenerator] = None, + signal_dynamics: Optional[SignalPropagator] = None, + reward_fn: Optional[RewardFn] = None, + viewer: Optional[Viewer[State]] = None, + 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 + 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 + signal-channels and 5x5 nest. + food_generator + Food state generator and updater. By default, initialises a + `BasicFoodGenerator` that creates 5x5 rectangular patches of food at + fixed intervals. + terrain_generator + Terrain generator. By default, initialises a `OpenTerrainGenerator` + that initialises a map with no obstacles. + signal_dynamics + Signal propagation functionality. By default, implements a + `BasicSignalPropagator`. + reward_fn + Reward function, by default rewards are supplied when an ant drops food on + the nest. + viewer + Environment visualiser. By default, initialises a viewer using a Matplotlib + backend. + max_steps + Maximum environment steps + carry_capacity + Maximum ant carrying capacity + take_food_amount + Amount of (attempted) food taken by a take food action + deposit_food_amount + Amount of (attempted) food deposited by a deposit food action + signal_deposit_amount + Amount of signal deposited by the deposit signal action + """ + colony_generator = colony_generator or BasicColonyGenerator(25, 2, (5, 5)) + colony_generator = SingleColonyWrapper(colony_generator) + self.env = ThantsMultiColony( + dims=dims, + colonies_generator=colony_generator, + food_generator=food_generator, + terrain_generator=terrain_generator, + signal_dynamics=signal_dynamics, + reward_fn=reward_fn, + viewer=viewer, + 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, + ) + super().__init__() + + def reset(self, key: chex.PRNGKey) -> tuple[State, TimeStep[Observations]]: + """ + Reset the environment state + + Parameters + ---------- + key + JAX random key + + Returns + ------- + tuple[State, TimeStep] + Tuple containing new environment state, and initial timestep + """ + state, timestep = self.env.reset(key) + return state, timestep[0] + + def step( + self, state: State, actions: chex.Array + ) -> tuple[State, TimeStep[Observations]]: + """ + Update the state of the environment + + Update performs the following steps + + - Unwrap actions into state updates + - Apply position updates + - Apply food pick-up/deposit updates + - Dissipate and propagate signals + - Apply signal deposit actions + - Clear any food returned to the nest + + Parameters + ---------- + state + Current environment state + actions + Integer array of individual ant actions + + Returns + ------- + tuple[State, TimeStep] + Tuple containing new state and TimeStep + """ + state, timestep = self.env.step(state, [actions]) + return state, timestep[0] + + @cached_property + def num_agents(self) -> int: + return self.env.num_agents[0] + + @cached_property + def observation_spec(self) -> specs.Spec[Observations]: + """ + Observation specification + + The observation consists of several components: + + - `[n-agents, 1, 9]` view of ants in the local vicinity + - `[n-agents, 9]` view of food deposits in local vicinity + - `[n-agents, n-signals, 9]` view of signals in the local vicinity + - `[n-agents, 1, 9]` view indicating nest locations in the vicinity + - `[n_agents,]` amount of food being carried by ants + + Returns + ------- + ObservationSpec + """ + return self.env.observation_spec[0] + + @cached_property + def action_spec(self) -> specs.BoundedArray: + """ + Action specification + + Actions are given by an array of integers indicating the discrete action + to be taken by each ant. + + Returns + ------- + ActionSpec + """ + return self.env.action_spec[0] + + @cached_property + def reward_spec(self) -> specs.Array: + """ + Reward specification + + Array of individual rewards for each ant agent + + Returns + ------- + RewardSpec + """ + return self.env.reward_spec[0] + + def render(self, state: State) -> None: + """Render a frame of the environment for a given state using matplotlib. + + Args: + state: State object. + """ + self.env.render(state) + + def animate( + self, + states: Sequence[State], + interval: int = 100, + save_path: Optional[str] = None, + ) -> FuncAnimation: + """Create an animation from a sequence of environment states. + + Args: + states: sequence of environment states corresponding to consecutive + timesteps. + interval: delay between frames in milliseconds. + save_path: the path where the animation file should be saved. If it + is None, the plot will not be saved. + + Returns: + Animation that can be saved as a GIF, MP4, or rendered with HTML. + """ + return self.env.animate(states, interval=interval, save_path=save_path) + + def close(self) -> None: + """Perform any necessary cleanup.""" + self.env.close() diff --git a/src/thants/multi/env.py b/src/thants/envs/multi.py similarity index 85% rename from src/thants/multi/env.py rename to src/thants/envs/multi.py index b25403d..3908e0a 100644 --- a/src/thants/multi/env.py +++ b/src/thants/envs/multi.py @@ -1,5 +1,5 @@ from functools import cached_property, partial -from typing import Optional, Sequence, Tuple +from typing import Optional, Sequence import chex import jax @@ -9,30 +9,26 @@ from jumanji.viewer import Viewer from matplotlib.animation import FuncAnimation -from thants.common.actions import derive_actions -from thants.common.generators.food import BasicFoodGenerator, FoodGenerator -from thants.common.generators.terrain import ( - OpenTerrainGenerator, - TerrainGenerator, -) -from thants.common.signals import BasicSignalPropagator, SignalPropagator -from thants.common.specs import ( - get_action_spec, - get_observation_spec, - get_reward_spec, -) -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 ( +from thants.actions import derive_actions +from thants.generators.colonies.multi 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, deposit_signals, merge_colonies -from thants.multi.types import Colonies, State -from thants.multi.viewer import ThantsMultiColonyViewer +from thants.generators.food import BasicFoodGenerator, FoodGenerator +from thants.generators.terrain import OpenTerrainGenerator, TerrainGenerator +from thants.observations import observations_from_state +from thants.rewards import DeliveredFoodRewards, RewardFn +from thants.signals import BasicSignalPropagator, SignalPropagator +from thants.specs import get_action_spec, get_observation_spec, get_reward_spec +from thants.steps import ( + clear_nest, + deposit_signals, + merge_colonies, + update_food, + update_positions, +) +from thants.types import Ants, Colonies, Observations, State +from thants.viewer import ThantsMultiColonyViewer class ThantsMultiColony(Environment): @@ -100,7 +96,7 @@ def __init__( self.signal_deposit_amount = signal_deposit_amount self.max_steps = max_steps self._colonies_generator = colonies_generator or BasicColoniesGenerator( - 25, 2, (5, 5) + [25, 25], 2, (5, 5) ) self._food_generator = food_generator or BasicFoodGenerator((5, 5), 100, 1.0) self._terrain_generator = terrain_generator or OpenTerrainGenerator() @@ -111,7 +107,9 @@ def __init__( self._viewer = viewer or ThantsMultiColonyViewer() super().__init__() - def reset(self, key: chex.PRNGKey) -> Tuple[State, list[TimeStep[Observations]]]: + def reset( + self, key: chex.PRNGKey + ) -> tuple[State, Sequence[TimeStep[Observations]]]: """ Reset the environment state @@ -122,8 +120,9 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, list[TimeStep[Observations]]] Returns ------- - tuple[State, TimeStep] - Tuple containing new environment state, and initial timestep + tuple[State, Sequence[TimeStep]] + Tuple containing new environment state, and initial timesteps + for each colony """ key, colony_key, food_key, terrain_key = jax.random.split(key, num=4) colonies = self._colonies_generator(self.dims, colony_key) @@ -145,8 +144,8 @@ def reset(self, key: chex.PRNGKey) -> Tuple[State, list[TimeStep[Observations]]] return state, time_steps def step( - self, state: State, actions: list[chex.Array] - ) -> Tuple[State, list[TimeStep[Observations]]]: + self, state: State, actions: Sequence[chex.Array] + ) -> tuple[State, Sequence[TimeStep[Observations]]]: """ Update the state of the environment @@ -240,15 +239,19 @@ def step( return new_state, timestep @cached_property - def num_agents(self) -> list[int]: + def num_agents(self) -> Sequence[int]: return self._colonies_generator.n_agents + @cached_property + def num_colonies(self) -> int: + return len(self.num_agents) + @cached_property def num_actions(self) -> int: return 7 + self._colonies_generator.n_signals @cached_property - def observation_spec(self) -> list[specs.Spec[Observations]]: + def observation_spec(self) -> Sequence[specs.Spec[Observations]]: """ List of observation specifications for each colony @@ -262,17 +265,21 @@ def observation_spec(self) -> list[specs.Spec[Observations]]: Returns ------- - list[ObservationSpec] + Sequence[ObservationSpec] + Sequence of observation specs for each colony """ return [ get_observation_spec( - n, self._colonies_generator.n_signals, self.carry_capacity + n, + self._colonies_generator.n_signals, + self.carry_capacity, + num_colonies=self.num_colonies, ) for n in self.num_agents ] @cached_property - def action_spec(self) -> list[specs.BoundedArray]: + def action_spec(self) -> Sequence[specs.BoundedArray]: """ List of action specifications for each colony @@ -281,7 +288,7 @@ def action_spec(self) -> list[specs.BoundedArray]: Returns ------- - list[ActionSpec] + Sequence[ActionSpec] """ return [ get_action_spec(n, self._colonies_generator.n_signals) @@ -289,7 +296,7 @@ def action_spec(self) -> list[specs.BoundedArray]: ] @cached_property - def reward_spec(self) -> list[specs.Array]: + def reward_spec(self) -> Sequence[specs.Array]: """ List of reward specifications for each colony @@ -297,7 +304,7 @@ def reward_spec(self) -> list[specs.Array]: Returns ------- - list[RewardSpec] + Sequence[RewardSpec] """ return [get_reward_spec(n) for n in self.num_agents] diff --git a/src/thants/common/__init__.py b/src/thants/generators/__init__.py similarity index 100% rename from src/thants/common/__init__.py rename to src/thants/generators/__init__.py diff --git a/src/thants/common/generators/__init__.py b/src/thants/generators/colonies/__init__.py similarity index 100% rename from src/thants/common/generators/__init__.py rename to src/thants/generators/colonies/__init__.py diff --git a/src/thants/mono/colony_generator.py b/src/thants/generators/colonies/mono.py similarity index 96% rename from src/thants/mono/colony_generator.py rename to src/thants/generators/colonies/mono.py index 2669931..316f04c 100644 --- a/src/thants/mono/colony_generator.py +++ b/src/thants/generators/colonies/mono.py @@ -2,8 +2,8 @@ import chex -from thants.common.types import Colony -from thants.common.utils import init_colony +from thants.generators.colonies.utils import init_colony +from thants.types import Colony class ColonyGenerator(abc.ABC): diff --git a/src/thants/multi/colonies_generator.py b/src/thants/generators/colonies/multi.py similarity index 60% rename from src/thants/multi/colonies_generator.py rename to src/thants/generators/colonies/multi.py index 1b33ca8..b9548b4 100644 --- a/src/thants/multi/colonies_generator.py +++ b/src/thants/generators/colonies/multi.py @@ -1,9 +1,11 @@ import abc +from typing import Sequence import chex -from thants.common.types import Colony -from thants.common.utils import init_colony +from thants.generators.colonies.mono import ColonyGenerator +from thants.generators.colonies.utils import init_colony +from thants.types import Colony class ColoniesGenerator(abc.ABC): @@ -11,7 +13,7 @@ class ColoniesGenerator(abc.ABC): Base generator of multiple colonies """ - def __init__(self, n_agents: list[int], n_signals: int) -> None: + def __init__(self, n_agents: Sequence[int], n_signals: int) -> None: """ Initialise base attributes @@ -26,7 +28,7 @@ def __init__(self, n_agents: list[int], n_signals: int) -> None: self.n_agents = n_agents @abc.abstractmethod - def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> list[Colony]: + def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> Sequence[Colony]: """ Initialise colonies @@ -39,18 +41,54 @@ def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> list[Colony]: Returns ------- - list[Colony] + Sequence[Colony] List of agent initial colony states """ +class SingleColonyWrapper(ColoniesGenerator): + """ + Wrapper class around a single colony generator + """ + + def __init__(self, generator: ColonyGenerator) -> None: + """ + Initialise the wrapper + + Parameters + ---------- + generator + Single colony generator + """ + self.generator = generator + super().__init__([generator.n_agents], generator.n_signals) + + def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> Sequence[Colony]: + """ + Initialise the colony + + Parameters + ---------- + dims + Environment dimensions + key + JAX random key + + Returns + ------- + Sequence[Colony] + List of agent initial colony states + """ + return [self.generator(dims, key)] + + class BasicColoniesGenerator(ColoniesGenerator): """ Basic generator that create 2 evenly spaced colonies """ def __init__( - self, n_agents: int, n_signals: int, nest_dims: tuple[int, int] + self, n_agents: tuple[int, int], n_signals: int, nest_dims: tuple[int, int] ) -> None: """ Initialise a basic generator @@ -65,9 +103,9 @@ def __init__( Rectangular nest dimensions """ self.nest_dims = nest_dims - super().__init__([n_agents, n_agents], n_signals) + super().__init__(n_agents, n_signals) - def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> list[Colony]: + def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> Sequence[Colony]: """ Initialise the pair of colonies @@ -80,7 +118,7 @@ def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> list[Colony]: Returns ------- - list[Colony] + Sequence[Colony] List of initialised colonies """ mid = dims[1] // 2 diff --git a/src/thants/common/utils.py b/src/thants/generators/colonies/utils.py similarity index 58% rename from src/thants/common/utils.py rename to src/thants/generators/colonies/utils.py index 719a070..4c39be4 100644 --- a/src/thants/common/utils.py +++ b/src/thants/generators/colonies/utils.py @@ -2,11 +2,8 @@ import chex import jax.numpy as jnp -from matplotlib import color_sequences -from matplotlib.axes import Axes -from matplotlib.figure import Figure -from thants.common.types import Ants, Colony, ColorScheme +from thants.types import Ants, Colony def get_rectangular_indices(rec_dims: tuple[int, int]) -> chex.Array: @@ -82,52 +79,3 @@ def init_colony( signals = jnp.zeros((n_signals, *dims)) return Colony(ants=ants, signals=signals, nest=nest) - - -def format_plot( - fig: Figure, ax: Axes, env_dims: tuple[float, float] -) -> tuple[Figure, Axes]: - """ - Format an environment plot, remove ticks and bound to the environment dimensions. - - Parameters - ---------- - fig - Matplotlib figure - ax - Matplotlib axes - env_dims - Environment dimensions - - Returns - ------- - tuple[Figure, Axes] - Formatted figure and axes - """ - ax.set_xticks([]) - ax.set_yticks([]) - ax.set_xlim(-0.5, env_dims[1] - 0.5) - ax.set_ylim(-0.5, env_dims[0] - 0.5) - - return fig, ax - - -def get_color_scheme(color_sequence: str, n_colonies: int) -> ColorScheme: - """ - Get a environment visualisation colour scheme from a matplotlib sequence - - Parameters - ---------- - color_sequence - Matplotlib color-sequence name - n_colonies - Number of colonies to visualise - - Returns - ------- - ColorScheme - Environment visualisation color-scheme - """ - colors = color_sequences[color_sequence] - colors = jnp.array([(*i, 1.0) for i in colors[: 3 + n_colonies]]) - return ColorScheme(terrain=colors[:2], food=colors[2], ants=colors[3:]) diff --git a/src/thants/common/generators/food.py b/src/thants/generators/food.py similarity index 98% rename from src/thants/common/generators/food.py rename to src/thants/generators/food.py index 5437ffc..2ae276f 100644 --- a/src/thants/common/generators/food.py +++ b/src/thants/generators/food.py @@ -4,7 +4,7 @@ import jax import jax.numpy as jnp -from thants.common.utils import get_rectangular_indices +from thants.generators.colonies.utils import get_rectangular_indices class FoodGenerator(abc.ABC): diff --git a/src/thants/common/generators/terrain.py b/src/thants/generators/terrain.py similarity index 100% rename from src/thants/common/generators/terrain.py rename to src/thants/generators/terrain.py diff --git a/src/thants/mono/__init__.py b/src/thants/mono/__init__.py deleted file mode 100644 index aafd760..0000000 --- a/src/thants/mono/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .env import ThantsMonoColony diff --git a/src/thants/mono/env.py b/src/thants/mono/env.py deleted file mode 100644 index 8026b7c..0000000 --- a/src/thants/mono/env.py +++ /dev/null @@ -1,308 +0,0 @@ -from functools import cached_property, partial -from typing import Optional, Sequence, Tuple - -import chex -import jax -from jumanji import Environment, specs -from jumanji.types import TimeStep, restart, termination, transition -from jumanji.viewer import Viewer -from matplotlib.animation import FuncAnimation - -from thants.common.actions import derive_actions -from thants.common.generators.food import BasicFoodGenerator, FoodGenerator -from thants.common.generators.terrain import ( - OpenTerrainGenerator, - TerrainGenerator, -) -from thants.common.signals import BasicSignalPropagator, SignalPropagator -from thants.common.specs import ( - get_action_spec, - get_observation_spec, - get_reward_spec, -) -from thants.common.steps import deposit_signals, update_food -from thants.common.types import Ants, Colony, Observations -from thants.mono.colony_generator import BasicColonyGenerator, ColonyGenerator -from thants.mono.observations import observations_from_state -from thants.mono.rewards import DeliveredFoodRewards, RewardFn -from thants.mono.steps import clear_nest, update_positions -from thants.mono.types import State -from thants.mono.viewer import ThantsViewer - - -class ThantsMonoColony(Environment): - """ - Thants single-colony environment - """ - - def __init__( - self, - dims: tuple[int, int] = (50, 50), - colony_generator: Optional[ColonyGenerator] = None, - food_generator: Optional[FoodGenerator] = None, - terrain_generator: Optional[TerrainGenerator] = None, - signal_dynamics: Optional[SignalPropagator] = None, - reward_fn: Optional[RewardFn] = None, - viewer: Optional[Viewer[State]] = None, - 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 - 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 - signal-channels and 5x5 nest. - food_generator - Food state generator and updater. By default, initialises a - `BasicFoodGenerator` that creates 5x5 rectangular patches of food at - fixed intervals. - terrain_generator - Terrain generator. By default, initialises a `OpenTerrainGenerator` - that initialises a map with no obstacles. - signal_dynamics - Signal propagation functionality. By default, implements a - `BasicSignalPropagator`. - reward_fn - Reward function, by default rewards are supplied when an ant drops food on - the nest. - viewer - Environment visualiser. By default, initialises a viewer using a Matplotlib - backend. - max_steps - Maximum environment steps - carry_capacity - Maximum ant carrying capacity - take_food_amount - Amount of (attempted) food taken by a take food action - deposit_food_amount - Amount of (attempted) food deposited by a deposit food action - signal_deposit_amount - Amount of signal deposited by the deposit signal action - """ - self.dims = dims - self.carry_capacity = carry_capacity - self.take_food_amount = take_food_amount - self.deposit_food_amount = deposit_food_amount - self.signal_deposit_amount = signal_deposit_amount - self.max_steps = max_steps - self._colony_generator = colony_generator or BasicColonyGenerator(25, 2, (5, 5)) - self._food_generator = food_generator or BasicFoodGenerator((5, 5), 100, 1.0) - self._terrain_generator = terrain_generator or OpenTerrainGenerator() - self._signal_dynamics = signal_dynamics or BasicSignalPropagator( - decay_rate=0.002, dissipation_rate=0.2 - ) - self._reward_fn = reward_fn or DeliveredFoodRewards() - self._viewer = viewer or ThantsViewer() - super().__init__() - - def reset(self, key: chex.PRNGKey) -> Tuple[State, TimeStep[Observations]]: - """ - Reset the environment state - - Parameters - ---------- - key - JAX random key - - Returns - ------- - tuple[State, TimeStep] - Tuple containing new environment state, and initial timestep - """ - key, colony_key, food_key, terrain_key = jax.random.split(key, num=4) - colony = self._colony_generator(self.dims, colony_key) - food = self._food_generator.init(self.dims, food_key) - # For safety clear any food placed on a nest - food = clear_nest(colony.nest, food) - terrain = self._terrain_generator(self.dims, terrain_key) - state = State( - step=0, - key=key, - colony=colony, - food=food, - terrain=terrain, - ) - observations = observations_from_state(state) - time_step = restart( - observation=observations, shape=(self._colony_generator.n_agents,) - ) - return state, time_step - - def step( - self, state: State, actions: chex.Array - ) -> Tuple[State, TimeStep[Observations]]: - """ - Update the state of the environment - - Update performs the following steps - - - Unwrap actions into state updates - - Apply position updates - - Apply food pick-up/deposit updates - - Dissipate and propagate signals - - Apply signal deposit actions - - Clear any food returned to the nest - - Parameters - ---------- - state - Current environment state - actions - Integer array of individual ant actions - - Returns - ------- - tuple[State, TimeStep] - Tuple containing new state and TimeStep - """ - key, food_key, signals_key = jax.random.split(state.key, num=3) - - # Unwrap 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, state.colony.ants.pos, state.terrain, actions.movements - ) - - # Pick up and drop-off food - new_food, new_carrying = update_food( - state.food, - new_pos, - actions.take_food, - actions.deposit_food, - state.colony.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, state.colony.signals) - # Deposit signals - new_signals = deposit_signals(new_signals, new_pos, actions.deposit_signals) - # Clear food dropped on the nest - new_food = clear_nest(state.colony.nest, new_food) - # Construct new state - ants = Ants(pos=new_pos, health=state.colony.ants.health, carrying=new_carrying) - colony = Colony(ants=ants, signals=new_signals, nest=state.colony.nest) - new_state = State( - step=state.step + 1, - key=key, - colony=colony, - food=new_food, - terrain=state.terrain, - ) - # Observations - observations = observations_from_state(new_state) - # Rewards - rewards = self._reward_fn(old_state=state, new_state=new_state) - timestep = jax.lax.cond( - state.step >= self.max_steps, - partial(termination, shape=(self.num_agents,)), - partial(transition, shape=(self.num_agents,)), - rewards, - observations, - ) - return new_state, timestep - - @cached_property - def num_agents(self) -> int: - return self._colony_generator.n_agents - - @cached_property - def num_actions(self) -> int: - return 7 + self._colony_generator.n_signals - - @cached_property - def observation_spec(self) -> specs.Spec[Observations]: - """ - Observation specification - - The observation consists of several components: - - - `[n-agents, 9]` view of ants in the local vicinity - - `[n-agents, 9]` view of food deposits in local vicinity - - `[n-agents, n-signals, 9]` view of signals in the local vicinity - - `[n-agents, n-signals, 9]` view indicating nest locations in the vicinity - - `[n_agents,]` amount of food being carried by ants - - Returns - ------- - ObservationSpec - """ - return get_observation_spec( - self.num_agents, self._colony_generator.n_signals, self.carry_capacity - ) - - @cached_property - def action_spec(self) -> specs.BoundedArray: - """ - Action specification - - Actions are given by an array of integers indicating the discrete action - to be taken by each ant. - - Returns - ------- - ActionSpec - """ - return get_action_spec(self.num_agents, self._colony_generator.n_signals) - - @cached_property - def reward_spec(self) -> specs.Array: - """ - Reward specification - - Array of rewards for each ant agent - - Returns - ------- - RewardSpec - """ - return get_reward_spec(self.num_agents) - - def render(self, state: State) -> None: - """Render a frame of the environment for a given state using matplotlib. - - Args: - state: State object. - """ - self._viewer.render(state) - - def animate( - self, - states: Sequence[State], - interval: int = 100, - save_path: Optional[str] = None, - ) -> FuncAnimation: - """Create an animation from a sequence of environment states. - - Args: - states: sequence of environment states corresponding to consecutive - timesteps. - interval: delay between frames in milliseconds. - save_path: the path where the animation file should be saved. If it - is None, the plot will not be saved. - - Returns: - Animation that can be saved as a GIF, MP4, or rendered with HTML. - """ - return self._viewer.animate(states, interval=interval, save_path=save_path) - - def close(self) -> None: - """Perform any necessary cleanup.""" - self._viewer.close() diff --git a/src/thants/mono/observations.py b/src/thants/mono/observations.py deleted file mode 100644 index 4ad41c4..0000000 --- a/src/thants/mono/observations.py +++ /dev/null @@ -1,65 +0,0 @@ -import chex -import jax -import jax.numpy as jnp - -from thants.common.types import Observations -from thants.mono.types import State - - -def observations_from_state(state: State) -> Observations: - """ - Generate individual agent observations from state - - Parameters - ---------- - state - Environment state - - Returns - ------- - Observations - Struct containing observation components - - - Local neighbourhood with flags indicating neighbouring ants - - Food amounts in the local neighbourhood - - Local neighbourhood indicating if a cell is a nest - - Deposited signals in the local neighbourhood - - The food currently carried by each ant - - Passable terrain in the local neighbourhood - """ - dims = state.colony.nest.shape - idxs = jnp.indices((3, 3)) - idxs = idxs.swapaxes(0, 2).reshape(9, 2) - 1 - dims_arr = jnp.array([dims]) - - view_idxs = jax.vmap(lambda x: (idxs + x) % dims_arr)(state.colony.ants.pos) - - 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() - - occupation = jnp.zeros(dims, dtype=float) - occupation = occupation.at[ - state.colony.ants.pos[:, 0], state.colony.ants.pos[:, 1] - ].set(1.0) - - ants = jax.vmap(get_view, in_axes=(None, 0))(occupation, view_idxs) - food = jax.vmap(get_view, in_axes=(None, 0))(state.food, view_idxs) - signals = jax.vmap(get_signals, in_axes=(None, 0))(state.colony.signals, view_idxs) - nest = jax.vmap(get_view, in_axes=(None, 0))(state.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=state.colony.ants.carrying, - terrain=terrain, - ) diff --git a/src/thants/mono/rewards.py b/src/thants/mono/rewards.py deleted file mode 100644 index 9394519..0000000 --- a/src/thants/mono/rewards.py +++ /dev/null @@ -1,80 +0,0 @@ -import abc - -import chex -import jax.numpy as jnp - -from thants.common.rewards import delivered_food -from thants.mono.types import State - - -class RewardFn(abc.ABC): - """ - Base reward function - """ - - @abc.abstractmethod - def __call__(self, old_state: State, new_state: State) -> chex.Array: - """ - Generates rewards from old and new states - - Parameters - ---------- - old_state - State at the start of the step - new_state - State at the end of the step - - Returns - ------- - chex.Array - Array of individual agent rewards - """ - - -class NullRewardFn(RewardFn): - """ - Dummy reward function returning 0 rewards - """ - - def __call__(self, old_state: State, new_state: State) -> chex.Array: - """ - Generates fixed 0 rewards for all agents - - Parameters - ---------- - old_state - State at the start of the step - new_state - State at the end of the step - - Returns - ------- - chex.Array - Array of individual agent rewards - """ - return jnp.zeros_like(old_state.colony.ants.health) - - -class DeliveredFoodRewards(RewardFn): - def __call__(self, old_state: State, new_state: State) -> chex.Array: - """ - Provide rewards when ants deposit food on the nest - - Parameters - ---------- - old_state - State at the start of the step - new_state - State at the end of the step - - Returns - ------- - chex.Array - Array of individual agent rewards - """ - return delivered_food( - new_state.colony.nest, - new_state.colony.ants.pos, - old_state.colony.ants.carrying, - new_state.colony.ants.carrying, - ) diff --git a/src/thants/mono/steps.py b/src/thants/mono/steps.py deleted file mode 100644 index b2f8dbe..0000000 --- a/src/thants/mono/steps.py +++ /dev/null @@ -1,59 +0,0 @@ -""" -Environment update steps -""" -import chex -import jax.numpy as jnp - - -def update_positions( - dims: tuple[int, int], pos: chex.Array, terrain: chex.Array, updates: chex.Array -) -> chex.Array: - """ - Update agent positions, checking for collisions - - Parameters - ---------- - dims - Environment dimensions - pos - Ant positions - terrain - Environment terrain - updates - Array of position updates - - Returns - ------- - chex.Array - Updated agent positions - """ - dims_arr = jnp.array([dims]) - new_pos = (pos + updates) % dims_arr - x = jnp.concatenate([pos[:, 1], new_pos[:, 1]]) - y = jnp.concatenate([pos[:, 0], new_pos[:, 0]]) - idxs = y * dims[1] + x - occupation = jnp.bincount(idxs, length=dims[0] * dims[1]).reshape(*dims) - move_occupied = occupation.at[new_pos[:, 0], new_pos[:, 1]].get() - 1 - passable = terrain.at[new_pos[:, 0], new_pos[:, 1]].get() - move_available = jnp.logical_and(move_occupied < 1, passable) - new_pos = jnp.where(move_available[:, jnp.newaxis], new_pos, pos) - return new_pos - - -def clear_nest(nest: chex.Array, food: chex.Array) -> chex.Array: - """ - Remove any food deposited on the nest - - Parameters - ---------- - nest - Array of nest flags - food - Food deposits state - - Returns - ------- - chex.Array - Update food state - """ - return jnp.where(nest, 0.0, food) diff --git a/src/thants/mono/types.py b/src/thants/mono/types.py deleted file mode 100644 index 6396fd1..0000000 --- a/src/thants/mono/types.py +++ /dev/null @@ -1,29 +0,0 @@ -from typing import TYPE_CHECKING - -import chex - -from thants.common.types import Colony - -if TYPE_CHECKING: - from dataclasses import dataclass -else: - from chex import dataclass - - -@dataclass -class State: - """ - Environment state - - step: Environment step - key: JAX random key - colony: Ant colony state - food: Environment food deposit state - terrain: Flag indicating if a cell is passable by ants - """ - - step: int - key: chex.PRNGKey - colony: Colony - food: chex.Array # [*env-size] - terrain: chex.Array # [*env-size] diff --git a/src/thants/mono/viewer.py b/src/thants/mono/viewer.py deleted file mode 100644 index b66c263..0000000 --- a/src/thants/mono/viewer.py +++ /dev/null @@ -1,171 +0,0 @@ -from typing import Optional, Sequence, Tuple - -import chex -import jax -import jax.numpy as jnp -import matplotlib.animation -import matplotlib.pyplot as plt -from jumanji.viewer import MatplotlibViewer -from matplotlib.image import AxesImage -from numpy.typing import NDArray - -from thants.common.types import ColorScheme -from thants.common.utils import format_plot, get_color_scheme -from thants.mono.types import State - - -def _draw_env( - state: State, colors: ColorScheme -) -> tuple[chex.Array, chex.Array, chex.Array]: - terrain = state.terrain.astype(int) - terrain = colors.terrain.at[terrain].get() - nest_colors = jnp.clip(colors.ants + 0.1, 0.0, 1.0) - nest = state.colony.nest[:, :, jnp.newaxis] * nest_colors[jnp.newaxis] - food = jnp.full((*state.food.shape, 4), colors.food) - return terrain, nest, food - - -@jax.jit -def _draw_ants(state: State, colors: ColorScheme) -> chex.Array: - dims = state.food.shape - ants = jnp.zeros((*dims, 4)) - ants = ants.at[state.colony.ants.pos[:, 0], state.colony.ants.pos[:, 1]].set( - colors.ants[0] - ) - return ants - - -class ThantsViewer(MatplotlibViewer[State]): - def __init__( - self, - name: str = "thants", - render_mode: str = "human", - color_sequence: str = "tab20", - ) -> None: - """ - Thants environment visualiser using Matplotlib - - Parameters - ---------- - name - Plot name, default ``thants`` - render_mode - Default ``human`` - color_sequence - Matplotlib colour sequence to sample from - """ - self.color_sequence = color_sequence - super().__init__(name, render_mode) - - def _set_figure_size(self, dims: tuple[int, int]) -> None: - longest = max(dims[0], dims[1]) - f_dims = (10.0 * dims[1] / longest, 10.0 * dims[0] / longest) - self.figure_size = f_dims - - def render( - self, state: State, save_path: Optional[str] = None - ) -> Optional[NDArray]: - """Render a frame of the environment for a given state using matplotlib - - Parameters - ---------- - state - State object containing the current dynamics of the environment - save_path - Optional path to save the rendered environment image to - - Returns - ------- - Array or None - RGB array if the render_mode is ``rgb_array`` - """ - self._clear_display() - dims = state.food.shape - self._set_figure_size(dims) - fig, ax = self._get_fig_ax(padding=0.01) - ax.clear() - fig, ax = format_plot(fig, ax, dims) - self._draw(ax, state) - - if save_path: - fig.savefig(save_path, bbox_inches="tight", pad_inches=0.2) - - return self._display(fig) - - def animate( - self, states: Sequence[State], interval: int, save_path: Optional[str] = None - ) -> matplotlib.animation.FuncAnimation: - """Create an animation from a sequence of states - - Parameters - ---------- - states - Sequence of ``State`` corresponding to subsequent timesteps - interval - Delay between frames in milliseconds, default to 200 - save_path - The path where the animation file should be saved. If it is None, - the plot will not be saved - - Returns - ------- - FuncAnimation - Animation object that can be saved as a GIF, MP4, or rendered with HTML - """ - if not states: - raise ValueError(f"The states argument has to be non-empty, got {states}.") - - env_dims = states[0].terrain.shape - self._set_figure_size(env_dims) - fig, ax = self._get_fig_ax(name_suffix="_animation", show=False, padding=0.01) - fig, ax = format_plot(fig, ax, env_dims) - plt.close(fig=fig) - - colors, ants_img, food_img = self._draw(ax, states[0]) - - def make_frame(state: State) -> tuple[AxesImage, AxesImage]: - step_ants = _draw_ants(state, colors) - ants_img.set_data(step_ants) - food_img.set_alpha(jnp.clip(state.food, 0.0, 1.0)) - return ants_img, food_img - - self._animation = matplotlib.animation.FuncAnimation( - fig, - make_frame, - frames=states, - interval=interval, - blit=True, - ) - - if save_path: - self._animation.save(save_path) - - return self._animation - - def _draw( - self, ax: plt.Axes, state: State - ) -> tuple[ColorScheme, AxesImage, AxesImage]: - colors = get_color_scheme(self.color_sequence, 1) - - terrain, nest, food = _draw_env(state, colors) - ants = _draw_ants(state, colors) - - ax.imshow(terrain) - ax.imshow(nest) - - food_img = ax.imshow(food, alpha=state.food) - ants_img = ax.imshow(ants) - - return colors, ants_img, food_img - - def _get_fig_ax( - self, - name_suffix: Optional[str] = None, - show: bool = True, - padding: float = 0.05, - **fig_kwargs: str, - ) -> Tuple[plt.Figure, plt.Axes]: - fig, ax = super()._get_fig_ax( - name_suffix=name_suffix, show=show, padding=padding - ) - return fig, ax diff --git a/src/thants/multi/__init__.py b/src/thants/multi/__init__.py deleted file mode 100644 index 7348008..0000000 --- a/src/thants/multi/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .env import ThantsMultiColony diff --git a/src/thants/multi/types.py b/src/thants/multi/types.py deleted file mode 100644 index a401216..0000000 --- a/src/thants/multi/types.py +++ /dev/null @@ -1,46 +0,0 @@ -from typing import TYPE_CHECKING - -import chex - -from thants.common.types import Ants - -if TYPE_CHECKING: - from dataclasses import dataclass -else: - from chex import dataclass - - -@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-colonies, n-channels, *env-size] - nests: chex.Array # [*env-size] - - -@dataclass -class State: - """ - Environment state - - step: Environment step - key: JAX random key - colonies: List ant colonies - food: Environment food deposit state - terrain: Flag indicating if a cell is passable by ants - """ - - step: int - key: chex.PRNGKey - colonies: Colonies - food: chex.Array # [*env-size] - terrain: chex.Array # [*env-size] diff --git a/src/thants/multi/observations.py b/src/thants/observations.py similarity index 80% rename from src/thants/multi/observations.py rename to src/thants/observations.py index 7bc700b..3141faf 100644 --- a/src/thants/multi/observations.py +++ b/src/thants/observations.py @@ -1,15 +1,16 @@ +from typing import Sequence + import chex import jax import jax.numpy as jnp import numpy as np -from thants.common.types import Observations -from thants.multi.types import State +from thants.types import Observations, State def observations_from_state( - colony_sizes: list[int], state: State -) -> list[Observations]: + colony_sizes: Sequence[int], state: State +) -> Sequence[Observations]: """ Generate individual agent observations from state for each colony @@ -22,7 +23,7 @@ def observations_from_state( Returns ------- - list[Observations] + Sequence[Observations] List of structs containing observation components for each colony - Local neighbourhood with flags indicating neighbouring ants @@ -33,6 +34,7 @@ def observations_from_state( - Local environment terrain """ n_colonies = len(colony_sizes) + n_signals = state.colonies.signals.shape[1] dims = state.food.shape dims_arr = jnp.array([dims]) idxs = jnp.indices((3, 3)) @@ -44,8 +46,8 @@ 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(i: int, arr: chex.Array, x: chex.Array) -> chex.Array: - return arr.at[i, :, x[:, 0], x[:, 1]].get() + def get_signals(arr: chex.Array, i: int, x: chex.Array) -> chex.Array: + return arr.at[i, jnp.arange(n_signals)[:, jnp.newaxis], x[:, 0], x[:, 1]].get() def get_nest(i: int, arr: chex.Array, x: chex.Array) -> chex.Array: return arr.at[x[:, 0], x[:, 1]].get() == (i + 1) @@ -63,17 +65,17 @@ def get_nest(i: int, arr: chex.Array, x: chex.Array) -> chex.Array: ) 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 + signals = jax.vmap(get_signals, in_axes=(None, 0, 0))( + state.colonies.signals, state.colonies.colony_idx, view_idxs ) nest = jax.vmap(get_nest, in_axes=(0, None, 0))( state.colonies.colony_idx, state.colonies.nests, view_idxs - ) + ).astype(float) terrain = jax.vmap(get_view, in_axes=(None, 0))(state.terrain, view_idxs).astype( float ) - boundaries = [0] + colony_sizes + boundaries = [0, *colony_sizes] boundaries = np.array(boundaries) boundaries = np.cumsum(boundaries) @@ -81,7 +83,7 @@ def get_nest(i: int, arr: chex.Array, x: chex.Array) -> chex.Array: Observations( ants=ants[a:b], food=food[a:b], - signals=signals[a:, b], + signals=signals[a:b], nest=nest[a:b], terrain=terrain[a:b], carrying=state.colonies.ants.carrying[a:b], diff --git a/src/thants/multi/rewards.py b/src/thants/rewards.py similarity index 57% rename from src/thants/multi/rewards.py rename to src/thants/rewards.py index 18191f8..378ec4e 100644 --- a/src/thants/multi/rewards.py +++ b/src/thants/rewards.py @@ -1,15 +1,55 @@ import abc +from typing import Optional, Sequence 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 +from thants.types import State + + +def delivered_food( + nest: chex.Array, + 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 + + Parameters + ---------- + nest + Array of nest flags + pos + Array of ant positions + carrying_before + 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 + + 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 + + +def _get_boundaries(colony_sizes: Sequence[int]) -> np.typing.NDArray: + boundaries = [0, *colony_sizes] boundaries = np.array(boundaries) boundaries = np.cumsum(boundaries) return boundaries @@ -18,8 +58,8 @@ def _get_boundaries(colony_sizes: list[int]) -> np.typing.NDArray: class RewardFn(abc.ABC): @abc.abstractmethod def __call__( - self, colony_sizes: list[int], old_state: State, new_state: State - ) -> list[chex.Array]: + self, colony_sizes: Sequence[int], old_state: State, new_state: State + ) -> Sequence[chex.Array]: """ Generate individual ant rewards for each colony @@ -41,8 +81,8 @@ def __call__( class NullRewardFn(RewardFn): def __call__( - self, colony_sizes: list[int], old_state: State, new_state: State - ) -> list[chex.Array]: + self, colony_sizes: Sequence[int], old_state: State, new_state: State + ) -> Sequence[chex.Array]: """ Assigns 0 reward to all agents @@ -57,7 +97,7 @@ def __call__( Returns ------- - list[chex.Array] + Sequence[chex.Array] List of rewards arrays per colony """ return [jnp.zeros(n) for n in colony_sizes] @@ -65,8 +105,8 @@ def __call__( class DeliveredFoodRewards(RewardFn): def __call__( - self, colony_sizes: list[int], old_state: State, new_state: State - ) -> list[chex.Array]: + self, colony_sizes: Sequence[int], old_state: State, new_state: State + ) -> Sequence[chex.Array]: """ Assigns rewards for ants depositing food on their own nest @@ -81,7 +121,7 @@ def __call__( Returns ------- - list[chex.Array] + Sequence[chex.Array] List of rewards arrays per colony """ boundaries = _get_boundaries(colony_sizes) diff --git a/src/thants/common/signals.py b/src/thants/signals.py similarity index 100% rename from src/thants/common/signals.py rename to src/thants/signals.py diff --git a/src/thants/common/specs.py b/src/thants/specs.py similarity index 96% rename from src/thants/common/specs.py rename to src/thants/specs.py index 9c16b80..d41fce8 100644 --- a/src/thants/common/specs.py +++ b/src/thants/specs.py @@ -3,7 +3,7 @@ import jax.numpy as jnp from jumanji import specs -from thants.common.types import Observations +from thants.types import Observations def get_observation_spec( @@ -23,6 +23,8 @@ def get_observation_spec( Number of signal channels carry_capacity Ant carrying capacity + num_colonies + Optional in the case of a multi-colony environment Returns ------- diff --git a/src/thants/multi/steps.py b/src/thants/steps.py similarity index 67% rename from src/thants/multi/steps.py rename to src/thants/steps.py index a2b38c7..a24c777 100644 --- a/src/thants/multi/steps.py +++ b/src/thants/steps.py @@ -1,75 +1,21 @@ -""" -Environment update steps -""" +from typing import Sequence + import chex 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( - occupation: chex.Array, pos: chex.Array, new_pos: chex.Array, terrain: chex.Array -) -> chex.Array: - move_occupied = occupation.at[new_pos[:, 0], new_pos[:, 1]].get() - 1 - passable = terrain.at[new_pos[:, 0], new_pos[:, 1]].get() - move_available = jnp.logical_and(move_occupied < 1, passable) - new_pos = jnp.where(move_available[:, jnp.newaxis], new_pos, pos) - return new_pos - - -def update_positions( - dims: tuple[int, int], - pos: list[chex.Array], - terrain: chex.Array, - updates: list[chex.Array], -) -> list[chex.Array]: - """ - Update agent positions - - Parameters - ---------- - dims - Environment dimensions - pos - List of colony positions - terrain - Environment terrain - updates - Array of colony position updates - - Returns - ------- - tuple[chex.Array, chex.Array] - Updated agent positions - """ - dims_arr = jnp.array([dims]) - - new_pos = [(p + update) % dims_arr for p, update in zip(pos, updates)] - - x = jnp.concatenate([a[:, 1] for b in zip(pos, new_pos) for a in b]) - y = jnp.concatenate([a[:, 0] for b in zip(pos, new_pos) for a in b]) - - idxs = y * dims[1] + x - - occupation = jnp.bincount(idxs, length=dims[0] * dims[1]).reshape(*dims) - - new_pos = [_move(occupation, p, np, terrain) for p, np in zip(pos, new_pos)] - - return new_pos +from thants.types import Ants, Colonies, Colony, SignalActions def update_food( food: chex.Array, - pos: list[chex.Array], - take: list[chex.Array], - deposit: list[chex.Array], - carrying: list[chex.Array], + pos: chex.Array, + take: chex.Array, + deposit: chex.Array, + carrying: chex.Array, capacity: float, -) -> tuple[chex.Array, list[chex.Array]]: +) -> tuple[chex.Array, chex.Array]: """ - Update food deposits due to ant actions + Update food piles due to ant actions Parameters ---------- @@ -88,24 +34,55 @@ def update_food( Returns ------- - tuple[chex.Array, list[chex.Array]] + tuple[chex.Array, chex.Array] Tuple containing - Update food deposit state - - List of updated ant carrying amounts for each colony + - Updated ant carrying amounts """ + available_food = food.at[pos[:, 0], pos[:, 1]].get() + available_capacity = capacity - carrying + taken_food = jnp.minimum(jnp.minimum(available_food, take), available_capacity) + deposited = jnp.minimum(deposit, carrying) + new_food = food.at[pos[:, 0], pos[:, 1]].add(deposited) + new_food = new_food.at[pos[:, 0], pos[:, 1]].subtract(taken_food) + new_carrying = carrying + taken_food - deposited + return new_food, new_carrying - updates = [ - _update_food(food, _pos, _take, _deposit, _carrying, capacity) - for _pos, _take, _deposit, _carrying in zip(pos, take, deposit, carrying) - ] - d_food = jnp.stack([x[0] - food for x in updates], axis=0) - d_food = jnp.sum(d_food, axis=0) - new_food = food + d_food - new_carrying = [x[1] for x in updates] +def update_positions( + dims: tuple[int, int], pos: chex.Array, terrain: chex.Array, updates: chex.Array +) -> chex.Array: + """ + Update agent positions, checking for collisions - return new_food, new_carrying + Parameters + ---------- + dims + Environment dimensions + pos + Ant positions + terrain + Environment terrain + updates + Array of position updates + + Returns + ------- + chex.Array + Updated agent positions + """ + dims_arr = jnp.array([dims]) + new_pos = (pos + updates) % dims_arr + x = jnp.concatenate([pos[:, 1], new_pos[:, 1]]) + y = jnp.concatenate([pos[:, 0], new_pos[:, 0]]) + idxs = y * dims[1] + x + occupation = jnp.bincount(idxs, length=dims[0] * dims[1]).reshape(*dims) + move_occupied = occupation.at[new_pos[:, 0], new_pos[:, 1]].get() - 1 + passable = terrain.at[new_pos[:, 0], new_pos[:, 1]].get() + move_available = jnp.logical_and(move_occupied < 1, passable) + new_pos = jnp.where(move_available[:, jnp.newaxis], new_pos, pos) + return new_pos def clear_nest(nests: chex.Array, food: chex.Array) -> chex.Array: @@ -127,7 +104,7 @@ def clear_nest(nests: chex.Array, food: chex.Array) -> chex.Array: return jnp.where(nests > 0, 0.0, food) -def merge_colonies(colonies: list[Colony]) -> Colonies: +def merge_colonies(colonies: Sequence[Colony]) -> Colonies: """ Merge a list of colonies into a single state diff --git a/src/thants/common/types.py b/src/thants/types.py similarity index 73% rename from src/thants/common/types.py rename to src/thants/types.py index 5da1151..23fa2e8 100644 --- a/src/thants/common/types.py +++ b/src/thants/types.py @@ -30,6 +30,42 @@ class Colony: nest: chex.Array # [*env-size] +@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-colonies, n-channels, *env-size] + nests: chex.Array # [*env-size] + + +@dataclass +class State: + """ + Environment state + + step: Environment step + key: JAX random key + colonies: List ant colonies + food: Environment food deposit state + terrain: Flag indicating if a cell is passable by ants + """ + + step: int + key: chex.PRNGKey + colonies: Colonies + food: chex.Array # [*env-size] + terrain: chex.Array # [*env-size] + + @dataclass class SignalActions: """ diff --git a/src/thants/multi/viewer.py b/src/thants/viewer.py similarity index 79% rename from src/thants/multi/viewer.py rename to src/thants/viewer.py index e3c88a5..83effa3 100644 --- a/src/thants/multi/viewer.py +++ b/src/thants/viewer.py @@ -6,12 +6,60 @@ import matplotlib.animation import matplotlib.pyplot as plt from jumanji.viewer import MatplotlibViewer +from matplotlib import color_sequences from matplotlib.image import AxesImage from numpy.typing import NDArray -from thants.common.types import ColorScheme -from thants.common.utils import format_plot, get_color_scheme -from thants.multi.types import State +from thants.types import ColorScheme, State + + +def format_plot( + fig: plt.Figure, ax: plt.Axes, env_dims: tuple[float, float] +) -> tuple[plt.Figure, plt.Axes]: + """ + Format an environment plot, remove ticks and bound to the environment dimensions. + + Parameters + ---------- + fig + Matplotlib figure + ax + Matplotlib axes + env_dims + Environment dimensions + + Returns + ------- + tuple[Figure, Axes] + Formatted figure and axes + """ + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_xlim(-0.5, env_dims[1] - 0.5) + ax.set_ylim(-0.5, env_dims[0] - 0.5) + + return fig, ax + + +def get_color_scheme(color_sequence: str, n_colonies: int) -> ColorScheme: + """ + Get a environment visualisation colour scheme from a matplotlib sequence + + Parameters + ---------- + color_sequence + Matplotlib color-sequence name + n_colonies + Number of colonies to visualise + + Returns + ------- + ColorScheme + Environment visualisation color-scheme + """ + colors = color_sequences[color_sequence] + colors = jnp.array([(*i, 1.0) for i in colors[: 3 + n_colonies]]) + return ColorScheme(terrain=colors[:2], food=colors[2], ants=colors[3:]) def _draw_env( @@ -148,7 +196,7 @@ def make_frame(state: State) -> tuple[AxesImage, AxesImage]: def _draw( self, ax: plt.Axes, state: State ) -> tuple[ColorScheme, AxesImage, AxesImage]: - colors = get_color_scheme(self.color_sequence, len(state.colonies)) + colors = get_color_scheme(self.color_sequence, state.colonies.signals.shape[0]) terrain, nests, food = _draw_env(state, colors) ants = _draw_ants(state, colors) diff --git a/tests/test_common/test_actions.py b/tests/test_actions.py similarity index 91% rename from tests/test_common/test_actions.py rename to tests/test_actions.py index 1353327..9a78a21 100644 --- a/tests/test_common/test_actions.py +++ b/tests/test_actions.py @@ -1,7 +1,7 @@ import jax.numpy as jnp -from thants.common.actions import derive_actions -from thants.common.types import Actions, SignalActions +from thants.actions import derive_actions +from thants.types import Actions, SignalActions def test_derive_actions() -> None: diff --git a/tests/test_common/test_rewards.py b/tests/test_common/test_rewards.py deleted file mode 100644 index 3d68470..0000000 --- a/tests/test_common/test_rewards.py +++ /dev/null @@ -1,17 +0,0 @@ -import jax.numpy as jnp - -from thants.common.rewards import delivered_food - - -def test_delivered_food_rewards(): - dims = (3, 1) - nest = jnp.ones(dims, dtype=bool).at[0, 0].set(False) - pos = jnp.array([[0, 0], [1, 0], [2, 0]]) - carrying_before = jnp.array([1.0, 1.0, 1.0]) - carrying_after = jnp.array([0.5, 0.5, 1.0]) - - rewards = delivered_food(nest, pos, carrying_before, carrying_after) - - expected = jnp.array([0.0, 0.5, 0.0]) - - assert jnp.allclose(rewards, expected) diff --git a/tests/test_common/test_update_food.py b/tests/test_common/test_update_food.py deleted file mode 100644 index d57b08a..0000000 --- a/tests/test_common/test_update_food.py +++ /dev/null @@ -1,51 +0,0 @@ -import jax.numpy as jnp -import pytest - -from thants.common.steps import update_food - - -@pytest.mark.parametrize( - ( - "food_pos, food_amount, agent_pos, take, deposit, " - "carrying, expected_food, expected_carrying" - ), - [ - # Attempt to take 0 food - ([(0, 0)], [0.0], [(0, 0)], [1.0], [0.0], [0.0], [0.0, 0.0, 0.0, 0.0], [0.0]), - # Take available amount - ([(0, 0)], [0.1], [(0, 0)], [1.0], [0.0], [0.0], [0.0, 0.0, 0.0, 0.0], [0.1]), - # Take wanted amount - ([(0, 0)], [1.0], [(0, 0)], [0.1], [0.0], [0.0], [0.9, 0.0, 0.0, 0.0], [0.1]), - # Deposit fixed - ([(1, 1)], [0.0], [(1, 1)], [0.0], [0.5], [1.0], [0.0, 0.0, 0.0, 0.5], [0.5]), - # Deposit carrying - ([(1, 1)], [0.0], [(1, 1)], [0.0], [1.0], [0.5], [0.0, 0.0, 0.0, 0.5], [0.0]), - ], -) -def test_update_food( - food_pos: list[tuple[int, int]], - food_amount: list[float], - agent_pos: list[tuple[int, int]], - take: list[float], - deposit: list[float], - carrying: list[float], - expected_food: list[float], - expected_carrying: list[float], -) -> None: - food_pos = jnp.array(food_pos) - food_amount = jnp.array(food_amount) - food = jnp.zeros((2, 2), dtype=float) - food = food.at[food_pos[:, 0], food_pos[:, 1]].set(food_amount) - - agent_pos = jnp.array(agent_pos) - take = jnp.array(take) - deposit = jnp.array(deposit) - carrying = jnp.array(carrying) - - new_food, new_carrying = update_food(food, agent_pos, take, deposit, carrying, 1.0) - - expected_food = jnp.array(expected_food).reshape(2, 2) - expected_carrying = jnp.array(expected_carrying) - - assert jnp.allclose(new_food, expected_food) - assert jnp.allclose(new_carrying, expected_carrying) diff --git a/tests/test_common/__init__.py b/tests/test_envs/__init__.py similarity index 100% rename from tests/test_common/__init__.py rename to tests/test_envs/__init__.py diff --git a/tests/test_mono/test_env.py b/tests/test_envs/test_mono.py similarity index 80% rename from tests/test_mono/test_env.py rename to tests/test_envs/test_mono.py index 6f827bf..b493fb9 100644 --- a/tests/test_mono/test_env.py +++ b/tests/test_envs/test_mono.py @@ -6,10 +6,10 @@ check_env_specs_does_not_smoke, ) -from thants.common.generators.food import BasicFoodGenerator -from thants.common.types import Observations -from thants.mono import ThantsMonoColony -from thants.mono.colony_generator import BasicColonyGenerator +from thants.envs.mono import ThantsMonoColony +from thants.generators.colonies.mono import BasicColonyGenerator +from thants.generators.food import BasicFoodGenerator +from thants.types import Observations @pytest.fixture @@ -27,7 +27,7 @@ def env() -> ThantsMonoColony: def test_env_does_not_smoke(env: ThantsMonoColony) -> None: """Test that we can run an episode without any errors.""" - env.max_steps = 100 + env.max_steps = 20 def select_action(action_key: chex.PRNGKey, _state: Observations) -> chex.Array: return jax.random.choice(action_key, 9, (env.num_agents,)) diff --git a/tests/test_multi/test_env.py b/tests/test_envs/test_multi.py similarity index 77% rename from tests/test_multi/test_env.py rename to tests/test_envs/test_multi.py index da9cbfc..e495801 100644 --- a/tests/test_multi/test_env.py +++ b/tests/test_envs/test_multi.py @@ -2,17 +2,16 @@ import jax.random import pytest -from thants.common.generators.food import BasicFoodGenerator -from thants.common.types import Observations -from thants.multi import ThantsMultiColony -from thants.multi.colonies_generator import BasicColoniesGenerator -from thants.multi.types import State +from thants.envs.multi import ThantsMultiColony +from thants.generators.colonies.multi import BasicColoniesGenerator +from thants.generators.food import BasicFoodGenerator +from thants.types import Observations, State @pytest.fixture def env() -> ThantsMultiColony: dims = (50, 100) - colony_generator = BasicColoniesGenerator(64, 2, (5, 5)) + colony_generator = BasicColoniesGenerator((64, 36), 2, (5, 5)) food_generator = BasicFoodGenerator( (2, 2), 50, diff --git a/tests/test_mono/__init__.py b/tests/test_generators/__init__.py similarity index 100% rename from tests/test_mono/__init__.py rename to tests/test_generators/__init__.py diff --git a/tests/test_common/test_generators.py b/tests/test_generators/test_food_generator.py similarity index 81% rename from tests/test_common/test_generators.py rename to tests/test_generators/test_food_generator.py index 140ff0f..53bf27d 100644 --- a/tests/test_common/test_generators.py +++ b/tests/test_generators/test_food_generator.py @@ -1,6 +1,6 @@ import jax.numpy as jnp -from thants.common.generators.food import BasicFoodGenerator +from thants.generators.food import BasicFoodGenerator def test_food_generator(key) -> None: diff --git a/tests/test_mono/test_colony_generator.py b/tests/test_generators/test_mono_colony_generators.py similarity index 90% rename from tests/test_mono/test_colony_generator.py rename to tests/test_generators/test_mono_colony_generators.py index 15273ab..291af36 100644 --- a/tests/test_mono/test_colony_generator.py +++ b/tests/test_generators/test_mono_colony_generators.py @@ -1,6 +1,6 @@ import jax.numpy as jnp -from thants.mono.colony_generator import BasicColonyGenerator +from thants.generators.colonies.mono import BasicColonyGenerator def test_basic_colony_generator(key) -> None: diff --git a/tests/test_multi/test_colony_generator.py b/tests/test_generators/test_multi_colony_generators.py similarity index 79% rename from tests/test_multi/test_colony_generator.py rename to tests/test_generators/test_multi_colony_generators.py index 9b027db..7849d43 100644 --- a/tests/test_multi/test_colony_generator.py +++ b/tests/test_generators/test_multi_colony_generators.py @@ -1,12 +1,12 @@ import jax.numpy as jnp -from thants.common.types import Colony -from thants.multi.colonies_generator import BasicColoniesGenerator +from thants.generators.colonies.multi import BasicColoniesGenerator +from thants.types import Colony def test_colony_generator(key): dims = (50, 100) - generator = BasicColoniesGenerator(25, 2, (5, 5)) + generator = BasicColoniesGenerator([25, 25], 2, (5, 5)) colonies = generator(dims, key) assert isinstance(colonies, list) assert len(colonies) == 2 diff --git a/tests/test_common/test_utils.py b/tests/test_generators/test_utils.py similarity index 76% rename from tests/test_common/test_utils.py rename to tests/test_generators/test_utils.py index 61d2315..225897c 100644 --- a/tests/test_common/test_utils.py +++ b/tests/test_generators/test_utils.py @@ -1,6 +1,6 @@ import jax.numpy as jnp -from thants.common.utils import get_rectangular_indices +from thants.generators.colonies.utils import get_rectangular_indices def test_rectangular_indices() -> None: diff --git a/tests/test_mono/test_observations.py b/tests/test_mono/test_observations.py deleted file mode 100644 index ca826dc..0000000 --- a/tests/test_mono/test_observations.py +++ /dev/null @@ -1,46 +0,0 @@ -import chex -import jax.numpy as jnp - -from thants.common.types import Ants, Colony, Observations -from thants.mono.observations import observations_from_state -from thants.mono.types import State - - -def test_observations_from_state(key: chex.PRNGKey) -> None: - dims = (3, 2) - - ant_pos = jnp.array([[0, 0], [1, 1]]) - ant_health = jnp.zeros(2) - ant_carry = jnp.zeros(2) - - food = jnp.zeros(dims) - signals = jnp.zeros((2, dims[0], dims[1])) - nest = jnp.zeros(dims) - terrain = jnp.ones(dims, dtype=bool) - - state = State( - step=0, - key=key, - colony=Colony( - ants=Ants( - pos=ant_pos, - health=ant_health, - carrying=ant_carry, - ), - signals=signals, - nest=nest, - ), - food=food, - terrain=terrain, - ) - - observations = observations_from_state(state) - - assert isinstance(observations, Observations) - - assert observations.ants.shape == (2, 9) - assert observations.food.shape == (2, 9) - assert observations.signals.shape == (2, 2, 9) - assert observations.nest.shape == (2, 9) - assert observations.terrain.shape == (2, 9) - assert observations.carrying.shape == (2,) diff --git a/tests/test_mono/test_steps.py b/tests/test_mono/test_steps.py deleted file mode 100644 index 34ccfa3..0000000 --- a/tests/test_mono/test_steps.py +++ /dev/null @@ -1,98 +0,0 @@ -import chex -import jax -import jax.numpy as jnp -import pytest - -from thants.mono.steps import clear_nest, update_positions - - -@pytest.mark.parametrize( - "pos, updates, expected", - [ - # Move - ([(0, 0)], [(1, 1)], [(1, 1)]), - # Wrapped move - ([(0, 0)], [(-1, -1)], [(1, 2)]), - # Move around - ([(0, 0), (1, 1)], [(1, 0), (0, 0)], [(1, 0), (1, 1)]), - # Blocked move - ([(0, 0), (1, 0)], [(1, 0), (0, 0)], [(0, 0), (1, 0)]), - # Blocked move - ([(0, 0), (1, 1)], [(1, 0), (0, -1)], [(0, 0), (1, 1)]), - ], -) -def test_movement( - pos: list[tuple[int, int]], - updates: list[tuple[int, int]], - expected: list[tuple[int, int]], -) -> None: - pos = jnp.array(pos) - updates = jnp.array(updates) - dims = (2, 3) - terrain = jnp.ones(dims, dtype=bool) - - new_pos = update_positions((2, 3), pos, terrain, updates) - - expected = jnp.array(expected) - - assert jnp.allclose(new_pos, expected) - - -def test_fuzz_movement(key: chex.PRNGKey) -> None: - n_agents = 20 - dims = (12, 5) - terrain = jnp.ones(dims, dtype=bool) - n_cells = dims[0] * dims[1] - pos_idx = jax.random.choice(key, n_cells, shape=(n_agents,), replace=False) - pos_x = pos_idx % dims[1] - pos_y = (pos_idx - pos_x) // dims[1] - start_pos = jnp.stack((pos_y, pos_x)).T - movements = jnp.array([[-1, 0], [1, 0], [0, -1], [0, 1], [0, 0]]) - - def step(carry, _): - k, pos = carry - k, move_key = jax.random.split(k) - moves = jax.random.choice(move_key, movements, shape=(n_agents,)) - new_pos = update_positions(dims, pos, terrain, moves) - return (k, new_pos), pos - - (_, final_pos), positions = jax.lax.scan(step, (key, start_pos), None, length=50) - - idxs = positions[:, :, 0] * dims[1] + positions[:, :, 1] - - def bin_count(x): - return jnp.bincount(x, length=n_cells) - - occupation = jax.vmap(bin_count)(idxs) - - assert jnp.max(occupation) == 1 - - -def test_terrain_blocking() -> None: - dims = (3, 3) - terrain = jnp.ones(dims, dtype=bool) - terrain = terrain.at[1, 1].set(False) - pos = jnp.array([[0, 1]]) - - # Should be blocked by terrain - updates = jnp.array([[1, 0]]) - new_pos = update_positions(dims, pos, terrain, updates) - assert jnp.array_equal(pos, new_pos) - - # Should not be blocked by terrain - updates = jnp.array([[0, 1]]) - new_pos = update_positions(dims, pos, terrain, updates) - expected_pos = jnp.array([[0, 2]]) - assert jnp.array_equal(expected_pos, new_pos) - - -def test_clear_nest() -> None: - dims = (2, 1) - - nest = jnp.zeros(dims, dtype=bool).at[0].set(True) - food = jnp.ones(dims) - - new_food = clear_nest(nest, food) - expected = jnp.array([[0.0], [1.0]]) - - assert jnp.allclose(new_food, expected) diff --git a/tests/test_multi/__init__.py b/tests/test_multi/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/test_multi/test_steps.py b/tests/test_multi/test_steps.py deleted file mode 100644 index 321f16e..0000000 --- a/tests/test_multi/test_steps.py +++ /dev/null @@ -1,118 +0,0 @@ -import jax.numpy as jnp - -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: - dims = (5, 5) - - pos = [ - jnp.array([[0, 0], [3, 3]]), - jnp.array([[0, 1]]), - ] - - terrain = jnp.ones(dims, dtype=bool) - - updates = [ - jnp.array([[0, 1], [0, 1]]), - jnp.array([[-1, 0]]), - ] - - new_pos = update_positions( - dims, - pos, - terrain, - updates, - ) - - assert isinstance(new_pos, list) - assert len(new_pos) == 2 - assert new_pos[0].shape == (2, 2) - assert new_pos[1].shape == (1, 2) - - expected_0 = jnp.array([[0, 0], [3, 4]]) - assert jnp.array_equal(new_pos[0], expected_0) - - expected_1 = jnp.array([[4, 1]]) - assert jnp.array_equal(new_pos[1], expected_1) - - -def test_colony_food_update() -> None: - food = jnp.zeros((2, 2)).at[1, 1].set(1.0) - - pos = [ - jnp.array([[0, 0], [1, 1]]), - jnp.array([[0, 1]]), - ] - - take = [jnp.array([1.0, 1.0]), jnp.array([0.0])] - deposit = [jnp.array([0.0, 0.0]), jnp.array([1.0])] - carrying = [jnp.array([0.0, 0.0]), jnp.array(0.5)] - - new_food, new_carrying = update_food(food, pos, take, deposit, carrying, 0.5) - - assert new_food.shape == (2, 2) - expected_food = jnp.array([[0.0, 0.5], [0.0, 0.5]]) - assert jnp.array_equal(new_food, expected_food) - - assert isinstance(carrying, list) - expected_carry_0 = jnp.array([0.0, 0.5]) - assert jnp.allclose(new_carrying[0], expected_carry_0) - expected_carry_1 = jnp.array([0.0]) - assert jnp.allclose(new_carrying[1], expected_carry_1) - - -def test_clear_food() -> None: - dims = (3, 1) - - 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) diff --git a/tests/test_multi/test_observations.py b/tests/test_observations.py similarity index 85% rename from tests/test_multi/test_observations.py rename to tests/test_observations.py index c7c82b9..563430d 100644 --- a/tests/test_multi/test_observations.py +++ b/tests/test_observations.py @@ -1,10 +1,9 @@ import chex import jax.numpy as jnp -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 +from thants.observations import observations_from_state +from thants.steps import merge_colonies +from thants.types import Ants, Colony, State def test_colony_observations(key: chex.PRNGKey) -> None: @@ -57,3 +56,6 @@ def test_colony_observations(key: chex.PRNGKey) -> None: expected_obs_1 = expected_obs_1.at[0, [0, 1], [4, 0]].set(1.0) assert jnp.allclose(observations[1].ants, expected_obs_1) + + assert observations[0].signals.shape == (2, 2, 9) + assert observations[1].signals.shape == (1, 2, 9) diff --git a/tests/test_multi/test_rewards.py b/tests/test_rewards.py similarity index 77% rename from tests/test_multi/test_rewards.py rename to tests/test_rewards.py index c2f9f18..8546cba 100644 --- a/tests/test_multi/test_rewards.py +++ b/tests/test_rewards.py @@ -1,9 +1,22 @@ import jax.numpy as jnp -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 +from thants.rewards import DeliveredFoodRewards, delivered_food +from thants.steps import merge_colonies +from thants.types import Ants, Colony, State + + +def test_delivered_food_rewards(): + dims = (3, 1) + nest = jnp.ones(dims, dtype=bool).at[0, 0].set(False) + pos = jnp.array([[0, 0], [1, 0], [2, 0]]) + carrying_before = jnp.array([1.0, 1.0, 1.0]) + carrying_after = jnp.array([0.5, 0.5, 1.0]) + + rewards = delivered_food(nest, pos, carrying_before, carrying_after) + + expected = jnp.array([0.0, 0.5, 0.0]) + + assert jnp.allclose(rewards, expected) def test_food_reward_function(key): diff --git a/tests/test_common/test_signals.py b/tests/test_signals.py similarity index 82% rename from tests/test_common/test_signals.py rename to tests/test_signals.py index dd2e97a..f04c2b9 100644 --- a/tests/test_common/test_signals.py +++ b/tests/test_signals.py @@ -2,9 +2,9 @@ import jax import jax.numpy as jnp -from thants.common.signals import BasicSignalPropagator -from thants.common.steps import deposit_signals -from thants.common.types import SignalActions +from thants.signals import BasicSignalPropagator +from thants.steps import deposit_signals +from thants.types import SignalActions def test_total_signals_preserved_single(key) -> None: @@ -52,12 +52,13 @@ def step(s: chex.Array, _: None) -> tuple[chex.Array, chex.Array]: def test_deposit_signals() -> None: - signals = jnp.zeros((2, 3, 3), dtype=float) + signals = jnp.zeros((1, 2, 3, 3), dtype=float) pos = jnp.array([[1, 1]]) deposits = SignalActions( channel=jnp.array([0]), amount=jnp.array([0.5]), ) - expected = signals.at[0, 1, 1].set(0.5) - new_signals = deposit_signals(signals, pos, deposits) + colony_idxs = jnp.array([0]) + expected = signals.at[0, 0, 1, 1].set(0.5) + new_signals = deposit_signals(signals, pos, colony_idxs, deposits) assert jnp.allclose(new_signals, expected) diff --git a/tests/test_steps.py b/tests/test_steps.py new file mode 100644 index 0000000..1b441d1 --- /dev/null +++ b/tests/test_steps.py @@ -0,0 +1,188 @@ +import chex +import jax +import jax.numpy as jnp +import pytest + +from thants.steps import ( + clear_nest, + merge_colonies, + update_food, + update_positions, +) +from thants.types import Ants, Colonies, Colony + + +@pytest.mark.parametrize( + ( + "food_pos, food_amount, agent_pos, take, deposit, " + "carrying, expected_food, expected_carrying" + ), + [ + # Attempt to take 0 food + ([(0, 0)], [0.0], [(0, 0)], [1.0], [0.0], [0.0], [0.0, 0.0, 0.0, 0.0], [0.0]), + # Take available amount + ([(0, 0)], [0.1], [(0, 0)], [1.0], [0.0], [0.0], [0.0, 0.0, 0.0, 0.0], [0.1]), + # Take wanted amount + ([(0, 0)], [1.0], [(0, 0)], [0.1], [0.0], [0.0], [0.9, 0.0, 0.0, 0.0], [0.1]), + # Deposit fixed + ([(1, 1)], [0.0], [(1, 1)], [0.0], [0.5], [1.0], [0.0, 0.0, 0.0, 0.5], [0.5]), + # Deposit carrying + ([(1, 1)], [0.0], [(1, 1)], [0.0], [1.0], [0.5], [0.0, 0.0, 0.0, 0.5], [0.0]), + ], +) +def test_update_food( + food_pos: list[tuple[int, int]], + food_amount: list[float], + agent_pos: list[tuple[int, int]], + take: list[float], + deposit: list[float], + carrying: list[float], + expected_food: list[float], + expected_carrying: list[float], +) -> None: + food_pos = jnp.array(food_pos) + food_amount = jnp.array(food_amount) + food = jnp.zeros((2, 2), dtype=float) + food = food.at[food_pos[:, 0], food_pos[:, 1]].set(food_amount) + + agent_pos = jnp.array(agent_pos) + take = jnp.array(take) + deposit = jnp.array(deposit) + carrying = jnp.array(carrying) + + new_food, new_carrying = update_food(food, agent_pos, take, deposit, carrying, 1.0) + + expected_food = jnp.array(expected_food).reshape(2, 2) + expected_carrying = jnp.array(expected_carrying) + + assert jnp.allclose(new_food, expected_food) + assert jnp.allclose(new_carrying, expected_carrying) + + +@pytest.mark.parametrize( + "pos, updates, expected", + [ + # Move + ([(0, 0)], [(1, 1)], [(1, 1)]), + # Wrapped move + ([(0, 0)], [(-1, -1)], [(1, 2)]), + # Move around + ([(0, 0), (1, 1)], [(1, 0), (0, 0)], [(1, 0), (1, 1)]), + # Blocked move + ([(0, 0), (1, 0)], [(1, 0), (0, 0)], [(0, 0), (1, 0)]), + # Blocked move + ([(0, 0), (1, 1)], [(1, 0), (0, -1)], [(0, 0), (1, 1)]), + ], +) +def test_movement( + pos: list[tuple[int, int]], + updates: list[tuple[int, int]], + expected: list[tuple[int, int]], +) -> None: + pos = jnp.array(pos) + updates = jnp.array(updates) + dims = (2, 3) + terrain = jnp.ones(dims, dtype=bool) + + new_pos = update_positions((2, 3), pos, terrain, updates) + + expected = jnp.array(expected) + + assert jnp.allclose(new_pos, expected) + + +def test_fuzz_movement(key: chex.PRNGKey) -> None: + n_agents = 20 + dims = (12, 5) + terrain = jnp.ones(dims, dtype=bool) + n_cells = dims[0] * dims[1] + pos_idx = jax.random.choice(key, n_cells, shape=(n_agents,), replace=False) + pos_x = pos_idx % dims[1] + pos_y = (pos_idx - pos_x) // dims[1] + start_pos = jnp.stack((pos_y, pos_x)).T + movements = jnp.array([[-1, 0], [1, 0], [0, -1], [0, 1], [0, 0]]) + + def step(carry, _): + k, pos = carry + k, move_key = jax.random.split(k) + moves = jax.random.choice(move_key, movements, shape=(n_agents,)) + new_pos = update_positions(dims, pos, terrain, moves) + return (k, new_pos), pos + + (_, final_pos), positions = jax.lax.scan(step, (key, start_pos), None, length=50) + + idxs = positions[:, :, 0] * dims[1] + positions[:, :, 1] + + def bin_count(x): + return jnp.bincount(x, length=n_cells) + + occupation = jax.vmap(bin_count)(idxs) + + assert jnp.max(occupation) == 1 + + +def test_terrain_blocking() -> None: + dims = (3, 3) + terrain = jnp.ones(dims, dtype=bool) + terrain = terrain.at[1, 1].set(False) + pos = jnp.array([[0, 1]]) + + # Should be blocked by terrain + updates = jnp.array([[1, 0]]) + new_pos = update_positions(dims, pos, terrain, updates) + assert jnp.array_equal(pos, new_pos) + + # Should not be blocked by terrain + updates = jnp.array([[0, 1]]) + new_pos = update_positions(dims, pos, terrain, updates) + expected_pos = jnp.array([[0, 2]]) + assert jnp.array_equal(expected_pos, new_pos) + + +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) + + +def test_clear_food() -> None: + dims = (3, 1) + + 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)