continual-foragax 0.32.0__py3-none-any.whl → 0.33.0__py3-none-any.whl

{continual_foragax-0.32.0.dist-info → continual_foragax-0.33.0.dist-info}/METADATA

@@ -1,17 +1,12 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.32.0
+Version: 0.33.0
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 Requires-Dist: gymnax
 Requires-Dist: six; python_version < "3.10"
-Provides-Extra: dev
-Requires-Dist: pre-commit; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: pytest-benchmark; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
 
 # foragax
 

{continual_foragax-0.32.0.dist-info → continual_foragax-0.33.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 foragax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 foragax/colors.py,sha256=rqNPiywP4Nvr0POhsGpasRk-nMMTS3DOwFRUgperlUk,2065
-foragax/env.py,sha256=OgZegbHnmVCchSEBHfwm4Kgc4rrwTEnx6r4WdAvK_P4,53800
+foragax/env.py,sha256=kf0B0D5CMUnOMiZ8diBlmlt_vf2Wh-M4gSNoeo4jfHY,55147
 foragax/objects.py,sha256=9wv0ZKT89dDkaeVwUwkVo4dwhRVeUxvsTyhoyYKfOEw,26508
 foragax/registry.py,sha256=hfzQHNgX6uoOdbf4_21iH25abQVQZIjBWn7h5bdrSBg,17981
 foragax/rendering.py,sha256=bms7wvBZTofoR-K-2QD2Ggeed7Viw8uwAEiEpEM3eSo,2768
@@ -128,8 +128,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.32.0.dist-info/METADATA,sha256=2FLbgAsQJg-W3DbOvDap6PRb5_ku7g8mfuXFuUZ2Ybs,4897
-continual_foragax-0.32.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.32.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.32.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.32.0.dist-info/RECORD,,
+continual_foragax-0.33.0.dist-info/METADATA,sha256=vEOZLNVNPhccIZDIrN-puYjVPQKWHtThbhoVdUjhF4A,4713
+continual_foragax-0.33.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.33.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.33.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.33.0.dist-info/RECORD,,

foragax/env.py

@@ -506,17 +506,6 @@ class ForagaxEnv(environment.Environment):
             lambda: object_state,
         )
 
-        # Clear color grid when object is collected
-        object_state = jax.lax.cond(
-            should_collect_now,
-            lambda: object_state.replace(
-                color=object_state.color.at[pos[1], pos[0]].set(
-                    jnp.full((3,), 255, dtype=jnp.uint8)
-                )
-            ),
-            lambda: object_state,
-        )
-
         # 3.5. HANDLE OBJECT EXPIRY
         # Only process expiry if there are objects that can expire
         key, object_state = self.expire_objects(key, state, object_state)
@@ -564,6 +553,23 @@ class ForagaxEnv(environment.Environment):
         info["biome_id"] = object_state.biome_id[pos[1], pos[0]]
         info["object_collected_id"] = jax.lax.select(should_collect, obj_at_pos, -1)
 
+        # Compute reward at each grid position
+        fixed_key = jax.random.key(0)  # Fixed key for deterministic reward computation
+
+        def compute_reward(obj_id, params):
+            return jax.lax.cond(
+                obj_id > 0,
+                lambda: jax.lax.switch(
+                    obj_id, self.reward_fns, state.time, fixed_key, params
+                ),
+                lambda: 0.0,
+            )
+
+        reward_grid = jax.vmap(jax.vmap(compute_reward))(
+            object_state.object_id, object_state.state_params
+        )
+        info["rewards"] = reward_grid
+
         # 4. UPDATE STATE
         state = EnvState(
             pos=pos,
@@ -648,12 +654,6 @@ class ForagaxEnv(environment.Environment):
                             rand_key,
                         )
 
-                        # Clear color grid when object expires
-                        empty_color = jnp.full((3,), 255, dtype=jnp.uint8)
-                        new_obj_state = new_obj_state.replace(
-                            color=new_obj_state.color.at[y, x].set(empty_color)
-                        )
-
                         return new_obj_state
 
                     def no_op():
@@ -1114,8 +1114,16 @@ class ForagaxEnv(environment.Environment):
             y_out = (y_coords < 0) | (y_coords >= self.size[1])
             x_out = (x_coords < 0) | (x_coords >= self.size[0])
             out_of_bounds = y_out | x_out
-            padding_index = self.object_ids[-1]
-            aperture = jnp.where(out_of_bounds, padding_index, values)
+
+            # Handle both object_id grids (2D) and color grids (3D)
+            if len(values.shape) == 3:
+                # Color grid: use PADDING color (0, 0, 0)
+                padding_value = jnp.array([0, 0, 0], dtype=values.dtype)
+                aperture = jnp.where(out_of_bounds[..., None], padding_value, values)
+            else:
+                # Object ID grid: use PADDING index
+                padding_index = self.object_ids[-1]
+                aperture = jnp.where(out_of_bounds, padding_index, values)
         else:
             aperture = values
 
@@ -1124,12 +1132,14 @@ class ForagaxEnv(environment.Environment):
     def get_obs(self, state: EnvState, params: EnvParams, key=None) -> jax.Array:
         """Get observation based on observation_type and full_world."""
         obs_grid = state.object_state.object_id
+        color_grid = state.object_state.color
 
         if self.full_world:
             return self._get_world_obs(obs_grid, state)
         else:
             grid = self._get_aperture(obs_grid, state.pos)
-            return self._get_aperture_obs(grid, state)
+            color_grid = self._get_aperture(color_grid, state.pos)
+            return self._get_aperture_obs(grid, color_grid, state)
 
     def _get_world_obs(self, obs_grid: jax.Array, state: EnvState) -> jax.Array:
         """Get world observation."""
@@ -1146,12 +1156,14 @@ class ForagaxEnv(environment.Environment):
             obs = obs.at[state.pos[1], state.pos[0], -1].set(1)
             return obs
         elif self.observation_type == "rgb":
-            obs = jax.nn.one_hot(obs_grid, num_obj_types)
-            # Agent position
-            obs = obs.at[state.pos[1], state.pos[0], :].set(0)
-            obs = obs.at[state.pos[1], state.pos[0], -1].set(1)
-            colors = self.object_colors / 255.0
-            obs = jnp.tensordot(obs, colors, axes=1)
+            # Use state colors directly (supports dynamic biomes)
+            colors = state.object_state.color / 255.0
+
+            # Mask empty cells (object_id == 0) to white
+            empty_mask = obs_grid == 0
+            white_color = jnp.ones((self.size[1], self.size[0], 3), dtype=jnp.float32)
+            obs = jnp.where(empty_mask[..., None], white_color, colors)
+
             return obs
         elif self.observation_type == "color":
             # Handle case with no objects (only EMPTY)
@@ -1168,17 +1180,24 @@ class ForagaxEnv(environment.Environment):
         else:
             raise ValueError(f"Unknown observation_type: {self.observation_type}")
 
-    def _get_aperture_obs(self, aperture: jax.Array, state: EnvState) -> jax.Array:
+    def _get_aperture_obs(
+        self, aperture: jax.Array, color_aperture: jax.Array, state: EnvState
+    ) -> jax.Array:
         """Get aperture observation."""
         if self.observation_type == "object":
             num_obj_types = len(self.object_ids)
             obs = jax.nn.one_hot(aperture, num_obj_types, axis=-1)
             return obs
         elif self.observation_type == "rgb":
-            num_obj_types = len(self.object_ids)
-            aperture_one_hot = jax.nn.one_hot(aperture, num_obj_types)
-            colors = self.object_colors / 255.0
-            obs = jnp.tensordot(aperture_one_hot, colors, axes=1)
+            # Use the color aperture that was passed in
+            aperture_colors = color_aperture / 255.0
+
+            # Mask empty cells (object_id == 0) to white
+            empty_mask = aperture == 0
+            white_color = jnp.ones(aperture_colors.shape, dtype=jnp.float32)
+
+            obs = jnp.where(empty_mask[..., None], white_color, aperture_colors)
+
             return obs
         elif self.observation_type == "color":
             # Handle case with no objects (only EMPTY)
@@ -1229,6 +1248,10 @@ class ForagaxEnv(environment.Environment):
             if self.dynamic_biomes:
                 # Use per-instance colors from state
                 img = state.object_state.color.copy()
+                # Mask empty cells (object_id == 0) to white
+                empty_mask = state.object_state.object_id == 0
+                white_color = jnp.array([255, 255, 255], dtype=jnp.uint8)
+                img = jnp.where(empty_mask[..., None], white_color, img)
             else:
                 # Use default object colors
                 img = jnp.zeros((self.size[1], self.size[0], 3))
@@ -1297,6 +1320,14 @@ class ForagaxEnv(environment.Environment):
                 )
                 img = state.object_state.color[y_coords_adj, x_coords_adj]
 
+                # Mask empty cells (object_id == 0) to white
+                aperture_object_ids = state.object_state.object_id[
+                    y_coords_adj, x_coords_adj
+                ]
+                empty_mask = aperture_object_ids == 0
+                white_color = jnp.array([255, 255, 255], dtype=jnp.uint8)
+                img = jnp.where(empty_mask[..., None], white_color, img)
+
                 if self.nowrap:
                     # For out-of-bounds, use padding object color
                     y_out = (y_coords < 0) | (y_coords >= self.size[1])