continual-foragax 0.10.3__py3-none-any.whl → 0.12.0__py3-none-any.whl

{continual_foragax-0.10.3.dist-info → continual_foragax-0.12.0.dist-info}/METADATA

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

{continual_foragax-0.10.3.dist-info → continual_foragax-0.12.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=EyWJ67GLCJl2Re8s4AoVRGXbccgOGn0J_XxuuWJaGrE,18702
-foragax/objects.py,sha256=CyBxrykTxpHCI_2hE9jE8mG4TU8R7VxzKdQ5mtxkEqU,6004
-foragax/registry.py,sha256=7_RDXvm_3RNO7culBLGkE0jH8Wk_q6jbMv72dZx4JO8,2722
+foragax/env.py,sha256=EyT6KY0d0mXNh6yw10V-8SJVAdyPAGKtRdFV4wXq-JM,19836
+foragax/objects.py,sha256=_TO7tBFCzH5L3JwzHK4bPIh090mtlWBSXcPsZ4y0gHg,6745
+foragax/registry.py,sha256=CjD1eRlY5956royMUPLYuZ3twVcuI9pbRN9M0TEagmo,3437
 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.10.3.dist-info/METADATA,sha256=km5hYBXYDVBJ5a6VhUN1Fjb92QHEQvjG-eThHhlPp84,4897
-continual_foragax-0.10.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.10.3.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.10.3.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.10.3.dist-info/RECORD,,
+continual_foragax-0.12.0.dist-info/METADATA,sha256=yHZYcu0knPEtTUHbAcIGNx9EEqkZpCs_UT3TwAh6bOM,4897
+continual_foragax-0.12.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.12.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.12.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.12.0.dist-info/RECORD,,

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

foragax/objects.py

@@ -179,6 +179,39 @@ 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,
+    collectable=True,
+    color=(63, 30, 25),
+    mean_regen_delay=300,
+    std_regen_delay=30,
+)
+BROWN_OYSTER = NormalRegenForagaxObject(
+    name="brown_oyster",
+    reward=1.0,
+    collectable=True,
+    color=(63, 30, 25),
+    mean_regen_delay=10,
+    std_regen_delay=1,
+)
+GREEN_DEATHCAP = DefaultForagaxObject(
+    name="green_deathcap",
+    reward=-1.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(10, 10),
+)
+GREEN_FAKE = DefaultForagaxObject(
+    name="green_fake",
+    reward=0.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(10, 10),
+)
+
 
 def create_weather_objects(
     file_index: int = 0, repeat: int = 500, multiplier: float = 1.0

foragax/registry.py

@@ -10,6 +10,10 @@ from foragax.env import (
     ForagaxWorldEnv,
 )
 from foragax.objects import (
+    BROWN_MOREL,
+    BROWN_OYSTER,
+    GREEN_DEATHCAP,
+    GREEN_FAKE,
     LARGE_MOREL,
     LARGE_OYSTER,
     MEDIUM_MOREL,
@@ -28,6 +32,19 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
             Biome(start=(0, 10), stop=(15, 12), object_frequencies=(0.0, 0.5)),
         ),
     },
+    "ForagaxTwoBiome-v1": {
+        "size": (15, 15),
+        "aperture_size": None,
+        "objects": (BROWN_MOREL, BROWN_OYSTER, GREEN_DEATHCAP, GREEN_FAKE),
+        "biomes": (
+            # Morel biome
+            Biome(start=(3, 0), stop=(5, 15), object_frequencies=(0.5, 0.0, 0.25, 0.0)),
+            # Oyster biome
+            Biome(
+                start=(10, 0), stop=(12, 15), object_frequencies=(0.0, 0.5, 0.0, 0.25)
+            ),
+        ),
+    },
     "ForagaxTwoBiomeSmall-v1": {
         "size": (16, 8),
         "aperture_size": None,
@@ -58,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.
 
@@ -66,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.
@@ -76,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)