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

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

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

@@ -1,7 +1,9 @@
 foragax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-foragax/env.py,sha256=mT9qfRfmzhFntJ1KKsaP4gKRNqijsJ4hh2sgbiyRBgM,15659
-foragax/objects.py,sha256=lR6QnWX8xBK-EA91BhJGjfA8MNmUg_RXBZO2KAoAWzE,6004
+foragax/colors.py,sha256=rqNPiywP4Nvr0POhsGpasRk-nMMTS3DOwFRUgperlUk,2065
+foragax/env.py,sha256=OtpcyqzBOQLdTvvRegD3SYm4mi4Ga2WE5eJ7OQmQOaw,18294
+foragax/objects.py,sha256=CyBxrykTxpHCI_2hE9jE8mG4TU8R7VxzKdQ5mtxkEqU,6004
 foragax/registry.py,sha256=7_RDXvm_3RNO7culBLGkE0jH8Wk_q6jbMv72dZx4JO8,2722
+foragax/rendering.py,sha256=KAoQpdndy5JDQlyG0c5QDHuH-_Tfy5RuVlDtndnHVjc,2765
 foragax/weather.py,sha256=KNAiwuFz8V__6G75vZIWQKPocLzXqxXn-Vt4TbHIpcA,1258
 foragax/data/ECA_non-blended_custom/TG_SOUID100897.txt,sha256=N7URbX6VlCZvCboUogYjMzy1I-0cfNPOn0QTLSHHfQ0,1776751
 foragax/data/ECA_non-blended_custom/TG_SOUID100928.txt,sha256=AaJMWisVu2YPlZFwvetzHZBI2DqWqUYmp8BVzx6gWZI,817991
@@ -126,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.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.10.0.dist-info/METADATA,sha256=Oo3pJnjoU7VLruPXrhNT-WfzF7cGvkGN8crW-BclMsk,4897
+continual_foragax-0.10.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.10.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.10.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.10.0.dist-info/RECORD,,

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)

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)
 

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.
 

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