continual-foragax 0.22.0__tar.gz → 0.23.1__tar.gz

{continual_foragax-0.22.0 → continual_foragax-0.23.1}/PKG-INFO

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

{continual_foragax-0.22.0 → continual_foragax-0.23.1}/pyproject.toml

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

{continual_foragax-0.22.0 → continual_foragax-0.23.1}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.22.0 → continual_foragax-0.23.1}/src/foragax/env.py

@@ -59,6 +59,7 @@ class EnvParams(environment.EnvParams):
 class EnvState(environment.EnvState):
     pos: jax.Array
     object_grid: jax.Array
+    biome_grid: jax.Array
     time: int
 
 
@@ -101,6 +102,7 @@ class ForagaxEnv(environment.Environment):
         self.object_blocking = jnp.array([o.blocking for o in objects])
         self.object_collectable = jnp.array([o.collectable for o in objects])
         self.object_colors = jnp.array([o.color for o in objects])
+        self.object_random_respawn = jnp.array([o.random_respawn for o in objects])
 
         self.reward_fns = [o.reward for o in objects]
         self.regen_delay_fns = [o.regen_delay for o in objects]
@@ -179,7 +181,7 @@ class ForagaxEnv(environment.Environment):
         is_collectable = self.object_collectable[obj_at_pos]
 
         # 3. HANDLE OBJECT COLLECTION AND RESPAWNING
-        key, subkey = jax.random.split(key)
+        key, subkey, rand_key = jax.random.split(key, 3)
 
         # Decrement timers (stored as negative values)
         is_timer = state.object_grid < 0
@@ -193,10 +195,49 @@ class ForagaxEnv(environment.Environment):
         )
         encoded_timer = obj_at_pos - ((regen_delay + 1) * num_obj_types)
 
+        def place_at_current_pos(current_grid, timer_val):
+            return current_grid.at[pos[1], pos[0]].set(timer_val)
+
+        def place_at_random_pos(current_grid, timer_val):
+            # Set the collected position to empty temporarily
+            grid = current_grid.at[pos[1], pos[0]].set(0)
+
+            # Find all valid spawn locations (empty cells within the same biome)
+            biome_id = state.biome_grid[pos[1], pos[0]]
+            biome_mask = state.biome_grid == biome_id
+            empty_mask = grid == 0
+            valid_spawn_mask = biome_mask & empty_mask
+
+            num_valid_spawns = jnp.sum(valid_spawn_mask)
+
+            # Get indices of valid spawn locations, padded to a static size
+            y_indices, x_indices = jnp.nonzero(
+                valid_spawn_mask, size=self.size[0] * self.size[1], fill_value=-1
+            )
+            valid_spawn_indices = jnp.stack([y_indices, x_indices], axis=1)
+
+            # Select a random valid location
+            random_idx = jax.random.randint(rand_key, (), 0, num_valid_spawns)
+            new_spawn_pos = valid_spawn_indices[random_idx]
+
+            # Place the timer at the new random position
+            return grid.at[new_spawn_pos[0], new_spawn_pos[1]].set(timer_val)
+
         # If collected, replace object with timer; otherwise, keep it
         val_at_pos = object_grid[pos[1], pos[0]]
-        new_val_at_pos = jax.lax.select(is_collectable, encoded_timer, val_at_pos)
-        object_grid = object_grid.at[pos[1], pos[0]].set(new_val_at_pos)
+        should_collect = is_collectable & (val_at_pos > 0)
+
+        # When not collecting, the value at the position remains unchanged.
+        # When collecting, we either place the timer at the current position or a random one.
+        object_grid = jax.lax.cond(
+            should_collect,
+            lambda: jax.lax.cond(
+                self.object_random_respawn[obj_at_pos],
+                lambda: place_at_random_pos(object_grid, encoded_timer),
+                lambda: place_at_current_pos(object_grid, encoded_timer),
+            ),
+            lambda: object_grid,
+        )
 
         info = {"discount": self.discount(state, params)}
         if self.weather_object is not None:
@@ -208,6 +249,7 @@ class ForagaxEnv(environment.Environment):
         state = EnvState(
             pos=pos,
             object_grid=object_grid,
+            biome_grid=state.biome_grid,
             time=state.time + 1,
         )
 
@@ -225,10 +267,12 @@ class ForagaxEnv(environment.Environment):
     ) -> Tuple[jax.Array, EnvState]:
         """Reset environment state."""
         object_grid = jnp.zeros((self.size[1], self.size[0]), dtype=int)
+        biome_grid = jnp.full((self.size[1], self.size[0]), -1, dtype=int)
         key, iter_key = jax.random.split(key)
         for i in range(self.biome_object_frequencies.shape[0]):
             iter_key, biome_key = jax.random.split(iter_key)
             mask = self.biome_masks[i]
+            biome_grid = jnp.where(mask, i, biome_grid)
 
             if self.deterministic_spawn:
                 biome_objects = self.generate_biome_new(i, biome_key)
@@ -237,13 +281,13 @@ class ForagaxEnv(environment.Environment):
                 biome_objects = self.generate_biome_old(i, biome_key)
                 object_grid = jnp.where(mask, biome_objects, object_grid)
 
-        # Place agent in the center of the world and ensure the cell is empty.
+        # Place agent in the center of the world
         agent_pos = jnp.array([self.size[0] // 2, self.size[1] // 2])
-        object_grid = object_grid.at[agent_pos[1], agent_pos[0]].set(0)
 
         state = EnvState(
             pos=agent_pos,
             object_grid=object_grid,
+            biome_grid=biome_grid,
             time=0,
         )
 
@@ -298,6 +342,12 @@ class ForagaxEnv(environment.Environment):
                     (self.size[1], self.size[0]),
                     int,
                 ),
+                "biome_grid": spaces.Box(
+                    0,
+                    self.biome_object_frequencies.shape[0],
+                    (self.size[1], self.size[0]),
+                    int,
+                ),
                 "time": spaces.Discrete(params.max_steps_in_episode),
             }
         )
@@ -455,7 +505,13 @@ class ForagaxObjectEnv(ForagaxEnv):
         deterministic_spawn: bool = False,
     ):
         super().__init__(
-            name, size, aperture_size, objects, biomes, nowrap, deterministic_spawn
+            name,
+            size,
+            aperture_size,
+            objects,
+            biomes,
+            nowrap,
+            deterministic_spawn,
         )
 
         # Compute unique colors and mapping for partial observability

{continual_foragax-0.22.0 → continual_foragax-0.23.1}/src/foragax/objects.py

@@ -16,11 +16,13 @@ class BaseForagaxObject:
         blocking: bool = False,
         collectable: bool = False,
         color: Tuple[int, int, int] = (0, 0, 0),
+        random_respawn: bool = False,
     ):
         self.name = name
         self.blocking = blocking
         self.collectable = collectable
         self.color = color
+        self.random_respawn = random_respawn
 
     @abc.abstractmethod
     def reward(self, clock: int, rng: jax.Array) -> float:
@@ -44,8 +46,9 @@ class DefaultForagaxObject(BaseForagaxObject):
         collectable: bool = False,
         regen_delay: Tuple[int, int] = (10, 100),
         color: Tuple[int, int, int] = (255, 255, 255),
+        random_respawn: bool = False,
     ):
-        super().__init__(name, blocking, collectable, color)
+        super().__init__(name, blocking, collectable, color, random_respawn)
         self.reward_val = reward
         self.regen_delay_range = regen_delay
 
@@ -70,6 +73,7 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
         mean_regen_delay: int = 10,
         std_regen_delay: int = 1,
         color: Tuple[int, int, int] = (0, 0, 0),
+        random_respawn: bool = False,
     ):
         super().__init__(
             name=name,
@@ -77,6 +81,7 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
             collectable=collectable,
             regen_delay=(mean_regen_delay, mean_regen_delay),
             color=color,
+            random_respawn=random_respawn,
         )
         self.mean_regen_delay = mean_regen_delay
         self.std_regen_delay = std_regen_delay
@@ -99,6 +104,7 @@ class WeatherObject(NormalRegenForagaxObject):
         mean_regen_delay: int = 10,
         std_regen_delay: int = 1,
         color: Tuple[int, int, int] = (0, 0, 0),
+        random_respawn: bool = False,
     ):
         super().__init__(
             name=name,
@@ -106,6 +112,7 @@ class WeatherObject(NormalRegenForagaxObject):
             mean_regen_delay=mean_regen_delay,
             std_regen_delay=std_regen_delay,
             color=color,
+            random_respawn=random_respawn,
         )
         self.rewards = rewards
         self.repeat = repeat
@@ -271,6 +278,40 @@ GREEN_FAKE_UNIFORM = DefaultForagaxObject(
     regen_delay=(9, 11),
 )
 
+# Random respawn variants
+BROWN_MOREL_UNIFORM_RANDOM = DefaultForagaxObject(
+    name="brown_morel",
+    reward=10.0,
+    collectable=True,
+    color=(63, 30, 25),
+    regen_delay=(90, 110),
+    random_respawn=True,
+)
+BROWN_OYSTER_UNIFORM_RANDOM = DefaultForagaxObject(
+    name="brown_oyster",
+    reward=1.0,
+    collectable=True,
+    color=(63, 30, 25),
+    regen_delay=(9, 11),
+    random_respawn=True,
+)
+GREEN_DEATHCAP_UNIFORM_RANDOM = DefaultForagaxObject(
+    name="green_deathcap",
+    reward=-5.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(9, 11),
+    random_respawn=True,
+)
+GREEN_FAKE_UNIFORM_RANDOM = DefaultForagaxObject(
+    name="green_fake",
+    reward=0.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(9, 11),
+    random_respawn=True,
+)
+
 
 def create_weather_objects(
     file_index: int = 0,

{continual_foragax-0.22.0 → continual_foragax-0.23.1}/src/foragax/registry.py

@@ -13,15 +13,19 @@ from foragax.objects import (
     BROWN_MOREL,
     BROWN_MOREL_2,
     BROWN_MOREL_UNIFORM,
+    BROWN_MOREL_UNIFORM_RANDOM,
     BROWN_OYSTER,
     BROWN_OYSTER_UNIFORM,
+    BROWN_OYSTER_UNIFORM_RANDOM,
     GREEN_DEATHCAP,
     GREEN_DEATHCAP_2,
     GREEN_DEATHCAP_3,
     GREEN_DEATHCAP_UNIFORM,
+    GREEN_DEATHCAP_UNIFORM_RANDOM,
     GREEN_FAKE,
     GREEN_FAKE_2,
     GREEN_FAKE_UNIFORM,
+    GREEN_FAKE_UNIFORM_RANDOM,
     LARGE_MOREL,
     LARGE_OYSTER,
     MEDIUM_MOREL,
@@ -176,6 +180,45 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "nowrap": True,
         "deterministic_spawn": True,
     },
+    "ForagaxTwoBiome-v10": {
+        "size": None,
+        "aperture_size": None,
+        "objects": (
+            BROWN_MOREL_UNIFORM_RANDOM,
+            BROWN_OYSTER_UNIFORM_RANDOM,
+            GREEN_DEATHCAP_UNIFORM_RANDOM,
+            GREEN_FAKE_UNIFORM_RANDOM,
+        ),
+        "biomes": None,
+        "nowrap": True,
+        "deterministic_spawn": True,
+    },
+    "ForagaxTwoBiome-v11": {
+        "size": None,
+        "aperture_size": None,
+        "objects": (
+            BROWN_MOREL_UNIFORM,
+            BROWN_OYSTER_UNIFORM,
+            GREEN_DEATHCAP_UNIFORM,
+            GREEN_FAKE_UNIFORM,
+        ),
+        "biomes": None,
+        "nowrap": True,
+        "deterministic_spawn": True,
+    },
+    "ForagaxTwoBiome-v12": {
+        "size": None,
+        "aperture_size": None,
+        "objects": (
+            BROWN_MOREL_UNIFORM_RANDOM,
+            BROWN_OYSTER_UNIFORM_RANDOM,
+            GREEN_DEATHCAP_UNIFORM_RANDOM,
+            GREEN_FAKE_UNIFORM_RANDOM,
+        ),
+        "biomes": None,
+        "nowrap": True,
+        "deterministic_spawn": True,
+    },
     "ForagaxTwoBiomeSmall-v1": {
         "size": (16, 8),
         "aperture_size": None,
@@ -246,7 +289,12 @@ def make(
     if nowrap is not None:
         config["nowrap"] = nowrap
 
-    if env_id in ("ForagaxTwoBiome-v7", "ForagaxTwoBiome-v8", "ForagaxTwoBiome-v9"):
+    if env_id in (
+        "ForagaxTwoBiome-v7",
+        "ForagaxTwoBiome-v8",
+        "ForagaxTwoBiome-v9",
+        "ForagaxTwoBiome-v10",
+    ):
         margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
         config["size"] = (width, 15)
@@ -265,6 +313,25 @@ def make(
             ),
         )
 
+    if env_id in ("ForagaxTwoBiome-v11", "ForagaxTwoBiome-v12"):
+        margin = aperture_size[1] // 2 + 1
+        width = 2 * margin + 9
+        config["size"] = (width, 15)
+        config["biomes"] = (
+            # Morel biome
+            Biome(
+                start=(margin, 0),
+                stop=(margin + 2, 15),
+                object_frequencies=(0.5, 0.0, 0.25, 0.0),
+            ),
+            # Oyster biome
+            Biome(
+                start=(margin + 7, 0),
+                stop=(margin + 9, 15),
+                object_frequencies=(0.0, 0.5, 0.0, 0.25),
+            ),
+        )
+
     if env_id == "ForagaxWeather-v3":
         margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9

{continual_foragax-0.22.0 → continual_foragax-0.23.1}/tests/test_foragax.py

@@ -257,7 +257,7 @@ def test_respawn():
         objects=object_types,
     )
     params = env.default_params
-    _, state = env.reset_env(key, params)
+    _, state = env.reset(key, params)
 
     flower_id = 1  # 0 is EMPTY
 
@@ -268,7 +268,7 @@ def test_respawn():
 
     # Collect the flower
     key, step_key = jax.random.split(key)
-    _, state, reward, _, _ = env.step_env(step_key, state, Actions.DOWN, params)
+    _, state, reward, _, _ = env.step(step_key, state, Actions.DOWN, params)
     assert reward == FLOWER.reward_val
     assert state.object_grid[4, 3] < 0
 
@@ -277,7 +277,7 @@ 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.DOWN, params)
+        _, state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
         assert state.object_grid[4, 3] < 0
 
     key, step_key = jax.random.split(key)
@@ -285,6 +285,102 @@ def test_respawn():
     assert state.object_grid[4, 3] == flower_id
 
 
+def test_random_respawn():
+    """Test that an object respawns at a random empty location within its biome."""
+    key = jax.random.key(0)
+
+    flower_random = DefaultForagaxObject(
+        name="flower",
+        reward=1.0,
+        collectable=True,
+        color=(0, 255, 0),
+        random_respawn=True,
+    )
+
+    object_types = (flower_random, WALL)
+    biome = Biome(start=(2, 2), stop=(5, 5), object_frequencies=(0.0, 0.0))
+    env = ForagaxObjectEnv(
+        size=7,
+        aperture_size=3,
+        objects=object_types,
+        biomes=(biome,),
+    )
+    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_grid=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)
+
+    assert reward == flower_random.reward_val
+    # Original position should be empty
+    assert new_state.object_grid[original_pos[1], original_pos[0]] == 0
+
+    # A timer should be placed somewhere
+    assert jnp.sum(new_state.object_grid < 0) == 1
+    timer_pos_flat = jnp.argmin(new_state.object_grid)
+    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)
+
+    # 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 on an empty cell (not the wall)
+    assert not jnp.array_equal(timer_pos, jnp.array([4, 4]))
+
+
+def test_random_respawn_no_empty_space():
+    """Test that an object respawns at the same spot if no empty space is available."""
+    key = jax.random.key(0)
+
+    flower_random = DefaultForagaxObject(
+        name="flower",
+        reward=1.0,
+        collectable=True,
+        color=(0, 255, 0),
+        random_respawn=True,
+    )
+
+    object_types = (WALL, flower_random)
+    # A 1x1 biome
+    biome = Biome(start=(3, 3), stop=(4, 4), object_frequencies=(0.0, 0.0))
+    env = ForagaxObjectEnv(
+        size=7,
+        objects=object_types,
+        biomes=(biome,),
+    )
+    params = env.default_params
+    _, state = env.reset(key, params)
+
+    flower_id = 2  # WALL is 1
+    original_pos = jnp.array([3, 3])
+
+    # Place a flower in the 1x1 biome
+    grid = jnp.zeros((7, 7), dtype=int)
+    grid = grid.at[original_pos[1], original_pos[0]].set(flower_id)
+    state = state.replace(object_grid=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)
+
+    assert reward == flower_random.reward_val
+    # The timer should be placed back at the original position
+    assert new_state.object_grid[original_pos[1], original_pos[0]] < 0
+
+
 def test_wrapping_dynamics():
     """Test that the agent wraps around the environment boundaries."""
     key = jax.random.key(0)