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

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

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

{continual_foragax-0.20.0.dist-info → continual_foragax-0.21.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=APPu0r31A9SlvmJYkoZzsrh59YFVoLW3Rmyb3JT-dHw,20613
-foragax/objects.py,sha256=iDFo_2CjpgErAm3QdQt5ixmQ8jdIvl7siIPqzGwVcGk,7665
-foragax/registry.py,sha256=5Ik61gIzTHc5A1_aOVbRXbVooRXz3GjmL7KdGsZTOTE,8642
+foragax/env.py,sha256=phGPMzwNKbwLLIyme1bjca3l2MLNKfYYBmly7JP8cLc,20498
+foragax/objects.py,sha256=8tBFMiquWCkhOpNndNmzovMjw7lE5P81OOlUvN2F65w,8301
+foragax/registry.py,sha256=lSTnLrERIpSgPIEmtUWjjTdUiH6oLV1aO9G58W31o9M,9072
 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.20.0.dist-info/METADATA,sha256=oTXP1c7cRZSnoUWUV6HvQb0FomwCpWnOtNrDp0M4u8I,4897
-continual_foragax-0.20.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.20.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.20.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.20.0.dist-info/RECORD,,
+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,,

foragax/env.py

@@ -361,19 +361,13 @@ class ForagaxEnv(environment.Environment):
             if self.nowrap:
                 y_coords = jnp.clip(y_coords_original, 0, self.size[1] - 1)
                 x_coords = jnp.clip(x_coords_original, 0, self.size[0] - 1)
-                in_bounds = (
-                    (y_coords_original >= 0)
-                    & (y_coords_original < self.size[1])
-                    & (x_coords_original >= 0)
-                    & (x_coords_original < self.size[0])
-                )
-                original_colors = img[y_coords, x_coords]
+                # Create tint mask: any in-bounds original position maps to a cell makes it tinted
+                tint_mask = jnp.zeros((self.size[1], self.size[0]), dtype=int)
+                tint_mask = tint_mask.at[y_coords, x_coords].set(1)
+                # Apply tint to masked positions
+                original_colors = img
                 tinted_colors = (1 - alpha) * original_colors + alpha * agent_color
-                img = img.at[y_coords, x_coords].set(
-                    jnp.where(
-                        in_bounds[..., None], tinted_colors, img[y_coords, x_coords]
-                    )
-                )
+                img = jnp.where(tint_mask[..., None], tinted_colors, img)
             else:
                 y_coords = jnp.mod(y_coords_original, self.size[1])
                 x_coords = jnp.mod(x_coords_original, self.size[0])

foragax/objects.py

@@ -242,6 +242,34 @@ GREEN_FAKE_2 = NormalRegenForagaxObject(
     mean_regen_delay=10,
     std_regen_delay=1,
 )
+BROWN_MOREL_UNIFORM = DefaultForagaxObject(
+    name="brown_morel",
+    reward=10.0,
+    collectable=True,
+    color=(63, 30, 25),
+    regen_delay=(90, 110),
+)
+BROWN_OYSTER_UNIFORM = DefaultForagaxObject(
+    name="brown_oyster",
+    reward=1.0,
+    collectable=True,
+    color=(63, 30, 25),
+    regen_delay=(9, 11),
+)
+GREEN_DEATHCAP_UNIFORM = DefaultForagaxObject(
+    name="green_deathcap",
+    reward=-5.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(9, 11),
+)
+GREEN_FAKE_UNIFORM = DefaultForagaxObject(
+    name="green_fake",
+    reward=0.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(9, 11),
+)
 
 
 def create_weather_objects(

foragax/registry.py

@@ -12,12 +12,16 @@ from foragax.env import (
 from foragax.objects import (
     BROWN_MOREL,
     BROWN_MOREL_2,
+    BROWN_MOREL_UNIFORM,
     BROWN_OYSTER,
+    BROWN_OYSTER_UNIFORM,
     GREEN_DEATHCAP,
     GREEN_DEATHCAP_2,
     GREEN_DEATHCAP_3,
+    GREEN_DEATHCAP_UNIFORM,
     GREEN_FAKE,
     GREEN_FAKE_2,
+    GREEN_FAKE_UNIFORM,
     LARGE_MOREL,
     LARGE_OYSTER,
     MEDIUM_MOREL,
@@ -147,6 +151,18 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "biomes": None,
         "nowrap": True,
     },
+    "ForagaxTwoBiome-v8": {
+        "size": None,
+        "aperture_size": None,
+        "objects": (
+            BROWN_MOREL_UNIFORM,
+            BROWN_OYSTER_UNIFORM,
+            GREEN_DEATHCAP_UNIFORM,
+            GREEN_FAKE_UNIFORM,
+        ),
+        "biomes": None,
+        "nowrap": True,
+    },
     "ForagaxTwoBiomeSmall-v1": {
         "size": (16, 8),
         "aperture_size": None,
@@ -217,7 +233,7 @@ def make(
     if nowrap is not None:
         config["nowrap"] = nowrap
 
-    if env_id == "ForagaxTwoBiome-v7":
+    if env_id in ("ForagaxTwoBiome-v7", "ForagaxTwoBiome-v8"):
         margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
         config["size"] = (width, 15)