Cleanup

README.md

@@ -128,7 +128,7 @@ by all the colonies
 
 - *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
+- *Terrain*: 2d array of flags indicating if a cell can be occupied by an ant. This
   allows obstacles to be placed on the environment.
 
 #### Updates
@@ -170,7 +170,7 @@ case there will be no change in state due to the chosen action.
 ### Observations
 
 Individual agent observations also consist of several components. Observations are
-individually made for the local neighbourhood of each ant, i.e. the 8 surrounding
+individually made for the local neighbourhood of each ant, i.e. the surrounding
 cells on the environment grid, and their own cell:
 
 - `ants`: Flag indicating if a cell in the neighbourhood is occupied by an ant,
@@ -183,8 +183,11 @@ cells on the environment grid, and their own cell:
 - `terrain`: Flag indicating if a neighbouring cell can be occupied.
 - `carrying`: Amount of food currently being carried by each ant.
 
+the number of local cells observed by each agent can be customised with the `view_distance`
+environment parameter.
+
 ### Rewards
 
 By default, rewards are granted to ants when they deposit food on their colonies nest.
-Reward signals can be customised by implementing the respective base classes
+Reward signals can be customised by implementing the respective base class
 [`thants.rewards.RewardFn`](src/thants/rewards.py).

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "thants"
-version = "0.3.3"
+version = "0.3.4"
 description = "Ants! On your GPU! With JAX!"
 authors = ["zombie-einstein <zombie-einstein@proton.me>"]
 license = "MIT"

src/thants/envs/mono.py

@@ -161,12 +161,14 @@ def observation_spec(self) -> specs.Spec[Observations]:
 
         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, 1, n-obs]` view of ants in the local vicinity
+        - `[n-agents, n-obs]` view of food deposits in local vicinity
+        - `[n-agents, n-signals, n-obs]` view of signals in the local vicinity
+        - `[n-agents, 1, n-obs]` view indicating nest locations in the vicinity
         - `[n_agents,]` amount of food being carried by ants
 
+        where `n-obs` is the number of observed cells, dependent on the view distance.
+
         Returns
         -------
         ObservationSpec

src/thants/envs/multi.py

@@ -134,6 +134,7 @@ def reset(  # type: ignore[override]
         merged_colonies = merge_colonies(colonies)
         food = self._food_generator.init(self.dims, food_key)
         terrain = self._terrain_generator(self.dims, terrain_key)
+        food = clear_nest(merged_colonies.nests, terrain, food)
         state = State(
             step=0,
             key=key,
@@ -271,12 +272,14 @@ def observation_spec(  # type: ignore[override]
 
         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, n-obs]` view of ants in the local vicinity
+        - `[n-agents, n-obs]` view of food deposits in local vicinity
+        - `[n-agents, n-signals, n-obs]` view of signals in the local vicinity
+        - `[n-agents, n-signals, n-obs]` view indicating nest locations in the vicinity
         - `[n_agents,]` amount of food being carried by ants
 
+        where `n-obs` is the number of observed cells, dependent on the view distance.
+
         Returns
         -------
         Sequence[ObservationSpec]

src/thants/envs/presets.py

@@ -19,6 +19,7 @@ def __init__(
         nest_dims: tuple[int, int] = (5, 5),
         food_drop_dims: tuple[int, int] = (5, 5),
         food_drop_interval: int = 50,
+        food_decay_rate: float = 0.0,
         max_steps: int = 10_000,
         carry_capacity: float = 1.0,
         take_food_amount: float = 0.1,
@@ -43,6 +44,8 @@ def __init__(
             Rectangular dimensions of food deposits
         food_drop_interval
             Interval between new randomly placed food deposits
+        food_decay_rate
+            Amount any depositied food is reduced each step
         max_steps
             Maximum environment steps
         carry_capacity
@@ -56,7 +59,9 @@ def __init__(
         view_distance
             Number of cells away from an agent observed by each agent
         """
-        food_generator = BasicFoodGenerator(food_drop_dims, food_drop_interval)
+        food_generator = BasicFoodGenerator(
+            food_drop_dims, food_drop_interval, decay_rate=food_decay_rate
+        )
         super().__init__(
             dims=dims,
             colonies_generator=DualBasicColoniesGenerator(
@@ -85,6 +90,7 @@ def __init__(
         nest_dims: tuple[int, int] = (5, 5),
         food_drop_dims: tuple[int, int] = (5, 5),
         food_drop_interval: int = 100,
+        food_decay_rate: float = 0.0,
         max_steps: int = 10_000,
         carry_capacity: float = 1.0,
         take_food_amount: float = 0.1,
@@ -109,6 +115,8 @@ def __init__(
             Rectangular dimensions of food deposits
         food_drop_interval
             Interval between new randomly placed food deposits
+        food_decay_rate
+            Amount any depositied food is reduced each step
         max_steps
             Maximum environment steps
         carry_capacity
@@ -122,7 +130,9 @@ def __init__(
         view_distance
             Number of cells away from an agent observed by each agent
         """
-        food_generator = BasicFoodGenerator(food_drop_dims, food_drop_interval)
+        food_generator = BasicFoodGenerator(
+            food_drop_dims, food_drop_interval, decay_rate=food_decay_rate
+        )
         super().__init__(
             dims=dims,
             colonies_generator=QuadBasicColoniesGenerator(