continual-foragax 0.19.0__tar.gz → 0.20.1__tar.gz

{continual_foragax-0.19.0 → continual_foragax-0.20.1}/PKG-INFO

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

{continual_foragax-0.19.0 → continual_foragax-0.20.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.19.0"
+version = "0.20.1"
 description = "A continual reinforcement learning benchmark"
 readme = "README.md"
 authors = [
@@ -30,7 +30,7 @@ build-backend = "setuptools.build_meta"
 [tool]
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.19.0"
+version = "0.20.1"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
 

{continual_foragax-0.19.0 → continual_foragax-0.20.1}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.19.0 → continual_foragax-0.20.1}/src/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])

{continual_foragax-0.19.0 → continual_foragax-0.20.1}/src/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,
@@ -216,13 +223,40 @@ def make(
         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)),
+            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)),
+            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)