continual-foragax 0.22.0__py3-none-any.whl → 0.23.1__py3-none-any.whl

{continual_foragax-0.22.0.dist-info → continual_foragax-0.23.1.dist-info}/METADATA

@@ -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.dist-info → continual_foragax-0.23.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
 foragax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 foragax/colors.py,sha256=rqNPiywP4Nvr0POhsGpasRk-nMMTS3DOwFRUgperlUk,2065
-foragax/env.py,sha256=Q-96fMoA_51TIJky2JXApITQHXC-1QfdeB5VZvNwe0o,21362
-foragax/objects.py,sha256=8tBFMiquWCkhOpNndNmzovMjw7lE5P81OOlUvN2F65w,8301
-foragax/registry.py,sha256=pRBWGP18jd4NKl1H-rwDYaAJKUgRWVfENQ9pvTS0tAw,9462
+foragax/env.py,sha256=t0tb4IiDFiC3xJLLv7s_IFOYgoyRdje1PcnEvJNcyVM,23578
+foragax/objects.py,sha256=CblI0NI7PQzeKk3MZa8sbaa9wB4pc_8CyOGbJOFWytE,9391
+foragax/registry.py,sha256=EbvPn2IpEMZ7HRjkAN7_JWaTJfYgjLhHkqspcT88DiY,11339
 foragax/rendering.py,sha256=bms7wvBZTofoR-K-2QD2Ggeed7Viw8uwAEiEpEM3eSo,2768
 foragax/weather.py,sha256=KNAiwuFz8V__6G75vZIWQKPocLzXqxXn-Vt4TbHIpcA,1258
 foragax/data/ECA_non-blended_custom/TG_SOUID100897.txt,sha256=N7URbX6VlCZvCboUogYjMzy1I-0cfNPOn0QTLSHHfQ0,1776751
@@ -128,8 +128,8 @@ foragax/data/ECA_non-blended_custom/TG_SOUID156887.txt,sha256=juzTPgJoJxfqmZkorL
 foragax/data/ECA_non-blended_custom/elements.txt,sha256=OtcUBoDAHxuln79BPKGu0tsQxG_5G2BfAX3Ck130kEA,4507
 foragax/data/ECA_non-blended_custom/metadata.txt,sha256=nudnmOCy5cPJfSXt_IjyX0S5-T7NkCZREICZSimqeqc,48260
 foragax/data/ECA_non-blended_custom/sources.txt,sha256=1j3lSmINAoCMqPqFrHfZJriOz6sTYZNOhXzUwvTLas0,20857
-continual_foragax-0.22.0.dist-info/METADATA,sha256=WCVeg6996zpBsLWNghDpibCAs70CDaI8KStzCFnSPNM,4897
-continual_foragax-0.22.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.22.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.22.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.22.0.dist-info/RECORD,,
+continual_foragax-0.23.1.dist-info/METADATA,sha256=8r-zJyus4axTuoqGPIOKo0-cKRIIT0Phtcd3adW2kVs,4897
+continual_foragax-0.23.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.23.1.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.23.1.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.23.1.dist-info/RECORD,,

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

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,

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