continual-foragax 0.42.0__tar.gz → 0.42.2__tar.gz

{continual_foragax-0.42.0 → continual_foragax-0.42.2}/PKG-INFO

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

{continual_foragax-0.42.0 → continual_foragax-0.42.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.42.0"
+version = "0.42.2"
 description = "A continual reinforcement learning benchmark"
 readme = "README.md"
 authors = [
@@ -12,6 +12,10 @@ dependencies = [
     "six; python_version < '3.10'",
 ]
 
+[tool.uv]
+index-strategy = "unsafe-best-match"
+find-links = ["https://storage.googleapis.com/jax-releases/jax_releases.html"]
+
 [dependency-groups]
 dev = [
     "pre-commit",
@@ -30,7 +34,7 @@ build-backend = "setuptools.build_meta"
 [tool]
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.42.0"
+version = "0.42.2"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
 

{continual_foragax-0.42.0 → continual_foragax-0.42.2}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.42.0 → continual_foragax-0.42.2}/src/foragax/env.py

@@ -22,7 +22,13 @@ from foragax.objects import (
     FourierObject,
     WeatherObject,
 )
-from foragax.rendering import apply_true_borders
+from foragax.rendering import (
+    apply_grid_lines,
+    apply_reward_overlay,
+    apply_true_borders,
+    get_base_image,
+    reward_to_color,
+)
 from foragax.weather import get_temperature
 
 
@@ -364,6 +370,19 @@ class ForagaxEnv(environment.Environment):
                     jnp.array(0, dtype=ID_DTYPE)
                 )
 
+                # Extract the actual state to move
+                obj_color = object_state.color[y, x]
+                obj_params = object_state.state_params[y, x]
+                obj_gen = object_state.generation[y, x]
+
+                # Clear visuals at old position
+                new_color = object_state.color.at[y, x].set(
+                    jnp.zeros(3, dtype=COLOR_DTYPE)
+                )
+                new_params = object_state.state_params.at[y, x].set(
+                    jnp.zeros_like(obj_params)
+                )
+
                 # Find valid spawn locations in the same biome
                 biome_id = object_state.biome_id[y, x]
                 biome_mask = object_state.biome_id == biome_id
@@ -381,7 +400,7 @@ class ForagaxEnv(environment.Environment):
                 )
                 new_spawn_pos = valid_spawn_indices[random_idx]
 
-                # Place timer at the new random position
+                # Place timer and move properties at the new random position
                 new_respawn_timer = new_respawn_timer.at[
                     new_spawn_pos[0], new_spawn_pos[1]
                 ].set(timer_val)
@@ -389,10 +408,24 @@ class ForagaxEnv(environment.Environment):
                     new_spawn_pos[0], new_spawn_pos[1]
                 ].set(object_type)
 
+                # Move properties to new position
+                new_color = new_color.at[new_spawn_pos[0], new_spawn_pos[1]].set(
+                    obj_color
+                )
+                new_params = new_params.at[new_spawn_pos[0], new_spawn_pos[1]].set(
+                    obj_params
+                )
+                new_generation = object_state.generation.at[
+                    new_spawn_pos[0], new_spawn_pos[1]
+                ].set(obj_gen)
+
                 return object_state.replace(
                     object_id=new_object_id,
                     respawn_timer=new_respawn_timer,
                     respawn_object_id=new_respawn_object_id,
+                    color=new_color,
+                    state_params=new_params,
+                    generation=new_generation,
                 )
 
             return jax.lax.cond(random_respawn, place_randomly, place_at_position)
@@ -604,22 +637,22 @@ class ForagaxEnv(environment.Environment):
         # 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 > jnp.array(0, dtype=ID_DTYPE),
-                lambda: jax.lax.switch(
-                    obj_id.astype(jnp.int32),
-                    self.reward_fns,
-                    state.time,
-                    fixed_key,
-                    params.astype(jnp.float32),
-                ),
-                lambda: 0.0,
+        def compute_reward(obj_id, params, timer):
+            reward = jax.lax.switch(
+                obj_id.astype(jnp.int32),
+                self.reward_fns,
+                state.time,
+                fixed_key,
+                params.astype(jnp.float32),
             )
+            # Only show reward for objects that are fully present (no timer)
+            mask = (obj_id > 0) & (timer == 0)
+            return jnp.where(mask, reward, 0.0)
 
         reward_grid = jax.vmap(jax.vmap(compute_reward))(
             object_state.object_id.astype(ID_DTYPE),
             object_state.state_params.astype(PARAM_DTYPE),
+            object_state.respawn_timer.astype(TIMER_DTYPE),
         )
         return reward_grid
 
@@ -1465,68 +1498,34 @@ class ForagaxEnv(environment.Environment):
         return spaces.Box(0, 1, obs_shape, float)
 
     def _compute_reward_grid(
-        self, state: EnvState, object_id=None, state_params=None
+        self, state: EnvState, object_id=None, state_params=None, respawn_timer=None
     ) -> jax.Array:
         """Compute rewards for given positions. If no grid provided, uses full world."""
         if object_id is None:
             object_id = state.object_state.object_id
         if state_params is None:
             state_params = state.object_state.state_params
+        if respawn_timer is None:
+            respawn_timer = state.object_state.respawn_timer
 
         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,
+        def compute_reward(obj_id, params, timer):
+            reward = jax.lax.switch(
+                obj_id.astype(jnp.int32),
+                self.reward_fns,
+                state.time,
+                fixed_key,
+                params.astype(jnp.float32),
             )
+            # Only show reward for objects that are fully present (no timer)
+            mask = (obj_id > 0) & (timer == 0)
+            return jnp.where(mask, reward, 0.0)
 
-        reward_grid = jax.vmap(jax.vmap(compute_reward))(object_id, state_params)
-        return reward_grid
-
-    def _reward_to_color(self, reward: jax.Array) -> jax.Array:
-        """Convert reward value to RGB color using diverging gradient.
-
-        Args:
-            reward: Reward value (typically -1 to +1)
-
-        Returns:
-            RGB color array with shape (..., 3) and dtype uint8
-        """
-        # Diverging gradient: +1 = green (0, 255, 0), 0 = white (255, 255, 255), -1 = magenta (255, 0, 255)
-        # Clamp reward to [-1, 1] range for color mapping
-        reward_clamped = jnp.clip(reward, -1.0, 1.0)
-
-        # For positive rewards: interpolate from white to green
-        # For negative rewards: interpolate from white to magenta
-        # At reward = 0: white (255, 255, 255)
-        # At reward = +1: green (0, 255, 0)
-        # At reward = -1: magenta (255, 0, 255)
-
-        red_component = jnp.where(
-            reward_clamped >= 0,
-            (1 - reward_clamped) * 255,  # Fade from white to green: 255 -> 0
-            255,  # Stay at 255 for all negative rewards
-        )
-
-        green_component = jnp.where(
-            reward_clamped >= 0,
-            255,  # Stay at 255 for all positive rewards
-            (1 + reward_clamped) * 255,  # Fade from white to magenta: 255 -> 0
-        )
-
-        blue_component = jnp.where(
-            reward_clamped >= 0,
-            (1 - reward_clamped) * 255,  # Fade from white to green: 255 -> 0
-            255,  # Stay at 255 for all negative rewards
+        reward_grid = jax.vmap(jax.vmap(compute_reward))(
+            object_id, state_params, respawn_timer
         )
-
-        return jnp.stack(
-            [red_component, green_component, blue_component], axis=-1
-        ).astype(jnp.uint8)
+        return reward_grid
 
     @partial(jax.jit, static_argnames=("self", "render_mode"))
     def render(
@@ -1535,13 +1534,7 @@ class ForagaxEnv(environment.Environment):
         params: EnvParams,
         render_mode: str = "world",
     ):
-        """Render the environment state.
-
-        Args:
-            state: Current environment state
-            params: Environment parameters
-            render_mode: One of "world", "world_true", "world_reward", "aperture", "aperture_true", "aperture_reward"
-        """
+        """Render the environment state."""
         is_world_mode = render_mode in ("world", "world_true", "world_reward")
         is_aperture_mode = render_mode in (
             "aperture",
@@ -1552,27 +1545,12 @@ class ForagaxEnv(environment.Environment):
         is_reward_mode = render_mode in ("world_reward", "aperture_reward")
 
         if is_world_mode:
-            # Create an RGB image from the object grid
-            # Use stateful object colors if dynamic_biomes is enabled, else use default colors
-            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))
-                render_grid = state.object_state.object_id
-
-                def update_image(i, img):
-                    color = self.object_colors[i]
-                    mask = render_grid == i
-                    img = jnp.where(mask[..., None], color, img)
-                    return img
-
-                img = jax.lax.fori_loop(0, len(self.object_ids), update_image, img)
+            img = get_base_image(
+                state.object_state.object_id,
+                state.object_state.color,
+                self.object_colors,
+                self.dynamic_biomes,
+            )
 
             # Define constants for all world modes
             alpha = 0.2
@@ -1582,26 +1560,18 @@ class ForagaxEnv(environment.Environment):
 
             if is_reward_mode:
                 # Construct 3x intermediate image
-                # Each cell is 3x3, with reward color in center
                 reward_grid = self._compute_reward_grid(state)
-                reward_colors = self._reward_to_color(reward_grid)
-
-                # Each cell has its base color in 8 pixels and reward color in 1 (center)
-                # Create a 3x3 pattern mask for center pixels
-                cell_mask = jnp.array(
-                    [[False, False, False], [False, True, False], [False, False, False]]
-                )
-                grid_reward_mask = jnp.tile(cell_mask, (self.size[1], self.size[0]))
+                reward_colors = reward_to_color(reward_grid)
 
-                # Repeat base colors and rewards to 3x3
+                # Repeat base colors to 3x scale
                 base_img_x3 = jnp.repeat(jnp.repeat(img, 3, axis=0), 3, axis=1)
-                reward_colors_x3 = jnp.repeat(
-                    jnp.repeat(reward_colors, 3, axis=0), 3, axis=1
-                )
 
-                # Composite base and reward colors
-                img = jnp.where(
-                    grid_reward_mask[..., None], reward_colors_x3, base_img_x3
+                # Composite base and reward colors using helper
+                img = apply_reward_overlay(
+                    base_img_x3,
+                    reward_colors,
+                    reward_grid,
+                    self.size,
                 )
 
                 # Tint the aperture region at 3x scale
@@ -1647,83 +1617,72 @@ class ForagaxEnv(environment.Environment):
                     img, state.object_state.object_id, self.size, len(self.object_ids)
                 )
 
-            # Add grid lines using masking instead of slice-setting
-            row_grid = (jnp.arange(self.size[1] * 24) % 24) == 0
-            col_grid = (jnp.arange(self.size[0] * 24) % 24) == 0
-            # skip first rows/cols as they are borders or managed by caller
-            row_grid = row_grid.at[0].set(False)
-            col_grid = col_grid.at[0].set(False)
-            grid_mask = row_grid[:, None] | col_grid[None, :]
-            img = jnp.where(grid_mask[..., None], self.grid_color_jax, img)
+            # Add grid lines
+            img = apply_grid_lines(img, self.size, self.grid_color_jax)
 
         elif is_aperture_mode:
             obs_grid = state.object_state.object_id
             aperture = self._get_aperture(obs_grid, state.pos)
 
-            if self.dynamic_biomes:
-                # Use per-instance colors from state - extract aperture view
-                y_coords, x_coords, y_coords_adj, x_coords_adj = (
-                    self._compute_aperture_coordinates(state.pos)
-                )
-                img = state.object_state.color[y_coords_adj, x_coords_adj]
+            y_coords, x_coords, y_coords_adj, x_coords_adj = (
+                self._compute_aperture_coordinates(state.pos)
+            )
+            color_state = 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])
-                    x_out = (x_coords < 0) | (x_coords >= self.size[0])
-                    out_of_bounds = y_out | x_out
-                    padding_color = jnp.array(self.objects[-1].color, dtype=jnp.float32)
-                    img = jnp.where(out_of_bounds[..., None], padding_color, img)
-            else:
-                # Use default object colors
-                aperture_one_hot = jax.nn.one_hot(aperture, len(self.object_ids))
-                img = jnp.tensordot(aperture_one_hot, self.object_colors, axes=1)
+            img = get_base_image(
+                aperture,
+                color_state,
+                self.object_colors,
+                self.dynamic_biomes,
+            )
+
+            if self.dynamic_biomes and self.nowrap:
+                # For out-of-bounds, use padding object color
+                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_color = jnp.array(self.objects[-1].color, dtype=jnp.float32)
+                img = jnp.where(out_of_bounds[..., None], padding_color, img)
 
             if is_reward_mode:
                 # Scale image by 3 to create space for reward visualization
                 img = img.astype(jnp.uint8)
                 img = jax.image.resize(
                     img,
-                    (self.aperture_size[0] * 3, self.aperture_size[1] * 3, 3),
+                    (
+                        self.aperture_size[0] * 3,
+                        self.aperture_size[1] * 3,
+                        3,
+                    ),
                     jax.image.ResizeMethod.NEAREST,
                 )
 
                 # Compute rewards for aperture region
-                y_coords, x_coords, y_coords_adj, x_coords_adj = (
-                    self._compute_aperture_coordinates(state.pos)
-                )
-
-                # Get reward grid only for aperture region
-                aperture_object_ids = state.object_state.object_id[
-                    y_coords_adj, x_coords_adj
-                ]
                 aperture_params = state.object_state.state_params[
                     y_coords_adj, x_coords_adj
                 ]
+                aperture_timer = self._get_aperture(
+                    state.object_state.respawn_timer, state.pos
+                )
                 aperture_rewards = self._compute_reward_grid(
-                    state, aperture_object_ids, aperture_params
+                    state, aperture, aperture_params, aperture_timer
                 )
 
                 # Convert rewards to colors
-                reward_colors = self._reward_to_color(aperture_rewards)
+                reward_colors = reward_to_color(aperture_rewards)
 
-                # Place reward colors in the middle cells (index 1 in each 3x3 block)
-                i_indices = jnp.arange(self.aperture_size[0])[:, None] * 3 + 1
-                j_indices = jnp.arange(self.aperture_size[1])[None, :] * 3 + 1
-                img = img.at[i_indices, j_indices].set(reward_colors)
+                # Apply reward overlay using helper
+                img = apply_reward_overlay(
+                    img,
+                    reward_colors,
+                    aperture_rewards,
+                    self.aperture_size,
+                )
 
                 # Draw agent in the center (all 9 cells of the 3x3 block)
                 center_y, center_x = (
-                    self.aperture_size[1] // 2,
                     self.aperture_size[0] // 2,
+                    self.aperture_size[1] // 2,
                 )
                 agent_offsets = jnp.array(
                     [[dy, dx] for dy in range(3) for dx in range(3)]
@@ -1757,17 +1716,13 @@ class ForagaxEnv(environment.Environment):
                 )
 
             if is_true_mode:
-                # Apply true object borders by overlaying true colors on border pixels
+                # Apply true object borders
                 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)
-            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)
+            # Add grid lines
+            img = apply_grid_lines(img, self.aperture_size, self.grid_color_jax)
 
         else:
             raise ValueError(f"Unknown render_mode: {render_mode}")

continual_foragax-0.42.2/src/foragax/rendering.py

@@ -0,0 +1,171 @@
+"""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],
+    num_objects: 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
+        num_objects: Number of object types
+
+    Returns:
+        Image with HSV borders overlaid on border pixels
+    """
+    # Create HSV border colors for each object type
+    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) using vectorized modulo operations
+    img_height, img_width = grid_size[0] * 24, grid_size[1] * 24
+    y_idx = jnp.arange(img_height) % 24
+    x_idx = jnp.arange(img_width) % 24
+
+    # Border pixels are those with offset 0, 1, 22, or 23 within each 24x24 cell
+    is_border_row = (y_idx < 2) | (y_idx >= 22)
+    is_border_col = (x_idx < 2) | (x_idx >= 22)
+    border_mask = is_border_row[:, None] | is_border_col[None, :]
+
+    # 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
+
+
+def reward_to_color(reward: jax.Array) -> jax.Array:
+    """Convert reward value to RGB color using diverging gradient.
+
+    Args:
+        reward: Reward value (typically -1 to +1)
+
+    Returns:
+        RGB color array with shape (..., 3) and dtype uint8
+    """
+    # Diverging gradient: +1 = green (0, 255, 0), 0 = white (255, 255, 255), -1 = magenta (255, 0, 255)
+    # Clamp reward to [-1, 1] range for color mapping
+    reward_clamped = jnp.clip(reward, -1.0, 1.0)
+
+    # For positive rewards: interpolate from white to green
+    # For negative rewards: interpolate from white to magenta
+    # At reward = 0: white (255, 255, 255)
+    # At reward = +1: green (0, 255, 0)
+    # At reward = -1: magenta (255, 0, 255)
+
+    red_component = jnp.where(
+        reward_clamped >= 0,
+        (1 - reward_clamped) * 255,  # Fade from white to green: 255 -> 0
+        255,  # Stay at 255 for all negative rewards
+    )
+
+    green_component = jnp.where(
+        reward_clamped >= 0,
+        255,  # Stay at 255 for all positive rewards
+        (1 + reward_clamped) * 255,  # Fade from white to magenta: 255 -> 0
+    )
+
+    blue_component = jnp.where(
+        reward_clamped >= 0,
+        (1 - reward_clamped) * 255,  # Fade from white to green: 255 -> 0
+        255,  # Stay at 255 for all negative rewards
+    )
+
+    return jnp.stack([red_component, green_component, blue_component], axis=-1).astype(
+        jnp.uint8
+    )
+
+
+def get_base_image(
+    object_id: jax.Array,
+    color_state: jax.Array,
+    object_colors: jax.Array,
+    dynamic_biomes: bool,
+) -> jax.Array:
+    """Construct base RGB image from object IDs or colors."""
+    if dynamic_biomes:
+        # Use per-instance colors from state
+        img = color_state.copy()
+        # Mask empty cells (object_id == 0) to white
+        empty_mask = object_id == 0
+        white_color = jnp.array([255, 255, 255], dtype=jnp.uint8)
+        img = jnp.where(empty_mask[..., None], white_color, img)
+    else:
+        # Map object IDs to colors
+        img = object_colors[object_id]
+
+    return img.astype(jnp.uint8)
+
+
+def apply_grid_lines(
+    img: jax.Array,
+    grid_size: Tuple[int, int],
+    grid_color: jax.Array,
+    cell_size: int = 24,
+) -> jax.Array:
+    """Apply grid lines to the image."""
+    row_grid = (jnp.arange(grid_size[0] * cell_size) % cell_size) == 0
+    col_grid = (jnp.arange(grid_size[1] * cell_size) % cell_size) == 0
+    # skip first rows/cols as they are borders or managed by caller
+    row_grid = row_grid.at[0].set(False)
+    col_grid = col_grid.at[0].set(False)
+    grid_mask = row_grid[:, None] | col_grid[None, :]
+    return jnp.where(grid_mask[..., None], grid_color, img)
+
+
+def apply_reward_overlay(
+    base_img: jax.Array,
+    reward_colors: jax.Array,
+    reward_grid: jax.Array,
+    grid_size: Tuple[int, int],
+) -> jax.Array:
+    """Apply reward visualization overlay (center dots) to the image.
+
+    Only applies dots where the reward is non-zero (abs > 1e-5).
+
+    Args:
+        base_img: Base image at 3x scale (each cell is 3x3)
+        reward_colors: Array of RGB colors for rewards
+        reward_grid: Grid of reward values
+        grid_size: (height, width) of the grid
+
+    Returns:
+        Image with reward dots overlaid
+    """
+    # Create a 3x3 pattern mask for center pixels
+    cell_mask = jnp.array(
+        [[False, False, False], [False, True, False], [False, False, False]]
+    )
+    grid_reward_mask = jnp.tile(cell_mask, grid_size)
+
+    # Only show reward where reward is meaningfully non-zero
+    reward_nonzero = jnp.abs(reward_grid) > 1e-5
+    # Expand to 3x scale
+    reward_nonzero_x3 = jnp.repeat(jnp.repeat(reward_nonzero, 3, axis=0), 3, axis=1)
+
+    # Final mask: center pixel of a cell AND cell has a non-zero reward
+    composite_mask = grid_reward_mask & reward_nonzero_x3
+
+    # Repeat reward colors to 3x to match image scale
+    reward_colors_x3 = jnp.repeat(jnp.repeat(reward_colors, 3, axis=0), 3, axis=1)
+
+    return jnp.where(composite_mask[..., None], reward_colors_x3, base_img)