Add variable view distance

pyproject.toml

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

src/thants/envs/mono.py

@@ -38,6 +38,7 @@ def __init__(
         take_food_amount: float = 0.1,
         deposit_food_amount: float = 0.1,
         signal_deposit_amount: float = 0.1,
+        view_distance: int = 1,
     ) -> None:
         """
         Initialise the environment
@@ -76,6 +77,8 @@ def __init__(
             Amount of (attempted) food deposited by a deposit food action
         signal_deposit_amount
             Amount of signal deposited by the deposit signal action
+        view_distance
+            Number of cells away from an agent observed by each agent
         """
         colony_generator = colony_generator or BasicColonyGenerator(25, 2, (5, 5))
         wrapped_colony_generator = SingleColonyWrapper(colony_generator)
@@ -92,6 +95,7 @@ def __init__(
             take_food_amount=take_food_amount,
             deposit_food_amount=deposit_food_amount,
             signal_deposit_amount=signal_deposit_amount,
+            view_distance=view_distance,
         )
         super().__init__()
 

src/thants/envs/multi.py

@@ -50,6 +50,7 @@ def __init__(
         take_food_amount: float = 0.1,
         deposit_food_amount: float = 0.1,
         signal_deposit_amount: float = 0.1,
+        view_distance: int = 1,
     ) -> None:
         """
         Initialise the environment
@@ -88,13 +89,17 @@ def __init__(
             Amount of (attempted) food deposited by a deposit food action
         signal_deposit_amount
             Amount of signal deposited by the deposit signal action
+        view_distance
+            Number of cells away from an agent observed by each agent
         """
         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
+        assert view_distance > 0, "view_distance should be an integer greater than 0"
+        self.view_distance = view_distance
         self._colonies_generator = colonies_generator or DualBasicColoniesGenerator(
             (25, 25), 2, (5, 5)
         )
@@ -136,7 +141,9 @@ def reset(  # type: ignore[override]
             food=food,
             terrain=terrain,
         )
-        observations = observations_from_state(self.num_agents, state)
+        observations = observations_from_state(
+            self.num_agents, state, view_distance=self.view_distance
+        )
         time_steps = [
             restart(observation=obs, shape=(n,))
             for obs, n in zip(observations, self._colonies_generator.n_agents)
@@ -228,7 +235,9 @@ def step(  # type: ignore[override]
         # Rewards
         rewards = self._reward_fn(self.num_agents, old_state=state, new_state=new_state)
         # Observations
-        observations = observations_from_state(self.num_agents, new_state)
+        observations = observations_from_state(
+            self.num_agents, new_state, view_distance=self.view_distance
+        )
         timestep = [
             jax.lax.cond(
                 state.step >= self.max_steps,
@@ -277,6 +286,7 @@ def observation_spec(  # type: ignore[override]
             get_observation_spec(
                 n,
                 self._colonies_generator.n_signals,
+                self.view_distance,
                 self.carry_capacity,
                 num_colonies=self.num_colonies,
             )

src/thants/envs/presets.py

@@ -24,6 +24,7 @@ def __init__(
         take_food_amount: float = 0.1,
         deposit_food_amount: float = 0.1,
         signal_deposit_amount: float = 0.1,
+        view_distance: int = 1,
     ) -> None:
         """
         Initialise the environment
@@ -52,6 +53,8 @@ def __init__(
             Max food that can be dropped by an ant in a single step
         signal_deposit_amount
             Amount of signal deposited in a single step
+        view_distance
+            Number of cells away from an agent observed by each agent
         """
         food_generator = BasicFoodGenerator(food_drop_dims, food_drop_interval)
         super().__init__(
@@ -65,6 +68,7 @@ def __init__(
             take_food_amount=take_food_amount,
             deposit_food_amount=deposit_food_amount,
             signal_deposit_amount=signal_deposit_amount,
+            view_distance=view_distance,
         )
 
 
@@ -86,6 +90,7 @@ def __init__(
         take_food_amount: float = 0.1,
         deposit_food_amount: float = 0.1,
         signal_deposit_amount: float = 0.1,
+        view_distance: int = 1,
     ) -> None:
         """
         Initialise the environment
@@ -114,6 +119,8 @@ def __init__(
             Max food that can be dropped by an ant in a single step
         signal_deposit_amount
             Amount of signal deposited in a single step
+        view_distance
+            Number of cells away from an agent observed by each agent
         """
         food_generator = BasicFoodGenerator(food_drop_dims, food_drop_interval)
         super().__init__(
@@ -129,4 +136,5 @@ def __init__(
             take_food_amount=take_food_amount,
             deposit_food_amount=deposit_food_amount,
             signal_deposit_amount=signal_deposit_amount,
+            view_distance=view_distance,
         )

src/thants/observations.py

@@ -9,7 +9,7 @@
 
 
 def observations_from_state(
-    colony_sizes: Sequence[int], state: State
+    colony_sizes: Sequence[int], state: State, view_distance: int = 1
 ) -> Sequence[Observations]:
     """
     Generate individual agent observations from state for each colony
@@ -20,6 +20,8 @@ def observations_from_state(
         List of colony sizes
     state
         Environment state
+    view_distance
+        Distance in cells observed around each agent, default 1
 
     Returns
     -------
@@ -37,8 +39,9 @@ def observations_from_state(
     n_signals = state.colonies.signals.shape[1]
     dims = state.food.shape
     dims_arr = jnp.array([dims])
-    idxs = jnp.indices((3, 3))
-    idxs = idxs.swapaxes(0, 2).reshape(9, 2) - 1
+    d = 2 * view_distance + 1
+    idxs = jnp.indices((d, d))
+    idxs = idxs.swapaxes(0, 2).reshape(d**2, 2) - view_distance
 
     def get_ant_view(arr: jax.Array, i: jax.Array, x: jax.Array) -> jax.Array:
         return arr.at[i[:, jnp.newaxis], x[:, 0], x[:, 1]].get()

src/thants/specs.py

@@ -9,6 +9,7 @@
 def get_observation_spec(
     num_agents: int,
     num_signals: int,
+    view_distance: int,
     carry_capacity: float,
     num_colonies: Optional[int] = None,
 ) -> specs.Spec[Observations]:
@@ -21,6 +22,8 @@ def get_observation_spec(
         Number of ants in the colony
     num_signals
         Number of signal channels
+    view_distance
+        View distance in cells each agent can see
     carry_capacity
         Ant carrying capacity
     num_colonies
@@ -31,46 +34,47 @@ def get_observation_spec(
     Spec
         Observation specification
     """
+    n_obs_cells = (2 * view_distance + 1) ** 2
     ants = (
         specs.BoundedArray(
-            shape=(num_agents, num_colonies, 9),
+            shape=(num_agents, num_colonies, n_obs_cells),
             minimum=0.0,
             maximum=1.0,
             dtype=float,
             name="ants",
         )
         if num_colonies
         else specs.BoundedArray(
-            shape=(num_agents, 9),
+            shape=(num_agents, n_obs_cells),
             minimum=0.0,
             maximum=1.0,
             dtype=float,
             name="ants",
         )
     )
     food = specs.BoundedArray(
-        shape=(num_agents, 9),
+        shape=(num_agents, n_obs_cells),
         minimum=0.0,
         maximum=jnp.inf,
         dtype=float,
         name="food",
     )
     signals = specs.BoundedArray(
-        shape=(num_agents, num_signals, 9),
+        shape=(num_agents, num_signals, n_obs_cells),
         minimum=0.0,
         maximum=jnp.inf,
         dtype=float,
         name="signals",
     )
     nest = specs.BoundedArray(
-        shape=(num_agents, 9),
+        shape=(num_agents, n_obs_cells),
         minimum=0.0,
         maximum=1.0,
         dtype=float,
         name="nest",
     )
     terrain = specs.BoundedArray(
-        shape=(num_agents, 9),
+        shape=(num_agents, n_obs_cells),
         minimum=0.0,
         maximum=1.0,
         dtype=float,

src/thants/viewer.py

@@ -204,7 +204,7 @@ def _draw(
 
         ax.imshow(terrain)
         ax.imshow(nests)
-        food_img = ax.imshow(food, alpha=state.food)
+        food_img = ax.imshow(jnp.clip(state.food, 0.0, 1.0), alpha=state.food)
         ants_img = ax.imshow(ants)
 
         return colors, ants_img, food_img

tests/test_observations.py

@@ -1,5 +1,6 @@
 import chex
 import jax.numpy as jnp
+import pytest
 
 from thants.observations import observations_from_state
 from thants.steps import merge_colonies
@@ -59,3 +60,56 @@ def test_colony_observations(key: chex.PRNGKey) -> None:
 
     assert observations[0].signals.shape == (2, 2, 9)
     assert observations[1].signals.shape == (1, 2, 9)
+
+
+@pytest.mark.parametrize(
+    "view_distance, x0, x1",
+    [
+        (1, [8, 7, 5, 4], [0, 2, 6, 8]),
+        (2, [4, 5, 7, 8], [8, 6, 2, 0]),
+        (3, [0, 0, 0, 0], [4, 4, 4, 4]),
+    ],
+)
+def test_observation_range(
+    key: chex.PRNGKey,
+    view_distance: int,
+    x0: list[int],
+    x1: list[int],
+) -> None:
+    dims = (3, 3)
+
+    colonies = [
+        Colony(
+            ants=Ants(
+                pos=jnp.array([[0, 0], [1, 1], [2, 2]]),
+                carrying=jnp.zeros((3,)),
+                health=jnp.ones((3,)),
+            ),
+            nest=jnp.zeros(dims),
+            signals=jnp.zeros((2, *dims)),
+        ),
+    ]
+
+    colonies = merge_colonies(colonies)
+
+    state = State(
+        step=0,
+        key=key,
+        colonies=colonies,
+        food=jnp.arange(9).reshape(3, 3),
+        terrain=jnp.ones(dims, dtype=bool),
+    )
+
+    observations = observations_from_state([3], state, view_distance=view_distance)
+
+    d = 2 * view_distance + 1
+    n = d**2
+
+    assert observations[0].ants.shape == (3, 1, n)
+    assert observations[0].food.shape == (3, n)
+    assert observations[0].nest.shape == (3, n)
+
+    idxs = jnp.array([0, d * (d - 1), d - 1, n - 1])
+
+    assert jnp.array_equal(observations[0].food[0, idxs], jnp.array(x0))
+    assert jnp.array_equal(observations[0].food[1, idxs], jnp.array(x1))