continual-foragax 0.10.0__tar.gz → 0.10.2__tar.gz

{continual_foragax-0.10.0 → continual_foragax-0.10.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.10.0
+Version: 0.10.2
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8

{continual_foragax-0.10.0 → continual_foragax-0.10.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.10.0"
+version = "0.10.2"
 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.10.0"
+version = "0.10.2"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
 

{continual_foragax-0.10.0 → continual_foragax-0.10.2}/src/continual_foragax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.10.0
+Version: 0.10.2
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8

{continual_foragax-0.10.0 → continual_foragax-0.10.2}/src/foragax/env.py

@@ -348,7 +348,9 @@ class ForagaxEnv(environment.Environment):
 
             if is_true_mode:
                 # Apply true object borders by overlaying true colors on border pixels
-                img = apply_true_borders(img, render_grid, self.size)
+                img = apply_true_borders(
+                    img, render_grid, self.size, len(self.object_ids)
+                )
 
             # Add grid lines for world mode
             grid_color = jnp.zeros(3, dtype=jnp.uint8)
@@ -376,7 +378,9 @@ class ForagaxEnv(environment.Environment):
 
             if is_true_mode:
                 # Apply true object borders by overlaying true colors on border pixels
-                img = apply_true_borders(img, aperture, self.aperture_size)
+                img = apply_true_borders(
+                    img, aperture, self.aperture_size, len(self.object_ids)
+                )
 
             # Add grid lines for aperture mode
             grid_color = jnp.zeros(3, dtype=jnp.uint8)
@@ -492,7 +496,14 @@ class ForagaxWorldEnv(ForagaxEnv):
         num_obj_types = len(self.object_ids)
         # Decode grid for observation
         obs_grid = jnp.maximum(0, state.object_grid)
-        obs = jax.nn.one_hot(obs_grid, num_obj_types)
+        obs = jnp.zeros((self.size[1], self.size[0], num_obj_types), dtype=jnp.float32)
+
+        num_object_channels = num_obj_types - 1
+        # create masks for all objects at once
+        object_ids = jnp.arange(1, num_obj_types)
+        object_masks = obs_grid[..., None] == object_ids[None, None, :]
+        obs = obs.at[:, :, :num_object_channels].set(object_masks.astype(float))
+
         obs = obs.at[state.pos[1], state.pos[0], -1].set(1)
         return obs
 

{continual_foragax-0.10.0 → continual_foragax-0.10.2}/src/foragax/rendering.py

@@ -9,7 +9,10 @@ from foragax.colors import hsv_to_rgb_255
 
 
 def apply_true_borders(
-    base_img: jax.Array, true_grid: jax.Array, grid_size: Tuple[int, int]
+    base_img: jax.Array,
+    true_grid: jax.Array,
+    grid_size: Tuple[int, int],
+    num_objects: int,
 ) -> jax.Array:
     """Apply true object borders by overlaying HSV border colors on border pixels.
 
@@ -17,12 +20,12 @@ def apply_true_borders(
         base_img: Base image with object colors
         true_grid: Grid of object IDs for determining border colors
         grid_size: (height, width) of the grid
+        num_objects: Number of object types
 
     Returns:
         Image with HSV borders overlaid on border pixels
     """
     # Create HSV border colors for each object type
-    num_objects = true_grid.max() + 1  # Assume object IDs start from 0
     hues = jnp.linspace(0, 1, num_objects, endpoint=False)
 
     # Convert HSV to RGB for border colors

{continual_foragax-0.10.0 → continual_foragax-0.10.2}/tests/test_foragax.py

@@ -143,12 +143,25 @@ def test_rgb_observation_mode_large_aperture():
 
 def test_world_observation_mode():
     # can use world observation mode
-    env = ForagaxWorldEnv(size=(10, 10))
+    env = ForagaxWorldEnv(
+        size=(10, 10),
+        aperture_size=5,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.1, 0.1)),),
+    )
     params = env.default_params
     key = jax.random.key(0)
     obs, state = env.reset(key, params)
 
-    assert obs.shape == (10, 10, 1)
+    assert obs.shape == (10, 10, 3)
+    # Check agent channel (last channel) is 0 except 1 at agent position
+    agent_channel = obs[:, :, 2]
+    assert jnp.sum(agent_channel) == 1
+    assert agent_channel[state.pos[1], state.pos[0]] == 1
+
+    # check non-zero values in other channels
+    assert jnp.sum(obs[:, :, 0]) > 0
+    assert jnp.sum(obs[:, :, 1]) > 0
 
 
 def test_basic_movement():