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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.8.2
+Version: 0.10.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.10.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.10.0}/pyproject.toml

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

{continual_foragax-0.8.2 → continual_foragax-0.10.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.10.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.10.0}/src/continual_foragax.egg-info/SOURCES.txt

@@ -7,9 +7,11 @@ src/continual_foragax.egg-info/entry_points.txt
 src/continual_foragax.egg-info/requires.txt
 src/continual_foragax.egg-info/top_level.txt
 src/foragax/__init__.py
+src/foragax/colors.py
 src/foragax/env.py
 src/foragax/objects.py
 src/foragax/registry.py
+src/foragax/rendering.py
 src/foragax/weather.py
 src/foragax/data/ECA_non-blended_custom/TG_SOUID100897.txt
 src/foragax/data/ECA_non-blended_custom/TG_SOUID100928.txt

continual_foragax-0.10.0/src/foragax/colors.py

@@ -0,0 +1,74 @@
+"""Color utility functions for Foragax."""
+
+import jax
+import jax.numpy as jnp
+
+
+def hsv_to_rgb(h: jax.Array, s: float = 1.0, v: float = 1.0) -> jax.Array:
+    """Convert HSV color values to RGB.
+
+    Args:
+        h: Hue values in range [0, 1]
+        s: Saturation value in range [0, 1], default 1.0
+        v: Value (brightness) in range [0, 1], default 1.0
+
+    Returns:
+        RGB values as array of shape (..., 3) with values in range [0, 1]
+    """
+    c = v * s
+    x = c * (1 - jnp.abs(jnp.mod(h * 6, 2) - 1))
+    m = v - c
+
+    # Create RGB arrays
+    r = jnp.zeros_like(h)
+    g = jnp.zeros_like(h)
+    b = jnp.zeros_like(h)
+
+    # Sector 0: 0-60 degrees (red to yellow)
+    mask0 = h < 1 / 6
+    r = jnp.where(mask0, c, r)
+    g = jnp.where(mask0, x, g)
+
+    # Sector 1: 60-120 degrees (yellow to green)
+    mask1 = (h >= 1 / 6) & (h < 2 / 6)
+    r = jnp.where(mask1, x, r)
+    g = jnp.where(mask1, c, g)
+
+    # Sector 2: 120-180 degrees (green to cyan)
+    mask2 = (h >= 2 / 6) & (h < 3 / 6)
+    g = jnp.where(mask2, c, g)
+    b = jnp.where(mask2, x, b)
+
+    # Sector 3: 180-240 degrees (cyan to blue)
+    mask3 = (h >= 3 / 6) & (h < 4 / 6)
+    g = jnp.where(mask3, x, g)
+    b = jnp.where(mask3, c, b)
+
+    # Sector 4: 240-300 degrees (blue to magenta)
+    mask4 = (h >= 4 / 6) & (h < 5 / 6)
+    r = jnp.where(mask4, x, r)
+    b = jnp.where(mask4, c, b)
+
+    # Sector 5: 300-360 degrees (magenta to red)
+    mask5 = h >= 5 / 6
+    r = jnp.where(mask5, c, r)
+    b = jnp.where(mask5, x, b)
+
+    # Add value offset
+    rgb = jnp.stack([r + m, g + m, b + m], axis=-1)
+    return rgb
+
+
+def hsv_to_rgb_255(h: jax.Array, s: float = 0.9, v: float = 0.8) -> jax.Array:
+    """Convert HSV to RGB with values scaled to 0-255 range for image rendering.
+
+    Args:
+        h: Hue values in range [0, 1]
+        s: Saturation value, default 0.9
+        v: Value (brightness), default 0.8
+
+    Returns:
+        RGB values as uint8 array of shape (..., 3) with values in range [0, 255]
+    """
+    rgb = hsv_to_rgb(h, s, v)
+    return (rgb * 255).astype(jnp.uint8)

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

@@ -14,6 +14,7 @@ from flax import struct
 from gymnax.environments import environment, spaces
 
 from foragax.objects import AGENT, EMPTY, BaseForagaxObject, WeatherObject
+from foragax.rendering import apply_true_borders
 from foragax.weather import get_temperature
 
 
@@ -295,7 +296,11 @@ class ForagaxEnv(environment.Environment):
     @partial(jax.jit, static_argnames=("self", "render_mode"))
     def render(self, state: EnvState, params: EnvParams, render_mode: str = "world"):
         """Render the environment state."""
-        if render_mode == "world":
+        is_world_mode = render_mode in ("world", "world_true")
+        is_aperture_mode = render_mode in ("aperture", "aperture_true")
+        is_true_mode = render_mode in ("world_true", "aperture_true")
+
+        if is_world_mode:
             # Create an RGB image from the object grid
             img = jnp.zeros((self.size[1], self.size[0], 3))
             # Decode grid for rendering: non-negative are objects, negative are empty
@@ -341,15 +346,18 @@ class ForagaxEnv(environment.Environment):
                 jax.image.ResizeMethod.NEAREST,
             )
 
+            if is_true_mode:
+                # Apply true object borders by overlaying true colors on border pixels
+                img = apply_true_borders(img, render_grid, self.size)
+
+            # Add grid lines for world mode
             grid_color = jnp.zeros(3, dtype=jnp.uint8)
             row_indices = jnp.arange(1, self.size[1]) * 24
             col_indices = jnp.arange(1, self.size[0]) * 24
             img = img.at[row_indices, :].set(grid_color)
             img = img.at[:, col_indices].set(grid_color)
 
-            return img
-
-        elif render_mode == "aperture":
+        elif is_aperture_mode:
             obs_grid = jnp.maximum(0, state.object_grid)
             aperture = self._get_aperture(obs_grid, state.pos)
             aperture_one_hot = jax.nn.one_hot(aperture, len(self.object_ids))
@@ -366,36 +374,89 @@ class ForagaxEnv(environment.Environment):
                 jax.image.ResizeMethod.NEAREST,
             )
 
+            if is_true_mode:
+                # Apply true object borders by overlaying true colors on border pixels
+                img = apply_true_borders(img, aperture, self.aperture_size)
+
+            # Add grid lines for aperture mode
             grid_color = jnp.zeros(3, dtype=jnp.uint8)
             row_indices = jnp.arange(1, self.aperture_size[0]) * 24
             col_indices = jnp.arange(1, self.aperture_size[1]) * 24
             img = img.at[row_indices, :].set(grid_color)
             img = img.at[:, col_indices].set(grid_color)
 
-            return img
         else:
             raise ValueError(f"Unknown render_mode: {render_mode}")
 
+        return img
+
 
 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.10.0}/src/foragax/objects.py

@@ -181,7 +181,7 @@ AGENT = DefaultForagaxObject(name="agent", blocking=True, color=(0, 0, 255))
 
 
 def create_weather_objects(
-    file_index: int = 0, repeat: int = 100, multiplier: float = 1.0
+    file_index: int = 0, repeat: int = 500, multiplier: float = 1.0
 ):
     """Create HOT and COLD WeatherObject instances using the specified file.
 

continual_foragax-0.10.0/src/foragax/rendering.py

@@ -0,0 +1,74 @@
+"""Rendering utilities for Foragax environments."""
+
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+
+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]
+) -> jax.Array:
+    """Apply true object borders by overlaying HSV border colors on border pixels.
+
+    Args:
+        base_img: Base image with object colors
+        true_grid: Grid of object IDs for determining border colors
+        grid_size: (height, width) of the grid
+
+    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
+    border_colors = hsv_to_rgb_255(hues[true_grid])
+
+    # Resize border colors to match rendered image size
+    border_img = jax.image.resize(
+        border_colors,
+        (grid_size[0] * 24, grid_size[1] * 24, 3),
+        jax.image.ResizeMethod.NEAREST,
+    )
+
+    # Create border mask (2-pixel thick borders) - vectorized like grid lines
+    height, width = grid_size
+    img_height, img_width = height * 24, width * 24
+
+    border_mask = jnp.zeros((img_height, img_width), dtype=bool)
+
+    # Create border row and column indices for all cells at once
+    cell_rows = jnp.arange(height)
+    cell_cols = jnp.arange(width)
+
+    # Top border rows: 2 rows per cell
+    top_border_rows = cell_rows[:, None] * 24 + jnp.arange(2)[None, :]
+    top_border_rows_flat = top_border_rows.flatten()
+
+    # Bottom border rows: 2 rows per cell
+    bottom_border_rows = cell_rows[:, None] * 24 + 22 + jnp.arange(2)[None, :]
+    bottom_border_rows_flat = bottom_border_rows.flatten()
+
+    # Left border columns: 2 columns per cell
+    left_border_cols = cell_cols[:, None] * 24 + jnp.arange(2)[None, :]
+    left_border_cols_flat = left_border_cols.flatten()
+
+    # Right border columns: 2 columns per cell
+    right_border_cols = cell_cols[:, None] * 24 + 22 + jnp.arange(2)[None, :]
+    right_border_cols_flat = right_border_cols.flatten()
+
+    # Set top and bottom borders (full width rectangles)
+    all_border_rows = jnp.concatenate([top_border_rows_flat, bottom_border_rows_flat])
+    border_mask = border_mask.at[all_border_rows, :].set(True)
+
+    # Set left and right borders (full height rectangles)
+    all_border_cols = jnp.concatenate([left_border_cols_flat, right_border_cols_flat])
+    border_mask = border_mask.at[:, all_border_cols].set(True)
+
+    # Apply border mask: use HSV border colors for border pixels, base colors elsewhere
+    result_img = jnp.where(border_mask[..., None], border_img, base_img)
+    return result_img

{continual_foragax-0.8.2 → continual_foragax-0.10.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