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

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

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

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

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

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

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

{continual_foragax-0.9.0 → continual_foragax-0.10.1}/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.1/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.9.0 → continual_foragax-0.10.1}/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,16 +374,22 @@ 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."""
@@ -478,7 +492,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.9.0 → continual_foragax-0.10.1}/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.1/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.9.0 → continual_foragax-0.10.1}/tests/test_foragax.py

@@ -13,7 +13,7 @@ from foragax.objects import (
     FLOWER,
     LARGE_MOREL,
     MEDIUM_MOREL,
-    MOREL, 
+    MOREL,
     OYSTER,
     THORNS,
     WALL,
@@ -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():