continual-foragax 0.38.0__tar.gz → 0.40.0__tar.gz

{continual_foragax-0.38.0 → continual_foragax-0.40.0}/PKG-INFO

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

{continual_foragax-0.38.0 → continual_foragax-0.40.0}/pyproject.toml

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

{continual_foragax-0.38.0 → continual_foragax-0.40.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.38.0 → continual_foragax-0.40.0}/src/foragax/env.py

@@ -6,7 +6,7 @@ Source: https://github.com/andnp/Forager
 from dataclasses import dataclass
 from enum import IntEnum
 from functools import partial
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Tuple, Union
 
 import jax
 import jax.numpy as jnp
@@ -128,7 +128,6 @@ class ForagaxEnv(environment.Environment):
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
         deterministic_spawn: bool = False,
-        teleport_interval: Optional[int] = None,
         observation_type: str = "object",
         dynamic_biomes: bool = False,
         biome_consumption_threshold: float = 0.9,
@@ -154,7 +153,6 @@ class ForagaxEnv(environment.Environment):
         self.observation_type = observation_type
         self.nowrap = nowrap
         self.deterministic_spawn = deterministic_spawn
-        self.teleport_interval = teleport_interval
         self.dynamic_biomes = dynamic_biomes
         self.biome_consumption_threshold = biome_consumption_threshold
         self.dynamic_biome_spawn_empty = dynamic_biome_spawn_empty
@@ -391,23 +389,6 @@ 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]]
         is_collectable = self.object_collectable[obj_at_pos]
@@ -774,20 +755,53 @@ class ForagaxEnv(environment.Environment):
             biome_mask = self.biome_masks_array[i]
             new_gen_value = new_biome_state.generation[i]
 
-            # Only update where new spawn has objects and biome should respawn
-            is_new_object = (
-                (all_new_objects[i] > 0) & biome_mask & should_respawn[i][..., None]
-            )
+            # Update mask: biome area AND needs respawn
+            should_update = biome_mask & should_respawn[i][..., None]
+
+            # 1. Merge: Overwrite with new objects if present, otherwise keep existing
+            # If the new spawn has an object, we take it. If it's empty, we keep whatever was there.
+            new_spawn_valid = all_new_objects[i] > 0
 
-            new_obj_id = jnp.where(is_new_object, all_new_objects[i], new_obj_id)
-            new_color = jnp.where(
-                is_new_object[..., None], all_new_colors[i], new_color
+            merged_objs = jnp.where(new_spawn_valid, all_new_objects[i], new_obj_id)
+            merged_colors = jnp.where(
+                new_spawn_valid[..., None], all_new_colors[i], new_color
             )
-            new_params = jnp.where(
-                is_new_object[..., None], all_new_params[i], new_params
+            merged_params = jnp.where(
+                new_spawn_valid[..., None], all_new_params[i], new_params
             )
-            new_gen = jnp.where(is_new_object, new_gen_value, new_gen)
-            new_spawn = jnp.where(is_new_object, current_time, new_spawn)
+
+            # For generation/spawn time, update only if we took the NEW object
+            merged_gen = jnp.where(new_spawn_valid, new_gen_value, new_gen)
+            merged_spawn = jnp.where(new_spawn_valid, current_time, new_spawn)
+
+            # 2. Apply dropout to the MERGED result (only where we are allowed to update)
+            if self.dynamic_biome_spawn_empty > 0:
+                key, dropout_key = jax.random.split(key)
+                # Dropout applies only to cells that are both in the biome AND need updating
+                keep_mask = jax.random.bernoulli(
+                    dropout_key, 1.0 - self.dynamic_biome_spawn_empty, merged_objs.shape
+                )
+                # Only apply dropout where should_update is true; elsewhere, keep merged_objs
+                dropout_mask = should_update & keep_mask
+                # Apply dropout only to the merged result and associated metadata
+                final_objs = jnp.where(dropout_mask, merged_objs, 0)
+                final_colors = jnp.where(dropout_mask[..., None], merged_colors, 0)
+                final_params = jnp.where(dropout_mask[..., None], merged_params, 0)
+                final_gen = jnp.where(dropout_mask, merged_gen, 0)
+                final_spawn = jnp.where(dropout_mask, merged_spawn, 0)
+            else:
+                final_objs = merged_objs
+                final_colors = merged_colors
+                final_params = merged_params
+                final_gen = merged_gen
+                final_spawn = merged_spawn
+
+            # 3. Write back: Only update where should_update is true
+            new_obj_id = jnp.where(should_update, final_objs, new_obj_id)
+            new_color = jnp.where(should_update[..., None], final_colors, new_color)
+            new_params = jnp.where(should_update[..., None], final_params, new_params)
+            new_gen = jnp.where(should_update, final_gen, new_gen)
+            new_spawn = jnp.where(should_update, final_spawn, new_spawn)
 
         # Clear timers in respawning biomes
         new_respawn_timer = object_state.respawn_timer
@@ -890,7 +904,7 @@ class ForagaxEnv(environment.Environment):
         biome_freqs = self.biome_object_frequencies[biome_idx]
         biome_mask = self.biome_masks_array[biome_idx]
 
-        key, spawn_key, color_key, params_key, dropout_key = jax.random.split(key, 5)
+        key, spawn_key, color_key, params_key = jax.random.split(key, 4)
 
         # Generate object IDs using deterministic or random spawn
         if deterministic:
@@ -929,13 +943,6 @@ class ForagaxEnv(environment.Environment):
                 jnp.searchsorted(cumulative_freqs, grid_rand, side="right") - 1
             )
 
-        # Apply random dropout if configured
-        if self.dynamic_biome_spawn_empty > 0:
-            dropout_mask = jax.random.bernoulli(
-                dropout_key, 1.0 - self.dynamic_biome_spawn_empty, object_grid.shape
-            )
-            object_grid = jnp.where(dropout_mask, object_grid, 0)
-
         # Initialize color grid
         color_grid = jnp.full((self.size[1], self.size[0], 3), 255, dtype=jnp.uint8)
 

{continual_foragax-0.38.0 → continual_foragax-0.40.0}/src/foragax/registry.py

@@ -147,12 +147,12 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "size": (15, 15),
         "aperture_size": None,
         "objects": None,
-        "biomes": (Biome(start=(0, 0), stop=(15, 15), object_frequencies=(0.2,)),),
+        "biomes": (Biome(start=(0, 0), stop=(15, 15), object_frequencies=(0.4,)),),
         "nowrap": False,
         "deterministic_spawn": True,
         "dynamic_biomes": True,
         "biome_consumption_threshold": 1000,
-        "dynamic_biome_spawn_empty": 0.1,
+        "dynamic_biome_spawn_empty": 0.4,
     },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
@@ -636,9 +636,6 @@ def make(
             biome2_deathcap,
         )
 
-    if env_id == "ForagaxTwoBiome-v16":
-        config["teleport_interval"] = 10000
-
     # Backward compatibility: map "world" to "object" with full world
     if observation_type == "world":
         # add deprecation warning

{continual_foragax-0.38.0 → continual_foragax-0.40.0}/tests/test_benchmark.py

@@ -2,6 +2,7 @@ import chex
 import jax
 import jax.numpy as jnp
 
+from foragax import registry
 from foragax.env import Actions, Biome, ForagaxEnv
 from foragax.objects import FLOWER, WALL
 
@@ -40,7 +41,7 @@ def test_benchmark_vision(benchmark):
     benchmark(benchmark_fn)
 
 
-def test_benchmark_creation(benchmark):
+def test_benchmark_reset(benchmark):
     env = ForagaxEnv(
         size=1_000,
         aperture_size=31,
@@ -63,6 +64,40 @@ def test_benchmark_creation(benchmark):
     benchmark(benchmark_fn)
 
 
+def test_benchmark_tiny_env(benchmark):
+    env = ForagaxEnv(
+        size=15,
+        aperture_size=11,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.1, 0.1)),),
+        observation_type="color",
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+    key, reset_key = jax.random.split(key)
+    _, state = env.reset(reset_key, params)
+
+    @jax.jit
+    def _run(state, key):
+        def f(carry, _):
+            state, key = carry
+            key, step_key = jax.random.split(key, 2)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+            return (new_state, key), None
+
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
+        return final_state
+
+    key, run_key = jax.random.split(key)
+    _run(state, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(state, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
 def test_benchmark_small_env(benchmark):
     env = ForagaxEnv(
         size=1_000,
@@ -84,7 +119,7 @@ def test_benchmark_small_env(benchmark):
             _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
-        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=1000)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
         return final_state
 
     key, run_key = jax.random.split(key)
@@ -120,7 +155,7 @@ def test_benchmark_big_env(benchmark):
             _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
-        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
         return final_state
 
     # warm-up compilation
@@ -136,7 +171,7 @@ def test_benchmark_big_env(benchmark):
 
 
 def test_benchmark_vmap_env(benchmark):
-    num_envs = 100
+    num_envs = 10
     env = ForagaxEnv(
         size=1_000,
         aperture_size=11,
@@ -163,7 +198,7 @@ def test_benchmark_vmap_env(benchmark):
             )
             return (new_states, key), None
 
-        (final_states, _), _ = jax.lax.scan(f, (states, key), None, length=1000)
+        (final_states, _), _ = jax.lax.scan(f, (states, key), None, length=10)
         return final_states
 
     # warm-up compilation
@@ -199,7 +234,7 @@ def test_benchmark_small_env_color(benchmark):
             _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
-        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
         return final_state
 
     key, run_key = jax.random.split(key)
@@ -233,7 +268,63 @@ def test_benchmark_small_env_world(benchmark):
             _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
             return (new_state, key), None
 
-        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
+        return final_state
+
+    key, run_key = jax.random.split(key)
+    _run(state, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(state, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_diwali_v5(benchmark):
+    env = registry.make("ForagaxDiwali-v5")
+    params = env.default_params
+    key = jax.random.key(0)
+    key, reset_key = jax.random.split(key)
+    _, state = env.reset(reset_key, params)
+
+    @jax.jit
+    def _run(state, key):
+        def f(carry, _):
+            state, key = carry
+            key, step_key = jax.random.split(key, 2)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+            return (new_state, key), None
+
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
+        return final_state
+
+    key, run_key = jax.random.split(key)
+    _run(state, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(state, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_sine_two_biome_v1(benchmark):
+    env = registry.make("ForagaxSineTwoBiome-v1")
+    params = env.default_params
+    key = jax.random.key(0)
+    key, reset_key = jax.random.split(key)
+    _, state = env.reset(reset_key, params)
+
+    @jax.jit
+    def _run(state, key):
+        def f(carry, _):
+            state, key = carry
+            key, step_key = jax.random.split(key, 2)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+            return (new_state, key), None
+
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
         return final_state
 
     key, run_key = jax.random.split(key)

{continual_foragax-0.38.0 → continual_foragax-0.40.0}/tests/test_foragax.py

@@ -977,66 +977,6 @@ def test_single_color_all_objects():
         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 = ForagaxEnv(
-        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,
-        observation_type="color",
-    )
-    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)
-
-
 def test_info_discount():
     """Test that info contains discount."""
     key = jax.random.key(0)