continual-foragax 0.11.0__tar.gz → 0.12.0__tar.gz

{continual_foragax-0.11.0 → continual_foragax-0.12.0}/PKG-INFO

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

{continual_foragax-0.11.0 → continual_foragax-0.12.0}/pyproject.toml

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

{continual_foragax-0.11.0 → continual_foragax-0.12.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.11.0 → continual_foragax-0.12.0}/src/foragax/env.py

@@ -13,7 +13,13 @@ import jax.numpy as jnp
 from flax import struct
 from gymnax.environments import environment, spaces
 
-from foragax.objects import AGENT, EMPTY, BaseForagaxObject, WeatherObject
+from foragax.objects import (
+    AGENT,
+    EMPTY,
+    PADDING,
+    BaseForagaxObject,
+    WeatherObject,
+)
 from foragax.rendering import apply_true_borders
 from foragax.weather import get_temperature
 
@@ -64,6 +70,7 @@ class ForagaxEnv(environment.Environment):
         aperture_size: Union[Tuple[int, int], int] = (5, 5),
         objects: Tuple[BaseForagaxObject, ...] = (),
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
+        nowrap: bool = False,
     ):
         super().__init__()
         if isinstance(size, int):
@@ -73,7 +80,10 @@ class ForagaxEnv(environment.Environment):
         if isinstance(aperture_size, int):
             aperture_size = (aperture_size, aperture_size)
         self.aperture_size = aperture_size
+        self.nowrap = nowrap
         objects = (EMPTY,) + objects
+        if self.nowrap:
+            objects = objects + (PADDING,)
         self.objects = objects
         self.weather_object = None
         for o in objects:
@@ -122,8 +132,12 @@ class ForagaxEnv(environment.Environment):
         direction = DIRECTIONS[action]
         new_pos = state.pos + direction
 
-        # Wrap around boundaries
-        new_pos = jnp.mod(new_pos, jnp.array(self.size))
+        if self.nowrap:
+            in_bounds = jnp.all((new_pos >= 0) & (new_pos < jnp.array(self.size)))
+            new_pos = jnp.where(in_bounds, new_pos, state.pos)
+        else:
+            # Wrap around boundaries
+            new_pos = jnp.mod(new_pos, jnp.array(self.size))
 
         # Check for blocking objects
         obj_at_new_pos = current_objects[new_pos[1], new_pos[0]]
@@ -288,10 +302,26 @@ class ForagaxEnv(environment.Environment):
 
         y_offsets = jnp.arange(ap_h)
         x_offsets = jnp.arange(ap_w)
-        y_coords = jnp.mod(start_y + y_offsets[:, None], self.size[1])
-        x_coords = jnp.mod(start_x + x_offsets, self.size[0])
+        y_coords = start_y + y_offsets[:, None]
+        x_coords = start_x + x_offsets
+
+        if self.nowrap:
+            # Clamp coordinates to bounds
+            y_coords_clamped = jnp.clip(y_coords, 0, self.size[1] - 1)
+            x_coords_clamped = jnp.clip(x_coords, 0, self.size[0] - 1)
+            values = object_grid[y_coords_clamped, x_coords_clamped]
+            # Mark out-of-bounds positions with -1
+            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)
+        else:
+            y_coords_mod = jnp.mod(y_coords, self.size[1])
+            x_coords_mod = jnp.mod(x_coords, self.size[0])
+            aperture = object_grid[y_coords_mod, x_coords_mod]
 
-        return object_grid[y_coords, x_coords]
+        return aperture
 
     @partial(jax.jit, static_argnames=("self", "render_mode"))
     def render(self, state: EnvState, params: EnvParams, render_mode: str = "world"):
@@ -404,8 +434,9 @@ class ForagaxObjectEnv(ForagaxEnv):
         aperture_size: Union[Tuple[int, int], int] = (5, 5),
         objects: Tuple[BaseForagaxObject, ...] = (),
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
+        nowrap: bool = False,
     ):
-        super().__init__(size, aperture_size, objects, biomes)
+        super().__init__(size, aperture_size, objects, biomes, nowrap)
 
         # Compute unique colors and mapping for partial observability
         # Exclude EMPTY (index 0) from color channels

{continual_foragax-0.11.0 → continual_foragax-0.12.0}/src/foragax/objects.py

@@ -179,6 +179,8 @@ DEATHCAP = DefaultForagaxObject(
 )
 AGENT = DefaultForagaxObject(name="agent", blocking=True, color=(0, 0, 255))
 
+PADDING = DefaultForagaxObject(name="padding", blocking=True, color=(0, 0, 0))
+
 BROWN_MOREL = NormalRegenForagaxObject(
     name="brown_morel",
     reward=30.0,

{continual_foragax-0.11.0 → continual_foragax-0.12.0}/src/foragax/registry.py

@@ -75,6 +75,7 @@ def make(
     observation_type: str = "object",
     aperture_size: Optional[Tuple[int, int]] = (5, 5),
     file_index: int = 0,
+    nowrap: bool = False,
 ) -> ForagaxEnv:
     """Create a Foragax environment.
 
@@ -83,6 +84,8 @@ def make(
         observation_type: The type of observation to use. One of "object", "rgb", or "world".
         aperture_size: The size of the agent's observation aperture. If None, the default
             for the environment is used.
+        file_index: File index for weather objects.
+        nowrap: If True, disables wrapping around environment boundaries.
 
     Returns:
         A Foragax environment instance.
@@ -93,6 +96,7 @@ def make(
     config = ENV_CONFIGS[env_id].copy()
 
     config["aperture_size"] = aperture_size
+    config["nowrap"] = nowrap
 
     if env_id.startswith("ForagaxWeather"):
         hot, cold = create_weather_objects(file_index=file_index)

{continual_foragax-0.11.0 → continual_foragax-0.12.0}/tests/test_foragax.py

@@ -364,6 +364,63 @@ def test_wrapping_dynamics():
     assert jnp.array_equal(state.pos, jnp.array([2, 2]))
 
 
+def test_no_wrapping_dynamics():
+    """Test that the agent does not wrap around the environment boundaries when nowrap=True."""
+    key = jax.random.key(0)
+    env = ForagaxObjectEnv(size=(5, 5), objects=(), nowrap=True)
+    params = env.default_params
+    _, state = env.reset(key, params)
+
+    # Agent starts at center (2,2)
+
+    # Move to top edge
+    for _ in range(2):
+        key, step_key = jax.random.split(key)
+        _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    assert jnp.array_equal(state.pos, jnp.array([2, 0]))  # at top
+
+    # Try to move up again, should stay put
+    key, step_key = jax.random.split(key)
+    _, state, _, _, _ = env.step(step_key, state, Actions.UP, params)
+    assert jnp.array_equal(state.pos, jnp.array([2, 0]))  # still at top
+
+    # Move to bottom edge
+    _, state = env.reset(key, params)
+    for _ in range(2):
+        key, step_key = jax.random.split(key)
+        _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    assert jnp.array_equal(state.pos, jnp.array([2, 4]))  # at bottom
+
+    # Try to move down again, should stay put
+    key, step_key = jax.random.split(key)
+    _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+    assert jnp.array_equal(state.pos, jnp.array([2, 4]))  # still at bottom
+
+    # Move to left edge
+    _, state = env.reset(key, params)
+    for _ in range(2):
+        key, step_key = jax.random.split(key)
+        _, state, _, _, _ = env.step(step_key, state, Actions.LEFT, params)
+    assert jnp.array_equal(state.pos, jnp.array([0, 2]))  # at left
+
+    # Try to move left again, should stay put
+    key, step_key = jax.random.split(key)
+    _, state, _, _, _ = env.step(step_key, state, Actions.LEFT, params)
+    assert jnp.array_equal(state.pos, jnp.array([0, 2]))  # still at left
+
+    # Move to right edge
+    _, state = env.reset(key, params)
+    for _ in range(2):
+        key, step_key = jax.random.split(key)
+        _, state, _, _, _ = env.step(step_key, state, Actions.RIGHT, params)
+    assert jnp.array_equal(state.pos, jnp.array([4, 2]))  # at right
+
+    # Try to move right again, should stay put
+    key, step_key = jax.random.split(key)
+    _, state, _, _, _ = env.step(step_key, state, Actions.RIGHT, params)
+    assert jnp.array_equal(state.pos, jnp.array([4, 2]))  # still at right
+
+
 def test_wrapping_vision():
     """Test that the agent's vision wraps around the environment boundaries."""
     key = jax.random.key(0)
@@ -408,6 +465,40 @@ def test_wrapping_vision():
     assert jnp.array_equal(obs, expected)
 
 
+def test_no_wrapping_vision():
+    """Test that the agent's vision does not wrap around boundaries when nowrap=True."""
+    key = jax.random.key(0)
+    object_types = (FLOWER,)
+    env_no_wrap = ForagaxObjectEnv(
+        size=(7, 7), aperture_size=(3, 3), objects=object_types, nowrap=True
+    )
+    params = env_no_wrap.default_params
+
+    # Place a flower at the opposite corner (6,6)
+    grid = jnp.zeros((7, 7), dtype=int)
+    grid = grid.at[6, 6].set(1)  # FLOWER
+
+    # Agent at (0,0)
+    state = env_no_wrap.reset(key, params)[1]
+    state = state.replace(object_grid=grid, pos=jnp.array([0, 0]))
+
+    # With no wrapping, should not see the flower, see padding
+    obs_no_wrap = env_no_wrap.get_obs(state, params)
+    assert env_no_wrap.num_color_channels == 2  # Flower + padding
+    # Check that padding channel is activated for out of bound positions
+    padding_mask = jnp.array(
+        [
+            [1, 1, 1],
+            [1, 0, 0],
+            [1, 0, 0],
+        ],
+        dtype=bool,
+    )
+    assert jnp.all(obs_no_wrap[padding_mask, 1] == 1)
+    # And flower not visible
+    assert jnp.all(obs_no_wrap[:, :, 0] == 0)
+
+
 def test_generate_objects_in_biome():
     """Test generating objects within a specific biome area."""
     object_types = (WALL, FLOWER, THORNS, MOREL, OYSTER)