continual-foragax 0.24.2__py3-none-any.whl → 0.26.0__py3-none-any.whl

{continual_foragax-0.24.2.dist-info → continual_foragax-0.26.0.dist-info}/METADATA

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

{continual_foragax-0.24.2.dist-info → continual_foragax-0.26.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=t0tb4IiDFiC3xJLLv7s_IFOYgoyRdje1PcnEvJNcyVM,23578
-foragax/objects.py,sha256=CblI0NI7PQzeKk3MZa8sbaa9wB4pc_8CyOGbJOFWytE,9391
-foragax/registry.py,sha256=iK-Pc9m7NGTds6UidE92LGNMGoMDwJHP09RcO5AwOZ0,12971
+foragax/env.py,sha256=w9KDLRm0xvJjnRP0C4gynk1At6LVdrz5xBEtDC3ePvM,24958
+foragax/objects.py,sha256=FCLZ-8d7qq9VMTG6G-TaRt842-sjgB0-DH0IoHwwngI,9503
+foragax/registry.py,sha256=7-6VDN1MKVEvX_1u5G8NkSpv9BccEmtjJa77-OTNg3A,14324
 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.24.2.dist-info/METADATA,sha256=UcZgs3ahcEv6m-4DIRv8E9giWBeR_g9adHROvoASUpI,4897
-continual_foragax-0.24.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.24.2.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.24.2.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.24.2.dist-info/RECORD,,
+continual_foragax-0.26.0.dist-info/METADATA,sha256=9cCIfMlbC9CAiRZBOoeiZwO45j-v60ZHjn3kjRTu2IE,4897
+continual_foragax-0.26.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.26.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.26.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.26.0.dist-info/RECORD,,

foragax/env.py

@@ -1,12 +1,12 @@
-"""JAX implementation of Foragax environment.
+"""JAX implementation of Forager environment.
 
-Source: https://github.com/andnp/Foragax
+Source: https://github.com/andnp/Forager
 """
 
 from dataclasses import dataclass
 from enum import IntEnum
 from functools import partial
-from typing import Any, Dict, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import jax
 import jax.numpy as jnp
@@ -75,6 +75,7 @@ class ForagaxEnv(environment.Environment):
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
         deterministic_spawn: bool = False,
+        teleport_interval: Optional[int] = None,
     ):
         super().__init__()
         self._name = name
@@ -87,6 +88,7 @@ class ForagaxEnv(environment.Environment):
         self.aperture_size = aperture_size
         self.nowrap = nowrap
         self.deterministic_spawn = deterministic_spawn
+        self.teleport_interval = teleport_interval
         objects = (EMPTY,) + objects
         if self.nowrap:
             objects = objects + (PADDING,)
@@ -117,6 +119,16 @@ class ForagaxEnv(environment.Environment):
             [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_starts_jax = jnp.array(self.biome_starts)
+        self.biome_stops_jax = jnp.array(self.biome_stops)
+        biome_centers = []
+        for i in range(len(self.biome_starts)):
+            start = self.biome_starts[i]
+            stop = self.biome_stops[i]
+            center_x = (start[0] + stop[0] - 1) // 2
+            center_y = (start[1] + stop[1] - 1) // 2
+            biome_centers.append((center_x, center_y))
+        self.biome_centers_jax = jnp.array(biome_centers)
         self.biome_masks = []
         for i in range(self.biome_object_frequencies.shape[0]):
             # Create mask for the biome
@@ -174,6 +186,23 @@ class ForagaxEnv(environment.Environment):
         is_blocking = self.object_blocking[obj_at_new_pos]
         pos = jax.lax.select(is_blocking, state.pos, new_pos)
 
+        # Check for automatic teleport
+        if self.teleport_interval is not None:
+            should_teleport = jnp.mod(state.time + 1, self.teleport_interval) == 0
+        else:
+            should_teleport = False
+
+        def teleport_fn():
+            # Calculate squared distances from current position to each biome center
+            diffs = self.biome_centers_jax - pos
+            distances = jnp.sum(diffs**2, axis=1)
+            # Find the index of the furthest biome center
+            furthest_idx = jnp.argmax(distances)
+            new_pos = self.biome_centers_jax[furthest_idx]
+            return new_pos
+
+        pos = jax.lax.cond(should_teleport, teleport_fn, lambda: pos)
+
         # 2. HANDLE COLLISIONS AND REWARDS
         obj_at_pos = current_objects[pos[1], pos[0]]
         key, subkey = jax.random.split(key)
@@ -503,6 +532,7 @@ class ForagaxObjectEnv(ForagaxEnv):
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
         deterministic_spawn: bool = False,
+        teleport_interval: Optional[int] = None,
     ):
         super().__init__(
             name,
@@ -512,6 +542,7 @@ class ForagaxObjectEnv(ForagaxEnv):
             biomes,
             nowrap,
             deterministic_spawn,
+            teleport_interval,
         )
 
         # Compute unique colors and mapping for partial observability

foragax/objects.py

@@ -318,6 +318,7 @@ def create_weather_objects(
     repeat: int = 500,
     multiplier: float = 1.0,
     same_color: bool = False,
+    random_respawn: bool = False,
 ):
     """Create HOT and COLD WeatherObject instances using the specified file.
 
@@ -346,6 +347,7 @@ def create_weather_objects(
         repeat=repeat,
         multiplier=multiplier,
         color=hot_color,
+        random_respawn=random_respawn,
     )
 
     cold_color = hot_color if same_color else (0, 255, 255)
@@ -355,6 +357,7 @@ def create_weather_objects(
         repeat=repeat,
         multiplier=-multiplier,
         color=cold_color,
+        random_respawn=random_respawn,
     )
 
     return hot, cold

foragax/registry.py

@@ -64,6 +64,27 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "biomes": None,
         "nowrap": True,
     },
+    "ForagaxWeather-v4": {
+        "size": None,
+        "aperture_size": None,
+        "objects": None,
+        "biomes": None,
+        "nowrap": True,
+        "deterministic_spawn": True,
+    },
+    "ForagaxWeather-v5": {
+        "size": (15, 15),
+        "aperture_size": None,
+        "objects": None,
+        "biomes": (
+            # Hot biome
+            Biome(start=(0, 3), stop=(15, 5), object_frequencies=(0.5, 0.0)),
+            # Cold biome
+            Biome(start=(0, 10), stop=(15, 12), object_frequencies=(0.0, 0.5)),
+        ),
+        "nowrap": False,
+        "deterministic_spawn": True,
+    },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
         "aperture_size": None,
@@ -272,6 +293,19 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "nowrap": True,
         "deterministic_spawn": True,
     },
+    "ForagaxTwoBiome-v16": {
+        "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,
@@ -348,6 +382,7 @@ def make(
         "ForagaxTwoBiome-v9",
         "ForagaxTwoBiome-v10",
         "ForagaxTwoBiome-v15",
+        "ForagaxTwoBiome-v16",
     ):
         margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
@@ -386,7 +421,7 @@ def make(
             ),
         )
 
-    if env_id == "ForagaxWeather-v3":
+    if env_id in ("ForagaxWeather-v3", "ForagaxWeather-v4"):
         margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
         config["size"] = (15, width)
@@ -406,10 +441,21 @@ def make(
         )
 
     if env_id.startswith("ForagaxWeather"):
-        same_color = env_id in ("ForagaxWeather-v2", "ForagaxWeather-v3")
-        hot, cold = create_weather_objects(file_index=file_index, same_color=same_color)
+        same_color = env_id in (
+            "ForagaxWeather-v2",
+            "ForagaxWeather-v3",
+            "ForagaxWeather-v4",
+            "ForagaxWeather-v5",
+        )
+        random_respawn = env_id in ("ForagaxWeather-v4", "ForagaxWeather-v5")
+        hot, cold = create_weather_objects(
+            file_index=file_index, same_color=same_color, random_respawn=random_respawn
+        )
         config["objects"] = (hot, cold)
 
+    if env_id == "ForagaxTwoBiome-v16":
+        config["teleport_interval"] = 10000
+
     env_class_map = {
         "object": ForagaxObjectEnv,
         "rgb": ForagaxRGBEnv,