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

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

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

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

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

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

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

{continual_foragax-0.10.1 → continual_foragax-0.10.3}/src/foragax/env.py

@@ -19,9 +19,9 @@ from foragax.weather import get_temperature
 
 
 class Actions(IntEnum):
-    UP = 0
+    DOWN = 0
     RIGHT = 1
-    DOWN = 2
+    UP = 2
     LEFT = 3
 
 
@@ -348,7 +348,9 @@ class ForagaxEnv(environment.Environment):
 
             if is_true_mode:
                 # Apply true object borders by overlaying true colors on border pixels
-                img = apply_true_borders(img, render_grid, self.size)
+                img = apply_true_borders(
+                    img, render_grid, self.size, len(self.object_ids)
+                )
 
             # Add grid lines for world mode
             grid_color = jnp.zeros(3, dtype=jnp.uint8)
@@ -376,7 +378,9 @@ class ForagaxEnv(environment.Environment):
 
             if is_true_mode:
                 # Apply true object borders by overlaying true colors on border pixels
-                img = apply_true_borders(img, aperture, self.aperture_size)
+                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)
@@ -430,7 +434,6 @@ class ForagaxObjectEnv(ForagaxEnv):
         # 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)
 
         # Handle case with no objects (only EMPTY)
         if self.num_color_channels == 0:

{continual_foragax-0.10.1 → continual_foragax-0.10.3}/src/foragax/rendering.py

@@ -9,7 +9,10 @@ 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]
+    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.
 
@@ -17,12 +20,12 @@ def apply_true_borders(
         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
-    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

{continual_foragax-0.10.1 → continual_foragax-0.10.3}/tests/test_foragax.py

@@ -185,7 +185,7 @@ def test_basic_movement():
 
     # stays still when bumping into a wall
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert jnp.array_equal(state.pos, jnp.array([3, 3]))
 
     key, step_key = jax.random.split(key)
@@ -197,7 +197,7 @@ def test_basic_movement():
     assert jnp.array_equal(state.pos, jnp.array([3, 3]))
 
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
     assert jnp.array_equal(state.pos, jnp.array([3, 2]))
 
 
@@ -222,10 +222,10 @@ def test_vision():
 
     # No movement
     key, step_key = jax.random.split(key)
-    obs, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    obs, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
 
     expected = jnp.zeros((3, 3, 1), dtype=int)
-    expected = expected.at[0, 1, 0].set(1)
+    expected = expected.at[2, 1, 0].set(1)
 
     chex.assert_trees_all_equal(state.pos, jnp.array([3, 3]))
     chex.assert_trees_all_equal(obs, expected)
@@ -234,10 +234,10 @@ def test_vision():
     key, step_key = jax.random.split(key)
     obs, state, _, _, _ = env.step(step_key, state, Actions.RIGHT, params)
     key, step_key = jax.random.split(key)
-    obs, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    obs, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     expected = jnp.zeros((3, 3, 1), dtype=int)
-    expected = expected.at[0, 0, 0].set(1)
     expected = expected.at[1, 0, 0].set(1)
+    expected = expected.at[2, 0, 0].set(1)
 
     chex.assert_trees_all_equal(state.pos, jnp.array([4, 4]))
     chex.assert_trees_all_equal(obs, expected)
@@ -264,7 +264,7 @@ def test_respawn():
 
     # Collect the flower
     key, step_key = jax.random.split(key)
-    _, state, reward, _, _ = env.step_env(step_key, state, Actions.UP, params)
+    _, state, reward, _, _ = env.step_env(step_key, state, Actions.DOWN, params)
     assert reward == FLOWER.reward_val
     assert state.object_grid[4, 3] < 0
 
@@ -273,11 +273,11 @@ def test_respawn():
     # Step until it respawns
     for i in range(steps_until_respawn):
         key, step_key = jax.random.split(key)
-        _, state, _, _, _ = env.step_env(step_key, state, Actions.UP, params)
+        _, state, _, _, _ = env.step_env(step_key, state, Actions.DOWN, params)
         assert state.object_grid[4, 3] < 0
 
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert state.object_grid[4, 3] == flower_id
 
 
@@ -291,38 +291,38 @@ def test_wrapping_dynamics():
     # Go up
     assert jnp.array_equal(state.pos, jnp.array([2, 2]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 3]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 4]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 0]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 1]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 2]))
 
     # Go down
     _, state = env.reset(key, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 2]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 1]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 0]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 4]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 3]))
     key, step_key = jax.random.split(key)
-    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
     assert jnp.array_equal(state.pos, jnp.array([2, 2]))
 
     # Go right
@@ -387,10 +387,10 @@ def test_wrapping_vision():
 
     # go down
     key, step_key = jax.random.split(key)
-    obs, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    obs, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
 
     expected = jnp.zeros((3, 3, 1), dtype=int)
-    expected = expected.at[2, 0, 0].set(1)
+    expected = expected.at[0, 0, 0].set(1)
 
     assert jnp.array_equal(state.pos, jnp.array([1, 1]))
     assert jnp.array_equal(obs, expected)
@@ -402,7 +402,7 @@ def test_wrapping_vision():
     obs, state, _, _, _ = env.step(step_key, state, Actions.LEFT, params)
 
     expected = jnp.zeros((3, 3, 1), dtype=int)
-    expected = expected.at[2, 2, 0].set(1)
+    expected = expected.at[0, 2, 0].set(1)
 
     assert jnp.array_equal(state.pos, jnp.array([4, 1]))
     assert jnp.array_equal(obs, expected)
@@ -476,7 +476,7 @@ def test_color_based_partial_observability():
     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)
+    flower_obs = obs[3, 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]))
@@ -616,14 +616,14 @@ def test_benchmark_vision(benchmark):
     @jax.jit
     def _run(state, key):
         key, step_key = jax.random.split(key)
-        obs, new_state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+        obs, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
         return obs, new_state
 
     # warm-up
     obs, new_state = _run(state, key)
 
     expected = jnp.zeros((3, 3, 1), dtype=int)
-    expected = expected.at[0, 1, 0].set(1)
+    expected = expected.at[2, 1, 0].set(1)
 
     chex.assert_trees_all_equal(new_state.pos, jnp.array([3, 3]))
     chex.assert_trees_all_equal(obs, expected)
@@ -674,7 +674,7 @@ def test_benchmark_small_env(benchmark):
         def f(carry, _):
             state, key = carry
             key, step_key = jax.random.split(key, 2)
-            _, new_state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
         (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=1000)
@@ -709,7 +709,7 @@ def test_benchmark_big_env(benchmark):
         def f(carry, _):
             state, key = carry
             key, step_key = jax.random.split(key, 2)
-            _, new_state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
         (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
@@ -750,7 +750,7 @@ def test_benchmark_vmap_env(benchmark):
             key, step_key = jax.random.split(key, 2)
             step_keys = jax.random.split(step_key, num_envs)
             _, new_states, _, _, _ = jax.vmap(env.step, in_axes=(0, 0, None, None))(
-                step_keys, states, Actions.UP, params
+                step_keys, states, Actions.DOWN, params
             )
             return (new_states, key), None
 
@@ -786,7 +786,7 @@ def test_benchmark_small_env_color(benchmark):
         def f(carry, _):
             state, key = carry
             key, step_key = jax.random.split(key, 2)
-            _, new_state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
         (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
@@ -818,7 +818,7 @@ def test_benchmark_small_env_world(benchmark):
         def f(carry, _):
             state, key = carry
             key, step_key = jax.random.split(key, 2)
-            _, new_state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
         (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)