continual-foragax 0.24.2__tar.gz → 0.25.0__tar.gz

{continual_foragax-0.24.2 → continual_foragax-0.25.0}/PKG-INFO

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

{continual_foragax-0.24.2 → continual_foragax-0.25.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.24.2"
+version = "0.25.0"
 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.24.2"
+version = "0.25.0"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
 

{continual_foragax-0.24.2 → continual_foragax-0.25.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.24.2 → continual_foragax-0.25.0}/src/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

{continual_foragax-0.24.2 → continual_foragax-0.25.0}/src/foragax/registry.py

@@ -272,6 +272,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 +361,7 @@ def make(
         "ForagaxTwoBiome-v9",
         "ForagaxTwoBiome-v10",
         "ForagaxTwoBiome-v15",
+        "ForagaxTwoBiome-v16",
     ):
         margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
@@ -410,6 +424,9 @@ def make(
         hot, cold = create_weather_objects(file_index=file_index, same_color=same_color)
         config["objects"] = (hot, cold)
 
+    if env_id == "ForagaxTwoBiome-v16":
+        config["teleport_interval"] = 10000
+
     env_class_map = {
         "object": ForagaxObjectEnv,
         "rgb": ForagaxRGBEnv,

{continual_foragax-0.24.2 → continual_foragax-0.25.0}/tests/test_foragax.py

@@ -906,3 +906,62 @@ def test_single_color_all_objects():
     for i in range(3):
         obj_obs = obs[1, i, :]  # Row 1, columns 0-2
         chex.assert_trees_all_equal(obj_obs, jnp.array([1.0]))
+
+
+def test_teleporting():
+    """Test automatic teleporting to the furthest biome center from current position."""
+    key = jax.random.key(0)
+
+    # Create environment with two biomes and teleport every 5 steps
+    env = ForagaxObjectEnv(
+        size=(10, 10),
+        aperture_size=(3, 3),
+        objects=(WALL,),
+        biomes=(
+            Biome(
+                start=(1, 1), stop=(5, 5), object_frequencies=(0.0,)
+            ),  # Biome 0 center at (2,2)
+            Biome(
+                start=(6, 6), stop=(10, 10), object_frequencies=(0.0,)
+            ),  # Biome 1 center at (7,7)
+        ),
+        teleport_interval=5,
+        nowrap=True,
+    )
+    params = env.default_params
+
+    # Reset and get initial state
+    obs, state = env.reset_env(key, params)
+
+    # Agent should start at center (5, 5), which is not in either biome initially
+    # But let's manually place it in biome 0 for testing
+    state = state.replace(pos=jnp.array([2, 2]))  # In biome 0
+
+    # Step 4 times (time will be 0,1,2,3,4 after these steps)
+    for i in range(4):
+        key, step_key = jax.random.split(key)
+        obs, state, _, _, _ = env.step_env(step_key, state, Actions.LEFT, params)
+
+    # After 4 steps, time=4, next step should teleport (4+1) % 5 == 0
+    key, step_key = jax.random.split(key)
+    obs, state, _, _, _ = env.step_env(step_key, state, Actions.LEFT, params)
+
+    # Should have teleported to the furthest biome center (biome 1 center)
+    # From (2,2), (7,7) is further than (2,2)
+    expected_pos = jnp.array([7, 7])
+    chex.assert_trees_all_equal(state.pos, expected_pos)
+
+    # Step another 4 times to reach time=9, then teleport back
+    # From [7,7], move right to stay in biome: [8,7], [9,7], then stay
+    for i in range(4):
+        key, step_key = jax.random.split(key)
+        obs, state, _, _, _ = env.step(step_key, state, Actions.RIGHT, params)
+
+    # After another 4 steps, time=9, next step should teleport (9+1) % 5 == 0
+    key, step_key = jax.random.split(key)
+    obs, state, _, _, _ = env.step(step_key, state, Actions.RIGHT, params)
+
+    # Should teleport to the furthest biome center (biome 0 center)
+    # From (7,7), (2,2) is further than (7,7)
+    expected_pos = jnp.array([2, 2])
+    chex.assert_trees_all_equal(state.pos, expected_pos)