continual-foragax 0.8.2__py3-none-any.whl → 0.9.0__py3-none-any.whl

{continual_foragax-0.8.2.dist-info → continual_foragax-0.9.0.dist-info}/METADATA

@@ -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.dist-info → continual_foragax-0.9.0.dist-info}/RECORD

@@ -1,5 +1,5 @@
 foragax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-foragax/env.py,sha256=mT9qfRfmzhFntJ1KKsaP4gKRNqijsJ4hh2sgbiyRBgM,15659
+foragax/env.py,sha256=5IJtONEGbW96bRCKxMy9efim-e8VGI1Ab4juZX3e0PY,17613
 foragax/objects.py,sha256=lR6QnWX8xBK-EA91BhJGjfA8MNmUg_RXBZO2KAoAWzE,6004
 foragax/registry.py,sha256=7_RDXvm_3RNO7culBLGkE0jH8Wk_q6jbMv72dZx4JO8,2722
 foragax/weather.py,sha256=KNAiwuFz8V__6G75vZIWQKPocLzXqxXn-Vt4TbHIpcA,1258
@@ -126,8 +126,8 @@ foragax/data/ECA_non-blended_custom/TG_SOUID156887.txt,sha256=juzTPgJoJxfqmZkorL
 foragax/data/ECA_non-blended_custom/elements.txt,sha256=OtcUBoDAHxuln79BPKGu0tsQxG_5G2BfAX3Ck130kEA,4507
 foragax/data/ECA_non-blended_custom/metadata.txt,sha256=nudnmOCy5cPJfSXt_IjyX0S5-T7NkCZREICZSimqeqc,48260
 foragax/data/ECA_non-blended_custom/sources.txt,sha256=1j3lSmINAoCMqPqFrHfZJriOz6sTYZNOhXzUwvTLas0,20857
-continual_foragax-0.8.2.dist-info/METADATA,sha256=RlHNKYqScJI1DV7QEUrkZHKtsCr4vgea5yDzo9PZ1D8,4574
-continual_foragax-0.8.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.8.2.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.8.2.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.8.2.dist-info/RECORD,,
+continual_foragax-0.9.0.dist-info/METADATA,sha256=2-IHrNN5vXoTGo557tCFRfD-49j9xf445rHn3a4My_c,4896
+continual_foragax-0.9.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.9.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.9.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.9.0.dist-info/RECORD,,

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)