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

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

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

{continual_foragax-0.18.0.dist-info → continual_foragax-0.20.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=EyT6KY0d0mXNh6yw10V-8SJVAdyPAGKtRdFV4wXq-JM,19836
+foragax/env.py,sha256=APPu0r31A9SlvmJYkoZzsrh59YFVoLW3Rmyb3JT-dHw,20613
 foragax/objects.py,sha256=iDFo_2CjpgErAm3QdQt5ixmQ8jdIvl7siIPqzGwVcGk,7665
-foragax/registry.py,sha256=65LMx_4tEe7e_N9Lh7AIclqlHkSBCwdesM_UTvMJDlY,7017
+foragax/registry.py,sha256=5Ik61gIzTHc5A1_aOVbRXbVooRXz3GjmL7KdGsZTOTE,8642
 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.18.0.dist-info/METADATA,sha256=DujsDAeSOHLhkyOPt4WXlB9b68bTQU2jsJnjhtV5t2o,4897
-continual_foragax-0.18.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.18.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.18.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.18.0.dist-info/RECORD,,
+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,,

foragax/env.py

@@ -355,17 +355,31 @@ class ForagaxEnv(environment.Environment):
             # Create indices for the aperture
             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])
-
-            # Get original colors from the aperture area
-            original_colors = img[y_coords, x_coords]
-
-            # Calculate tinted colors
-            tinted_colors = (1 - alpha) * original_colors + alpha * agent_color
-
-            # Update the image with tinted colors
-            img = img.at[y_coords, x_coords].set(tinted_colors)
+            y_coords_original = start_y + y_offsets[:, None]
+            x_coords_original = start_x + x_offsets
+
+            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]
+                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]
+                    )
+                )
+            else:
+                y_coords = jnp.mod(y_coords_original, self.size[1])
+                x_coords = jnp.mod(x_coords_original, self.size[0])
+                original_colors = img[y_coords, x_coords]
+                tinted_colors = (1 - alpha) * original_colors + alpha * agent_color
+                img = img.at[y_coords, x_coords].set(tinted_colors)
 
             # Agent color
             img = img.at[state.pos[1], state.pos[0]].set(jnp.array(AGENT.color))

foragax/registry.py

@@ -49,6 +49,13 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         ),
         "nowrap": True,
     },
+    "ForagaxWeather-v3": {
+        "size": None,
+        "aperture_size": None,
+        "objects": None,
+        "biomes": None,
+        "nowrap": True,
+    },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
         "aperture_size": None,
@@ -133,6 +140,13 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         ),
         "nowrap": True,
     },
+    "ForagaxTwoBiome-v7": {
+        "size": None,
+        "aperture_size": None,
+        "objects": (BROWN_MOREL_2, BROWN_OYSTER, GREEN_DEATHCAP_3, GREEN_FAKE_2),
+        "biomes": None,
+        "nowrap": True,
+    },
     "ForagaxTwoBiomeSmall-v1": {
         "size": (16, 8),
         "aperture_size": None,
@@ -197,13 +211,52 @@ def make(
         raise ValueError(f"Unknown env_id: {env_id}")
 
     config = ENV_CONFIGS[env_id].copy()
-
+    if isinstance(aperture_size, int):
+        aperture_size = (aperture_size, aperture_size)
     config["aperture_size"] = aperture_size
     if nowrap is not None:
         config["nowrap"] = nowrap
 
+    if env_id == "ForagaxTwoBiome-v7":
+        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.25, 0.0, 0.5, 0.0),
+            ),
+            # Oyster biome
+            Biome(
+                start=(margin + 7, 0),
+                stop=(margin + 9, 15),
+                object_frequencies=(0.0, 0.25, 0.0, 0.5),
+            ),
+        )
+
+    if env_id == "ForagaxWeather-v3":
+        margin = aperture_size[1] // 2 + 1
+        width = 2 * margin + 9
+        config["size"] = (15, width)
+        config["biomes"] = (
+            # Hot biome
+            Biome(
+                start=(0, margin),
+                stop=(15, margin + 2),
+                object_frequencies=(0.5, 0.0),
+            ),
+            # Cold biome
+            Biome(
+                start=(0, margin + 7),
+                stop=(15, margin + 9),
+                object_frequencies=(0.0, 0.5),
+            ),
+        )
+
     if env_id.startswith("ForagaxWeather"):
-        same_color = env_id == "ForagaxWeather-v2"
+        same_color = env_id in ("ForagaxWeather-v2", "ForagaxWeather-v3")
         hot, cold = create_weather_objects(file_index=file_index, same_color=same_color)
         config["objects"] = (hot, cold)