continual-foragax 0.32.0__tar.gz → 0.33.0__tar.gz

{continual_foragax-0.32.0 → continual_foragax-0.33.0}/PKG-INFO

@@ -1,17 +1,12 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.32.0
+Version: 0.33.0
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 Requires-Dist: gymnax
 Requires-Dist: six; python_version < "3.10"
-Provides-Extra: dev
-Requires-Dist: pre-commit; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: pytest-benchmark; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
 
 # foragax
 

{continual_foragax-0.32.0 → continual_foragax-0.33.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.32.0"
+version = "0.33.0"
 description = "A continual reinforcement learning benchmark"
 readme = "README.md"
 authors = [
@@ -12,7 +12,7 @@ dependencies = [
     "six; python_version < '3.10'",
 ]
 
-[project.optional-dependencies]
+[dependency-groups]
 dev = [
     "pre-commit",
     "pytest",
@@ -30,7 +30,7 @@ build-backend = "setuptools.build_meta"
 [tool]
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.32.0"
+version = "0.33.0"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
 

{continual_foragax-0.32.0 → continual_foragax-0.33.0}/src/continual_foragax.egg-info/PKG-INFO

@@ -1,17 +1,12 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.32.0
+Version: 0.33.0
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 Requires-Dist: gymnax
 Requires-Dist: six; python_version < "3.10"
-Provides-Extra: dev
-Requires-Dist: pre-commit; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: pytest-benchmark; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
 
 # foragax
 

continual_foragax-0.33.0/src/continual_foragax.egg-info/requires.txt

@@ -0,0 +1,4 @@
+gymnax
+
+[:python_version < "3.10"]
+six

{continual_foragax-0.32.0 → continual_foragax-0.33.0}/src/foragax/env.py

@@ -506,17 +506,6 @@ class ForagaxEnv(environment.Environment):
             lambda: object_state,
         )
 
-        # Clear color grid when object is collected
-        object_state = jax.lax.cond(
-            should_collect_now,
-            lambda: object_state.replace(
-                color=object_state.color.at[pos[1], pos[0]].set(
-                    jnp.full((3,), 255, dtype=jnp.uint8)
-                )
-            ),
-            lambda: object_state,
-        )
-
         # 3.5. HANDLE OBJECT EXPIRY
         # Only process expiry if there are objects that can expire
         key, object_state = self.expire_objects(key, state, object_state)
@@ -564,6 +553,23 @@ class ForagaxEnv(environment.Environment):
         info["biome_id"] = object_state.biome_id[pos[1], pos[0]]
         info["object_collected_id"] = jax.lax.select(should_collect, obj_at_pos, -1)
 
+        # 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 > 0,
+                lambda: jax.lax.switch(
+                    obj_id, self.reward_fns, state.time, fixed_key, params
+                ),
+                lambda: 0.0,
+            )
+
+        reward_grid = jax.vmap(jax.vmap(compute_reward))(
+            object_state.object_id, object_state.state_params
+        )
+        info["rewards"] = reward_grid
+
         # 4. UPDATE STATE
         state = EnvState(
             pos=pos,
@@ -648,12 +654,6 @@ class ForagaxEnv(environment.Environment):
                             rand_key,
                         )
 
-                        # Clear color grid when object expires
-                        empty_color = jnp.full((3,), 255, dtype=jnp.uint8)
-                        new_obj_state = new_obj_state.replace(
-                            color=new_obj_state.color.at[y, x].set(empty_color)
-                        )
-
                         return new_obj_state
 
                     def no_op():
@@ -1114,8 +1114,16 @@ class ForagaxEnv(environment.Environment):
             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_index = self.object_ids[-1]
-            aperture = jnp.where(out_of_bounds, padding_index, values)
+
+            # Handle both object_id grids (2D) and color grids (3D)
+            if len(values.shape) == 3:
+                # Color grid: use PADDING color (0, 0, 0)
+                padding_value = jnp.array([0, 0, 0], dtype=values.dtype)
+                aperture = jnp.where(out_of_bounds[..., None], padding_value, values)
+            else:
+                # Object ID grid: use PADDING index
+                padding_index = self.object_ids[-1]
+                aperture = jnp.where(out_of_bounds, padding_index, values)
         else:
             aperture = values
 
@@ -1124,12 +1132,14 @@ class ForagaxEnv(environment.Environment):
     def get_obs(self, state: EnvState, params: EnvParams, key=None) -> jax.Array:
         """Get observation based on observation_type and full_world."""
         obs_grid = state.object_state.object_id
+        color_grid = state.object_state.color
 
         if self.full_world:
             return self._get_world_obs(obs_grid, state)
         else:
             grid = self._get_aperture(obs_grid, state.pos)
-            return self._get_aperture_obs(grid, state)
+            color_grid = self._get_aperture(color_grid, state.pos)
+            return self._get_aperture_obs(grid, color_grid, state)
 
     def _get_world_obs(self, obs_grid: jax.Array, state: EnvState) -> jax.Array:
         """Get world observation."""
@@ -1146,12 +1156,14 @@ class ForagaxEnv(environment.Environment):
             obs = obs.at[state.pos[1], state.pos[0], -1].set(1)
             return obs
         elif self.observation_type == "rgb":
-            obs = jax.nn.one_hot(obs_grid, num_obj_types)
-            # Agent position
-            obs = obs.at[state.pos[1], state.pos[0], :].set(0)
-            obs = obs.at[state.pos[1], state.pos[0], -1].set(1)
-            colors = self.object_colors / 255.0
-            obs = jnp.tensordot(obs, colors, axes=1)
+            # Use state colors directly (supports dynamic biomes)
+            colors = state.object_state.color / 255.0
+
+            # Mask empty cells (object_id == 0) to white
+            empty_mask = obs_grid == 0
+            white_color = jnp.ones((self.size[1], self.size[0], 3), dtype=jnp.float32)
+            obs = jnp.where(empty_mask[..., None], white_color, colors)
+
             return obs
         elif self.observation_type == "color":
             # Handle case with no objects (only EMPTY)
@@ -1168,17 +1180,24 @@ class ForagaxEnv(environment.Environment):
         else:
             raise ValueError(f"Unknown observation_type: {self.observation_type}")
 
-    def _get_aperture_obs(self, aperture: jax.Array, state: EnvState) -> jax.Array:
+    def _get_aperture_obs(
+        self, aperture: jax.Array, color_aperture: jax.Array, state: EnvState
+    ) -> jax.Array:
         """Get aperture observation."""
         if self.observation_type == "object":
             num_obj_types = len(self.object_ids)
             obs = jax.nn.one_hot(aperture, num_obj_types, axis=-1)
             return obs
         elif self.observation_type == "rgb":
-            num_obj_types = len(self.object_ids)
-            aperture_one_hot = jax.nn.one_hot(aperture, num_obj_types)
-            colors = self.object_colors / 255.0
-            obs = jnp.tensordot(aperture_one_hot, colors, axes=1)
+            # Use the color aperture that was passed in
+            aperture_colors = color_aperture / 255.0
+
+            # Mask empty cells (object_id == 0) to white
+            empty_mask = aperture == 0
+            white_color = jnp.ones(aperture_colors.shape, dtype=jnp.float32)
+
+            obs = jnp.where(empty_mask[..., None], white_color, aperture_colors)
+
             return obs
         elif self.observation_type == "color":
             # Handle case with no objects (only EMPTY)
@@ -1229,6 +1248,10 @@ class ForagaxEnv(environment.Environment):
             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))
@@ -1297,6 +1320,14 @@ class ForagaxEnv(environment.Environment):
                 )
                 img = 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])

{continual_foragax-0.32.0 → continual_foragax-0.33.0}/tests/test_foragax.py

@@ -308,7 +308,7 @@ def test_respawn():
 
 
 def test_random_respawn():
-    """Test that an object respawns at a random empty location within its biome."""
+    """Test that an object can respawn at random empty locations within its biome."""
     key = jax.random.key(0)
 
     flower_random = DefaultForagaxObject(
@@ -329,46 +329,60 @@ def test_random_respawn():
         observation_type="color",
     )
     params = env.default_params
-    _, state = env.reset(key, params)
 
     flower_id = 1  # 0 is EMPTY
     original_pos = jnp.array([3, 3])
 
-    # Place a flower and move the agent to it
-    grid = jnp.zeros((7, 7), dtype=int)
-    grid = grid.at[original_pos[1], original_pos[0]].set(flower_id)
-    # Add a wall to make sure it doesn't spawn there
-    grid = grid.at[4, 4].set(2)  # Use a fixed ID for the wall
-    state = state.replace(
-        object_state=state.object_state.replace(object_id=grid), pos=jnp.array([2, 3])
-    )
+    # Test multiple times to verify randomness
+    timer_positions = []
+    for i in range(20):
+        key, reset_key = jax.random.split(key)
+        _, state = env.reset(reset_key, params)
+
+        # Place a flower and move the agent to it
+        grid = jnp.zeros((7, 7), dtype=int)
+        grid = grid.at[original_pos[1], original_pos[0]].set(flower_id)
+        # Add a wall to make sure it doesn't spawn there
+        grid = grid.at[4, 4].set(2)  # Use a fixed ID for the wall
+        state = state.replace(
+            object_state=state.object_state.replace(object_id=grid),
+            pos=jnp.array([2, 3]),
+        )
 
-    # Collect the flower
-    key, step_key = jax.random.split(key)
-    _, new_state, reward, _, _ = env.step(step_key, state, Actions.RIGHT, params)
+        # Collect the flower
+        key, step_key = jax.random.split(key)
+        _, new_state, reward, _, _ = env.step(step_key, state, Actions.RIGHT, params)
 
-    assert reward == flower_random.reward_val
-    # Original position should be empty
-    assert new_state.object_state.object_id[original_pos[1], original_pos[0]] == 0
+        assert reward == flower_random.reward_val
+        # Original position should be empty
+        assert new_state.object_state.object_id[original_pos[1], original_pos[0]] == 0
 
-    # A timer should be placed somewhere (check respawn_timer instead of object_id)
-    assert jnp.sum(new_state.object_state.respawn_timer > 0) == 1
-    timer_pos_flat = jnp.argmax(new_state.object_state.respawn_timer)
-    timer_pos = jnp.array(jnp.unravel_index(timer_pos_flat, (7, 7)))
-    # New position should not be the original position
-    assert not jnp.array_equal(timer_pos, original_pos)
+        # A timer should be placed somewhere (check respawn_timer instead of object_id)
+        assert jnp.sum(new_state.object_state.respawn_timer > 0) == 1
+        timer_pos_flat = jnp.argmax(new_state.object_state.respawn_timer)
+        timer_pos_array = jnp.unravel_index(timer_pos_flat, (7, 7))
+        timer_pos = (int(timer_pos_array[0]), int(timer_pos_array[1]))
 
-    # New position should be within the biome
-    assert jnp.all(timer_pos >= jnp.array(biome.start))
-    assert jnp.all(timer_pos < jnp.array(biome.stop))
+        # New position should be within the biome
+        assert timer_pos[0] >= biome.start[0] and timer_pos[0] < biome.stop[0]
+        assert timer_pos[1] >= biome.start[1] and timer_pos[1] < biome.stop[1]
 
-    # New position should be on an empty cell (not the wall)
-    assert not jnp.array_equal(timer_pos, jnp.array([4, 4]))
+        # New position should be on an empty cell (not the wall)
+        assert timer_pos != (4, 4)
 
-    # Verify respawn_object_id is set correctly
-    assert (
-        new_state.object_state.respawn_object_id[timer_pos[0], timer_pos[1]]
-        == flower_id
+        # Verify respawn_object_id is set correctly
+        assert (
+            new_state.object_state.respawn_object_id[timer_pos[0], timer_pos[1]]
+            == flower_id
+        )
+
+        timer_positions.append(timer_pos)
+
+    # Verify that we get multiple different positions (randomness works)
+    # With 20 trials and multiple valid positions, we should see at least 2 different locations
+    unique_positions = set(timer_positions)
+    assert len(unique_positions) >= 2, (
+        f"Expected at least 2 unique positions, got {len(unique_positions)}: {unique_positions}"
     )
 
 
@@ -1774,7 +1788,7 @@ def test_object_no_individual_respawn():
 
 
 def test_object_color_grid_cleared_on_collection():
-    """Test that object_color_grid is cleared when objects are collected."""
+    """Test that object colors are preserved in state but masked in rendering when collected."""
     key = jax.random.key(0)
     env = ForagaxEnv(
         size=(7, 7),
@@ -1793,6 +1807,9 @@ def test_object_color_grid_cleared_on_collection():
     flower_pos = flower_positions[0]
     y, x = flower_pos
 
+    # Store original color
+    original_color = state.object_state.color[y, x].copy()
+
     # Move agent to flower and collect it
     # Position agent above the flower so moving down collects it
     state = state.replace(pos=jnp.array([x, y - 1]))
@@ -1801,10 +1818,15 @@ def test_object_color_grid_cleared_on_collection():
 
     assert reward == FLOWER.reward_val
 
-    # Check that the color grid at the collected position is cleared (should be [255, 255, 255])
+    # NEW BEHAVIOR: Color is preserved in state (not cleared to white)
     collected_color = state.object_state.color[y, x]
-    expected_empty_color = jnp.array([255, 255, 255], dtype=jnp.uint8)
-    chex.assert_trees_all_equal(collected_color, expected_empty_color)
+    chex.assert_trees_all_equal(collected_color, original_color)
+
+    # But object_id should be 0 (or timer should be set)
+    assert (
+        state.object_state.object_id[y, x] == 0
+        or state.object_state.respawn_timer[y, x] > 0
+    )
 
     # Check that other positions with objects still have their colors
     other_flower_positions = jnp.argwhere(
@@ -1814,11 +1836,31 @@ def test_object_color_grid_cleared_on_collection():
     other_pos = other_flower_positions[0]
     other_color = state.object_state.color[other_pos[0], other_pos[1]]
     # Should not be the empty color
+    expected_empty_color = jnp.array([255, 255, 255], dtype=jnp.uint8)
     assert not jnp.allclose(other_color, expected_empty_color)
 
+    # Test RGB observation: empty cells should appear white
+    # Create an RGB environment with the same setup
+    env_rgb = ForagaxEnv(
+        size=(7, 7),
+        objects=(FLOWER,),
+        biomes=(Biome(object_frequencies=(1.0,)),),
+        observation_type="rgb",
+        dynamic_biomes=True,
+        aperture_size=-1,  # Use world view to get full grid observation
+    )
+    # Use the same state to get observation
+    obs_rgb = env_rgb.get_obs(state, params)
+
+    # Check that the collected position shows white in RGB observation
+    # RGB observations are normalized to [0, 1], so white is [1.0, 1.0, 1.0]
+    collected_rgb = obs_rgb[y, x]
+    expected_white_rgb = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
+    chex.assert_trees_all_close(collected_rgb, expected_white_rgb, rtol=1e-5)
+
 
 def test_object_color_grid_cleared_on_expiry():
-    """Test that object_color_grid is cleared when objects expire."""
+    """Test that object colors are preserved in state but masked in rendering when expired."""
     # Create an object that expires quickly
     expiring_obj = DefaultForagaxObject(
         name="expiring",
@@ -1847,6 +1889,9 @@ def test_object_color_grid_cleared_on_expiry():
     obj_pos = obj_positions[0]
     y, x = obj_pos
 
+    # Store original color
+    original_color = state.object_state.color[y, x].copy()
+
     # Step until expiry (4 steps: expiry happens when age >= expiry_time)
     for _ in range(4):
         key, key_step = jax.random.split(key)
@@ -1854,10 +1899,34 @@ def test_object_color_grid_cleared_on_expiry():
             key_step, state, Actions.DOWN, env.default_params
         )
 
-    # Object should have expired and color should be cleared
+    # NEW BEHAVIOR: Color is preserved in state (not cleared to white)
     expired_color = state.object_state.color[y, x]
-    expected_empty_color = jnp.array([255, 255, 255], dtype=jnp.uint8)
-    chex.assert_trees_all_equal(expired_color, expected_empty_color)
+    chex.assert_trees_all_equal(expired_color, original_color)
+
+    # But object_id should be 0 (or timer should be set)
+    assert (
+        state.object_state.object_id[y, x] == 0
+        or state.object_state.respawn_timer[y, x] > 0
+    )
+
+    # Test RGB observation: empty/expired cells should appear white
+    # Create an RGB environment with the same setup
+    env_rgb = ForagaxEnv(
+        size=(5, 5),
+        aperture_size=-1,  # Use world view to get full grid observation
+        objects=(expiring_obj,),
+        biomes=(Biome(object_frequencies=(0.5,)),),
+        observation_type="rgb",
+        dynamic_biomes=True,
+    )
+    # Use the same state to get observation
+    obs_rgb = env_rgb.get_obs(state, env.default_params)
+
+    # Check that the expired position shows white in RGB observation
+    # RGB observations are normalized to [0, 1], so white is [1.0, 1.0, 1.0]
+    expired_rgb = obs_rgb[y, x]
+    expected_white_rgb = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
+    chex.assert_trees_all_close(expired_rgb, expected_white_rgb, rtol=1e-5)
 
 
 def test_biome_regeneration_preserves_old_objects():
@@ -2728,3 +2797,35 @@ def test_sine_object():
         r_orig = sine_obj.reward(t, key, None)
         r_shift = sine_obj_shifted.reward(t, key, None)
         assert jnp.allclose(r_orig + r_shift, 0.0, atol=0.01)
+
+
+def test_info_rewards():
+    """Test that info contains rewards with reward values for each grid position."""
+    key = jax.random.key(0)
+    env = ForagaxEnv(
+        size=(5, 5),
+        objects=(FLOWER,),
+        biomes=(Biome(object_frequencies=(0.5,)),),
+        observation_type="object",
+    )
+    params = env.default_params
+    obs, state = env.reset(key, params)
+
+    key, step_key = jax.random.split(key)
+    _, _, _, _, info = env.step(step_key, state, Actions.UP, params)
+
+    # Check that rewards is in info
+    assert "rewards" in info
+
+    # Check shape matches environment size
+    assert info["rewards"].shape == (5, 5)
+
+    # Check that positions with objects have non-zero rewards
+    object_mask = state.object_state.object_id > 0
+    rewards_at_objects = info["rewards"][object_mask]
+    assert jnp.all(rewards_at_objects == FLOWER.reward_val)
+
+    # Check that empty positions have zero rewards
+    empty_mask = state.object_state.object_id == 0
+    rewards_at_empty = info["rewards"][empty_mask]
+    assert jnp.all(rewards_at_empty == 0.0)

continual_foragax-0.32.0/src/continual_foragax.egg-info/requires.txt

@@ -1,10 +0,0 @@
-gymnax
-
-[:python_version < "3.10"]
-six
-
-[dev]
-pre-commit
-pytest
-pytest-benchmark
-ruff