continual-foragax 0.21.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

{continual_foragax-0.21.0.dist-info → continual_foragax-0.23.0.dist-info}/METADATA

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

{continual_foragax-0.21.0.dist-info → continual_foragax-0.23.0.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=phGPMzwNKbwLLIyme1bjca3l2MLNKfYYBmly7JP8cLc,20498
-foragax/objects.py,sha256=8tBFMiquWCkhOpNndNmzovMjw7lE5P81OOlUvN2F65w,8301
-foragax/registry.py,sha256=lSTnLrERIpSgPIEmtUWjjTdUiH6oLV1aO9G58W31o9M,9072
+foragax/env.py,sha256=7csniZeHv2ZX6-8CoYm0pMue7tXmmDwr2XU5msH0n0Q,23680
+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.21.0.dist-info/METADATA,sha256=_i2CbUxmj7yQX6QIv3D6QswoWkRAT_c_EejTUUYjd1w,4897
-continual_foragax-0.21.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.21.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.21.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.21.0.dist-info/RECORD,,
+continual_foragax-0.23.0.dist-info/METADATA,sha256=EIaFemphhlG9KTITgvT_qTQfjKS7sQl7hdrxCceAsW4,4897
+continual_foragax-0.23.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.23.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.23.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.23.0.dist-info/RECORD,,

foragax/env.py

@@ -10,6 +10,7 @@ from typing import Any, Dict, Tuple, Union
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 from flax import struct
 from gymnax.environments import environment, spaces
 
@@ -58,6 +59,7 @@ class EnvParams(environment.EnvParams):
 class EnvState(environment.EnvState):
     pos: jax.Array
     object_grid: jax.Array
+    biome_grid: jax.Array
     time: int
 
 
@@ -66,13 +68,16 @@ class ForagaxEnv(environment.Environment):
 
     def __init__(
         self,
+        name: str = "Foragax-v0",
         size: Union[Tuple[int, int], int] = (10, 10),
         aperture_size: Union[Tuple[int, int], int] = (5, 5),
         objects: Tuple[BaseForagaxObject, ...] = (),
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
+        deterministic_spawn: bool = False,
     ):
         super().__init__()
+        self._name = name
         if isinstance(size, int):
             size = (size, size)
         self.size = size
@@ -81,6 +86,7 @@ class ForagaxEnv(environment.Environment):
             aperture_size = (aperture_size, aperture_size)
         self.aperture_size = aperture_size
         self.nowrap = nowrap
+        self.deterministic_spawn = deterministic_spawn
         objects = (EMPTY,) + objects
         if self.nowrap:
             objects = objects + (PADDING,)
@@ -96,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]
@@ -103,12 +110,35 @@ class ForagaxEnv(environment.Environment):
         self.biome_object_frequencies = jnp.array(
             [b.object_frequencies for b in biomes]
         )
-        self.biome_starts = jnp.array(
+        self.biome_starts = np.array(
             [b.start if b.start is not None else (-1, -1) for b in biomes]
         )
-        self.biome_stops = jnp.array(
+        self.biome_stops = np.array(
             [b.stop if b.stop is not None else (-1, -1) for b in biomes]
         )
+        self.biome_sizes = np.prod(self.biome_stops - self.biome_starts, axis=1)
+        self.biome_masks = []
+        for i in range(self.biome_object_frequencies.shape[0]):
+            # Create mask for the biome
+            start = jax.lax.select(
+                self.biome_starts[i, 0] == -1,
+                jnp.array([0, 0]),
+                self.biome_starts[i],
+            )
+            stop = jax.lax.select(
+                self.biome_stops[i, 0] == -1,
+                jnp.array(self.size),
+                self.biome_stops[i],
+            )
+            rows = jnp.arange(self.size[1])[:, None]
+            cols = jnp.arange(self.size[0])
+            mask = (
+                (rows >= start[1])
+                & (rows < stop[1])
+                & (cols >= start[0])
+                & (cols < stop[0])
+            )
+            self.biome_masks.append(mask)
 
     @property
     def default_params(self) -> EnvParams:
@@ -151,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
@@ -165,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:
@@ -180,6 +249,7 @@ class ForagaxEnv(environment.Environment):
         state = EnvState(
             pos=pos,
             object_grid=object_grid,
+            biome_grid=state.biome_grid,
             time=state.time + 1,
         )
 
@@ -196,57 +266,20 @@ class ForagaxEnv(environment.Environment):
         self, key: jax.Array, params: EnvParams
     ) -> Tuple[jax.Array, EnvState]:
         """Reset environment state."""
-        key, subkey = jax.random.split(key)
-
         object_grid = jnp.zeros((self.size[1], self.size[0]), dtype=int)
-
-        iter_key = subkey
+        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)
-            # Generate random layout
-            grid_rand = jax.random.uniform(biome_key, (self.size[1], self.size[0]))
+            mask = self.biome_masks[i]
+            biome_grid = jnp.where(mask, i, biome_grid)
 
-            # Create mask for the biome
-            start = jax.lax.select(
-                self.biome_starts[i, 0] == -1,
-                jnp.array([0, 0]),
-                self.biome_starts[i],
-            )
-            stop = jax.lax.select(
-                self.biome_stops[i, 0] == -1,
-                jnp.array(self.size),
-                self.biome_stops[i],
-            )
-
-            rows = jnp.arange(self.size[1])[:, None]
-            cols = jnp.arange(self.size[0])
-
-            mask = (
-                (rows >= start[1])
-                & (rows < stop[1])
-                & (cols >= start[0])
-                & (cols < stop[0])
-            )
-
-            # Generate objects for this biome and update the main grid
-            biome_freqs = self.biome_object_frequencies[i]
-            empty_freq = 1.0 - jnp.sum(biome_freqs)
-            all_freqs = jnp.concatenate([jnp.array([empty_freq]), biome_freqs])
-
-            cumulative_freqs = jnp.cumsum(
-                jnp.concatenate([jnp.array([0.0]), all_freqs])
-            )
-
-            # Determine which object to place in each cell
-            # The last object ID will be used for any value of grid_rand >= cumulative_freqs[-1]
-            # so we don't need to cap grid_rand
-            obj_ids_for_biome = jnp.arange(len(all_freqs))
-            cell_obj_ids = (
-                jnp.searchsorted(cumulative_freqs, grid_rand, side="right") - 1
-            )
-            biome_objects = obj_ids_for_biome[cell_obj_ids]
-
-            object_grid = jnp.where(mask, biome_objects, object_grid)
+            if self.deterministic_spawn:
+                biome_objects = self.generate_biome_new(i, biome_key)
+                object_grid = object_grid.at[mask].set(biome_objects)
+            else:
+                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.
         agent_pos = jnp.array([self.size[0] // 2, self.size[1] // 2])
@@ -255,11 +288,31 @@ class ForagaxEnv(environment.Environment):
         state = EnvState(
             pos=agent_pos,
             object_grid=object_grid,
+            biome_grid=biome_grid,
             time=0,
         )
 
         return self.get_obs(state, params), state
 
+    def generate_biome_old(self, i: int, biome_key: jax.Array):
+        biome_freqs = self.biome_object_frequencies[i]
+        grid_rand = jax.random.uniform(biome_key, (self.size[1], self.size[0]))
+        empty_freq = 1.0 - jnp.sum(biome_freqs)
+        all_freqs = jnp.concatenate([jnp.array([empty_freq]), biome_freqs])
+        cumulative_freqs = jnp.cumsum(jnp.concatenate([jnp.array([0.0]), all_freqs]))
+        biome_objects = jnp.searchsorted(cumulative_freqs, grid_rand, side="right") - 1
+        return biome_objects
+
+    def generate_biome_new(self, i: int, biome_key: jax.Array):
+        biome_freqs = self.biome_object_frequencies[i]
+        grid = jnp.linspace(0, 1, self.biome_sizes[i], endpoint=False)
+        biome_objects = len(biome_freqs) - jnp.searchsorted(
+            jnp.cumsum(biome_freqs[::-1]), grid, side="right"
+        )
+        flat_biome_objects = biome_objects.flatten()
+        shuffled_objects = jax.random.permutation(biome_key, flat_biome_objects)
+        return shuffled_objects
+
     def is_terminal(self, state: EnvState, params: EnvParams) -> jax.Array:
         """Foragax is a continuing environment."""
         return False
@@ -267,7 +320,7 @@ class ForagaxEnv(environment.Environment):
     @property
     def name(self) -> str:
         """Environment name."""
-        return "Foragax-v0"
+        return self._name
 
     @property
     def num_actions(self) -> int:
@@ -290,6 +343,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),
             }
         )
@@ -438,13 +497,23 @@ class ForagaxObjectEnv(ForagaxEnv):
 
     def __init__(
         self,
+        name: str = "Foragax-v0",
         size: Union[Tuple[int, int], int] = (10, 10),
         aperture_size: Union[Tuple[int, int], int] = (5, 5),
         objects: Tuple[BaseForagaxObject, ...] = (),
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
+        deterministic_spawn: bool = False,
     ):
-        super().__init__(size, aperture_size, objects, biomes, nowrap)
+        super().__init__(
+            name,
+            size,
+            aperture_size,
+            objects,
+            biomes,
+            nowrap,
+            deterministic_spawn,
+        )
 
         # Compute unique colors and mapping for partial observability
         # Exclude EMPTY (index 0) from color channels

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,
@@ -163,6 +167,58 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "biomes": None,
         "nowrap": True,
     },
+    "ForagaxTwoBiome-v9": {
+        "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-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,
@@ -233,7 +289,12 @@ def make(
     if nowrap is not None:
         config["nowrap"] = nowrap
 
-    if env_id in ("ForagaxTwoBiome-v7", "ForagaxTwoBiome-v8"):
+    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)
@@ -252,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
@@ -286,5 +366,6 @@ def make(
         raise ValueError(f"Unknown observation type: {observation_type}")
 
     env_class = env_class_map[observation_type]
+    config["name"] = env_id
 
     return env_class(**config)