continual-foragax 0.8.2__tar.gz → 0.9.0__tar.gz

{continual_foragax-0.8.2 → continual_foragax-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.8.2
+Version: 0.9.0
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8
@@ -119,3 +119,12 @@ class into registry configs or construct environments programmatically.
 ## Development
 
 Run unit tests via pytest.
+
+## Acknowledgments
+
+We acknowledge the data providers in the ECA&D project. Klein Tank, A.M.G. and
+Coauthors, 2002. Daily dataset of 20th-century surface air temperature and
+precipitation series for the European Climate Assessment. Int. J. of Climatol.,
+22, 1441-1453.
+
+Data and metadata available at https://www.ecad.eu

{continual_foragax-0.8.2 → continual_foragax-0.9.0}/README.md

@@ -104,3 +104,12 @@ class into registry configs or construct environments programmatically.
 ## Development
 
 Run unit tests via pytest.
+
+## Acknowledgments
+
+We acknowledge the data providers in the ECA&D project. Klein Tank, A.M.G. and
+Coauthors, 2002. Daily dataset of 20th-century surface air temperature and
+precipitation series for the European Climate Assessment. Int. J. of Climatol.,
+22, 1441-1453.
+
+Data and metadata available at https://www.ecad.eu

{continual_foragax-0.8.2 → continual_foragax-0.9.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.8.2"
+version = "0.9.0"
 description = "A continual reinforcement learning benchmark"
 readme = "README.md"
 authors = [
@@ -30,7 +30,7 @@ build-backend = "setuptools.build_meta"
 [tool]
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.8.2"
+version = "0.9.0"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
 

{continual_foragax-0.8.2 → continual_foragax-0.9.0}/src/continual_foragax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.8.2
+Version: 0.9.0
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8
@@ -119,3 +119,12 @@ class into registry configs or construct environments programmatically.
 ## Development
 
 Run unit tests via pytest.
+
+## Acknowledgments
+
+We acknowledge the data providers in the ECA&D project. Klein Tank, A.M.G. and
+Coauthors, 2002. Daily dataset of 20th-century surface air temperature and
+precipitation series for the European Climate Assessment. Int. J. of Climatol.,
+22, 1441-1453.
+
+Data and metadata available at https://www.ecad.eu

{continual_foragax-0.8.2 → continual_foragax-0.9.0}/src/foragax/env.py

@@ -380,22 +380,69 @@ class ForagaxEnv(environment.Environment):
 class ForagaxObjectEnv(ForagaxEnv):
     """Foragax environment with object-based aperture observation."""
 
+    def __init__(
+        self,
+        size: Union[Tuple[int, int], int] = (10, 10),
+        aperture_size: Union[Tuple[int, int], int] = (5, 5),
+        objects: Tuple[BaseForagaxObject, ...] = (),
+        biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
+    ):
+        super().__init__(size, aperture_size, objects, biomes)
+
+        # Compute unique colors and mapping for partial observability
+        # Exclude EMPTY (index 0) from color channels
+        object_colors_no_empty = self.object_colors[1:]
+
+        # Find unique colors in order of first appearance
+        unique_colors = []
+        color_indices = jnp.zeros(len(object_colors_no_empty), dtype=int)
+        color_map = {}
+        next_channel = 0
+
+        for i, color in enumerate(object_colors_no_empty):
+            color_tuple = tuple(color.tolist())
+            if color_tuple not in color_map:
+                color_map[color_tuple] = next_channel
+                unique_colors.append(color)
+                next_channel += 1
+            color_indices = color_indices.at[i].set(color_map[color_tuple])
+
+        self.unique_colors = jnp.array(unique_colors)
+        self.num_color_channels = len(unique_colors)
+        # color_indices maps from object_id-1 to color_channel_index
+        self.object_to_color_map = color_indices
+
     def get_obs(self, state: EnvState, params: EnvParams, key=None) -> jax.Array:
-        num_obj_types = len(self.object_ids)
         # Decode grid for observation
         obs_grid = jnp.maximum(0, state.object_grid)
         aperture = self._get_aperture(obs_grid, state.pos)
         aperture = jnp.flip(aperture, axis=0)
-        obs = jax.nn.one_hot(aperture, num_obj_types)
-        obs = obs[:, :, 1:]
+
+        # Handle case with no objects (only EMPTY)
+        if self.num_color_channels == 0:
+            return jnp.zeros(aperture.shape + (0,), dtype=jnp.float32)
+
+        # Map object IDs to color channel indices
+        # aperture contains object IDs (0 = EMPTY, 1+ = objects)
+        # For EMPTY (0), we want no color channel activated
+        # For objects (1+), map to color channel using object_to_color_map
+        color_channels = jnp.where(
+            aperture == 0,
+            -1,  # Special value for EMPTY
+            jnp.take(self.object_to_color_map, aperture - 1, axis=0),
+        )
+
+        # Create one-hot encoding for color channels
+        # jax.nn.one_hot produces all zeros for -1 (EMPTY positions)
+        obs = jax.nn.one_hot(color_channels, self.num_color_channels, axis=-1)
+
         return obs
 
     def observation_space(self, params: EnvParams) -> spaces.Box:
-        num_obj_types = len(self.object_ids)
         obs_shape = (
             self.aperture_size[0],
             self.aperture_size[1],
-            num_obj_types - 1,
+            self.num_color_channels,
         )
         return spaces.Box(0, 1, obs_shape, float)
 

{continual_foragax-0.8.2 → continual_foragax-0.9.0}/tests/test_foragax.py

@@ -9,7 +9,16 @@ from foragax.env import (
     ForagaxRGBEnv,
     ForagaxWorldEnv,
 )
-from foragax.objects import FLOWER, MOREL, OYSTER, THORNS, WALL
+from foragax.objects import (
+    FLOWER,
+    LARGE_MOREL,
+    MEDIUM_MOREL,
+    MOREL, 
+    OYSTER,
+    THORNS,
+    WALL,
+    DefaultForagaxObject,
+)
 
 
 def test_observation_shape():
@@ -423,6 +432,162 @@ def test_generate_objects_in_biome():
     assert thorns_id not in unique_objects
 
 
+def test_color_based_partial_observability():
+    """Test that objects with the same color are grouped into the same observation channel."""
+    env = ForagaxObjectEnv(
+        size=(10, 10),
+        aperture_size=(5, 5),
+        objects=(MOREL, LARGE_MOREL, MEDIUM_MOREL, FLOWER),
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+
+    # Check that morels (all same color) use 1 channel, plus flower = 2 total
+    assert env.num_color_channels == 2
+
+    # Create a state with different morels and a flower
+    state = env.reset(key, params)[1]
+
+    # Manually place objects
+    grid = jnp.zeros((10, 10), dtype=int)
+    grid = grid.at[5, 5].set(1)  # MOREL
+    grid = grid.at[5, 6].set(2)  # LARGE_MOREL
+    grid = grid.at[5, 7].set(3)  # MEDIUM_MOREL
+    grid = grid.at[6, 5].set(4)  # FLOWER
+    state = state.replace(object_grid=grid)
+
+    obs = env.get_obs(state, params)
+
+    # All morels should activate the same channel (channel 0 for brown)
+    # Flower should activate channel 1 (green)
+    center_obs = obs[2, 2, :]  # MOREL at center
+    morel_obs = obs[2, 3, :]  # LARGE_MOREL
+    med_morel_obs = obs[2, 4, :]  # MEDIUM_MOREL
+    flower_obs = obs[1, 2, :]  # FLOWER (flipped coordinates)
+
+    # All morels should have the same observation (channel 0 activated)
+    chex.assert_trees_all_equal(center_obs, jnp.array([1.0, 0.0]))
+    chex.assert_trees_all_equal(center_obs, morel_obs)
+    chex.assert_trees_all_equal(center_obs, med_morel_obs)
+
+    # Flower should have different observation (channel 1 activated)
+    chex.assert_trees_all_equal(flower_obs, jnp.array([0.0, 1.0]))
+
+
+def test_color_channel_mapping():
+    """Test that the color channel mapping is correct."""
+    env = ForagaxObjectEnv(
+        size=(10, 10),
+        aperture_size=(3, 3),
+        objects=(WALL, FLOWER, THORNS),
+    )
+
+    # WALL (gray), FLOWER (green), THORNS (red) - all different colors
+    assert env.num_color_channels == 3
+
+    # Check that unique colors are correctly identified in order of first appearance
+    expected_colors = jnp.array(
+        [
+            [127, 127, 127],  # gray (WALL) - appears first
+            [0, 255, 0],  # green (FLOWER) - appears second
+            [255, 0, 0],  # red (THORNS) - appears third
+        ]
+    )
+    chex.assert_trees_all_equal(env.unique_colors, expected_colors)
+
+
+def test_same_color_objects_same_channel():
+    """Test that objects with identical colors produce identical observations."""
+    obj1 = DefaultForagaxObject(name="obj1", color=(100, 50, 25))
+    obj2 = DefaultForagaxObject(name="obj2", color=(100, 50, 25))  # Same color
+    obj3 = DefaultForagaxObject(name="obj3", color=(200, 100, 50))  # Different color
+
+    env = ForagaxObjectEnv(
+        size=(7, 7),
+        aperture_size=(3, 3),
+        objects=(obj1, obj2, obj3),
+    )
+    params = env.default_params
+
+    # Should have 2 color channels
+    assert env.num_color_channels == 2
+
+    # Create test state
+    key = jax.random.key(0)
+    state = env.reset(key, params)[1]
+
+    # Place objects
+    grid = jnp.zeros((7, 7), dtype=int)
+    grid = grid.at[3, 3].set(1)  # obj1
+    grid = grid.at[3, 4].set(2)  # obj2 (same color as obj1)
+    grid = grid.at[4, 3].set(3)  # obj3 (different color)
+    state = state.replace(object_grid=grid)
+
+    obs = env.get_obs(state, params)
+
+    # obj1 and obj2 should have identical observations
+    obj1_obs = obs[1, 1, :]  # Center
+    obj2_obs = obs[1, 2, :]  # Right of center
+    obj3_obs = obs[2, 1, :]  # Below center
+
+    chex.assert_trees_all_equal(obj1_obs, obj2_obs)
+    assert not jnp.allclose(obj1_obs, obj3_obs)
+
+
+def test_empty_environment_observation():
+    """Test observation shape when no objects are present."""
+    env = ForagaxObjectEnv(
+        size=(5, 5),
+        aperture_size=(3, 3),
+        objects=(),  # No objects
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+    obs, state = env.reset(key, params)
+
+    # Should have 0 color channels
+    assert env.num_color_channels == 0
+    assert obs.shape == (3, 3, 0)
+
+
+def test_single_color_all_objects():
+    """Test when all objects have the same color."""
+    # Create multiple objects with same color
+    from foragax.objects import DefaultForagaxObject
+
+    obj1 = DefaultForagaxObject(name="obj1", color=(50, 100, 150))
+    obj2 = DefaultForagaxObject(name="obj2", color=(50, 100, 150))
+    obj3 = DefaultForagaxObject(name="obj3", color=(50, 100, 150))
+
+    env = ForagaxObjectEnv(
+        size=(7, 7),
+        aperture_size=(3, 3),
+        objects=(obj1, obj2, obj3),
+    )
+    params = env.default_params
+
+    # Should have 1 color channel
+    assert env.num_color_channels == 1
+
+    # Create test state
+    key = jax.random.key(0)
+    state = env.reset(key, params)[1]
+
+    # Place all objects
+    grid = jnp.zeros((7, 7), dtype=int)
+    grid = grid.at[3, 2].set(1)  # obj1
+    grid = grid.at[3, 3].set(2)  # obj2
+    grid = grid.at[3, 4].set(3)  # obj3
+    state = state.replace(object_grid=grid)
+
+    obs = env.get_obs(state, params)
+
+    # All objects should activate the same channel
+    for i in range(3):
+        obj_obs = obs[1, i, :]  # Row 1, columns 0-2
+        chex.assert_trees_all_equal(obj_obs, jnp.array([1.0]))
+
+
 def test_benchmark_vision(benchmark):
     env = ForagaxObjectEnv(size=7, aperture_size=3, objects=(WALL,))
     params = env.default_params