continual-foragax 0.36.0__tar.gz → 0.37.0__tar.gz

{continual_foragax-0.36.0 → continual_foragax-0.37.0}/PKG-INFO

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

{continual_foragax-0.36.0 → continual_foragax-0.37.0}/pyproject.toml

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

{continual_foragax-0.36.0 → continual_foragax-0.37.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.36.0 → continual_foragax-0.37.0}/src/foragax/env.py

@@ -132,6 +132,7 @@ class ForagaxEnv(environment.Environment):
         observation_type: str = "object",
         dynamic_biomes: bool = False,
         biome_consumption_threshold: float = 0.9,
+        max_expiries_per_step: int = 1,
     ):
         super().__init__()
         self._name = name
@@ -155,6 +156,9 @@ class ForagaxEnv(environment.Environment):
         self.teleport_interval = teleport_interval
         self.dynamic_biomes = dynamic_biomes
         self.biome_consumption_threshold = biome_consumption_threshold
+        if max_expiries_per_step < 1:
+            raise ValueError("max_expiries_per_step must be at least 1")
+        self.max_expiries_per_step = max_expiries_per_step
 
         objects = (EMPTY,) + objects
         if self.nowrap and not self.full_world:
@@ -608,63 +612,63 @@ class ForagaxEnv(environment.Environment):
                 & (current_objects_for_expiry > 0)
             )
 
-            # Count how many objects actually need to expire
-            num_expiring = jnp.sum(should_expire)
-
             # Only process expiry if there are actually objects to expire
-            def process_expiries():
-                # Get positions of objects that should expire
-                # Use nonzero with fixed size to maintain JIT compatibility
-                max_objects = self.size[0] * self.size[1]
-                y_indices, x_indices = jnp.nonzero(
-                    should_expire, size=max_objects, fill_value=-1
-                )
+            has_expiring = jnp.any(should_expire)
+
+            # Precompute the first expiring index in flat space so the work inside cond is minimal.
+            overage = jnp.where(
+                should_expire,
+                object_ages - expiry_times,
+                -jnp.inf,
+            ).reshape(-1)
+            sorted_flat_indices = jnp.argsort(overage)[::-1]
+            selected_flat_indices = jnp.where(
+                overage[sorted_flat_indices] > -jnp.inf,
+                sorted_flat_indices,
+                -jnp.ones_like(sorted_flat_indices),
+            )[: self.max_expiries_per_step]
 
+            def process_expiries():
                 key_local, expiry_key = jax.random.split(key)
 
-                def process_one_expiry(carry, i):
-                    obj_state = carry
-                    y = y_indices[i]
-                    x = x_indices[i]
+                def body_fn(i, obj_state):
+                    flat_idx = selected_flat_indices[i]
 
-                    # Skip if this is a padding index (from fill_value)
-                    is_valid = (y >= 0) & (x >= 0)
-
-                    def expire_one():
+                    def expire_at(obj_state):
+                        y = flat_idx // self.size[0]
+                        x = flat_idx % self.size[0]
                         obj_id = current_objects_for_expiry[y, x]
-                        exp_key = jax.random.fold_in(expiry_key, y * self.size[0] + x)
+                        exp_key = jax.random.fold_in(expiry_key, flat_idx)
                         exp_delay = jax.lax.switch(
                             obj_id, self.expiry_regen_delay_fns, state.time, exp_key
                         )
                         timer_countdown = jax.lax.cond(
                             exp_delay == jnp.iinfo(jnp.int32).max,
-                            lambda: 0,  # No timer (permanent removal)
-                            lambda: exp_delay + 1,  # Timer countdown
+                            lambda: 0,
+                            lambda: exp_delay + 1,
                         )
 
-                        # Use unified timer placement method
-                        rand_key = jax.random.split(exp_key)[1]
-                        new_obj_state = self._place_timer(
+                        respawn_random = self.object_random_respawn[obj_id]
+                        rand_key = jax.random.fold_in(exp_key, 1)
+                        return self._place_timer(
                             obj_state,
                             y,
                             x,
                             obj_id,
                             timer_countdown,
-                            self.object_random_respawn[obj_id],
+                            respawn_random,
                             rand_key,
                         )
 
-                        return new_obj_state
-
-                    def no_op():
-                        return obj_state
-
-                    return jax.lax.cond(is_valid, expire_one, no_op), None
+                    return jax.lax.cond(
+                        flat_idx >= 0,
+                        expire_at,
+                        lambda obj_state: obj_state,
+                        obj_state,
+                    )
 
-                new_object_state, _ = jax.lax.scan(
-                    process_one_expiry,
-                    object_state,
-                    jnp.arange(max_objects),
+                new_object_state = jax.lax.fori_loop(
+                    0, self.max_expiries_per_step, body_fn, object_state
                 )
                 return key_local, new_object_state
 
@@ -672,7 +676,7 @@ class ForagaxEnv(environment.Environment):
                 return key, object_state
 
             key, object_state = jax.lax.cond(
-                num_expiring > 0,
+                has_expiring,
                 process_expiries,
                 no_expiries,
             )

{continual_foragax-0.36.0 → continual_foragax-0.37.0}/src/foragax/objects.py

@@ -241,6 +241,7 @@ class FourierObject(BaseForagaxObject):
         reward_delay: int = 0,
         max_reward_delay: Optional[int] = None,
         regen_delay: Optional[Tuple[int, int]] = None,
+        reward_repeat: int = 1,
     ):
         if max_reward_delay is None:
             max_reward_delay = reward_delay
@@ -257,6 +258,7 @@ class FourierObject(BaseForagaxObject):
         self.base_magnitude = base_magnitude
         self.reward_delay_val = reward_delay
         self.regen_delay_range = regen_delay
+        self.reward_repeat = reward_repeat
 
     def get_state(self, key: jax.Array) -> jax.Array:
         """Generate random Fourier series parameters.
@@ -321,7 +323,7 @@ class FourierObject(BaseForagaxObject):
         y_max = params[2]
 
         # Normalize time to [0, 2π] using the object's period
-        t = 2.0 * jnp.pi * (clock % period) / period
+        t = 2.0 * jnp.pi * ((clock // self.reward_repeat) % period) / period
 
         # Extract interleaved coefficients: [a1, b1, a2, b2, ...]
         ab_coeffs = params[3:]
@@ -715,6 +717,7 @@ def create_fourier_objects(
     base_magnitude: float = 1.0,
     reward_delay: int = 0,
     regen_delay: Optional[Tuple[int, int]] = None,
+    reward_repeat: int = 1,
 ):
     """Create HOT and COLD FourierObject instances.
 
@@ -733,6 +736,7 @@ def create_fourier_objects(
         color=(0, 0, 0),
         reward_delay=reward_delay,
         regen_delay=regen_delay,
+        reward_repeat=reward_repeat,
     )
 
     cold = FourierObject(
@@ -742,6 +746,7 @@ def create_fourier_objects(
         color=(0, 0, 0),
         reward_delay=reward_delay,
         regen_delay=regen_delay,
+        reward_repeat=reward_repeat,
     )
 
     return hot, cold

{continual_foragax-0.36.0 → continual_foragax-0.37.0}/src/foragax/registry.py

@@ -128,6 +128,21 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "deterministic_spawn": True,
         "dynamic_biomes": True,
     },
+    "ForagaxDiwali-v4": {
+        "size": (15, 15),
+        "aperture_size": None,
+        "objects": None,
+        "biomes": (
+            # Hot biome
+            Biome(start=(0, 2), stop=(15, 6), object_frequencies=(0.5, 0.0)),
+            # Cold biome
+            Biome(start=(0, 9), stop=(15, 13), object_frequencies=(0.0, 0.5)),
+        ),
+        "nowrap": False,
+        "deterministic_spawn": True,
+        "dynamic_biomes": True,
+        "biome_consumption_threshold": 1000,
+    },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
         "aperture_size": None,
@@ -575,6 +590,13 @@ def make(
             reward_delay=reward_delay,
             regen_delay=(9, 11),
         )[:1]
+    if env_id == "ForagaxDiwali-v4":
+        config["objects"] = create_fourier_objects(
+            num_fourier_terms=10,
+            reward_delay=reward_delay,
+            regen_delay=(9, 11),
+            reward_repeat=100,
+        )
 
     if env_id == "ForagaxSineTwoBiome-v1":
         biome1_oyster, biome1_deathcap, biome2_oyster, biome2_deathcap = (

{continual_foragax-0.36.0 → continual_foragax-0.37.0}/tests/test_foragax.py

@@ -1268,6 +1268,7 @@ def test_basic_object_expiry():
         objects=(expiring_object,),
         biomes=(Biome(object_frequencies=(1.0,)),),
         observation_type="object",
+        max_expiries_per_step=10,  # Ensure all expiries processed each step
     )
 
     key = jax.random.key(0)
@@ -1316,6 +1317,112 @@ def test_basic_object_expiry():
     assert state.time == 9  # Current time after step completes
 
 
+def test_basic_object_expiry_one():
+    """Test that objects expire and respawn after expiry_time steps."""
+    # Create an object that expires after 5 steps with expiry_regen_delay=(2, 2)
+    # Note: Due to timer encoding, delay=2 actually takes 3 steps (consistent with regen_delay)
+    expiring_object = DefaultForagaxObject(
+        name="expiring",
+        reward=1.0,
+        collectable=False,  # Not collectable, so it won't be removed by collection
+        color=(255, 0, 0),
+        expiry_time=10,
+        expiry_regen_delay=(8, 8),  # Takes 9 steps due to +1 in timer encoding
+    )
+
+    env = ForagaxEnv(
+        size=(3, 3),
+        aperture_size=(3, 3),
+        objects=(expiring_object,),
+        biomes=(Biome(object_frequencies=(1.0,)),),
+        observation_type="object",
+    )
+
+    key = jax.random.key(0)
+    key, key_reset = jax.random.split(key)
+    obs, state = env.reset(key_reset, env.default_params)
+
+    # Initial state - all objects present
+    assert jnp.all(state.object_state.object_id == 1)
+    assert jnp.all(state.object_state.spawn_time == 0)
+
+    # Step through 10 times - objects should still be present
+    for _ in range(10):
+        key, key_step = jax.random.split(key)
+        obs, state, _, _, _ = env.step(
+            key_step, state, Actions.DOWN, env.default_params
+        )
+
+    assert jnp.all(state.object_state.object_id == 1)
+    assert state.time == 10
+
+    # Step once more - objects should expire and become timers
+    key, key_step = jax.random.split(key)
+    obs, state, _, _, _ = env.step(key_step, state, Actions.DOWN, env.default_params)
+
+    assert (
+        jnp.count_nonzero(state.object_state.object_id == 0) == 1
+    )  # One objects removed
+    assert (
+        jnp.count_nonzero(state.object_state.respawn_timer > 0) == 1
+    )  # One have timers
+    assert (
+        jnp.count_nonzero(state.object_state.respawn_object_id == 1) == 1
+    )  # Will respawn as object 1
+    assert jnp.all(
+        state.object_state.spawn_time == 0
+    )  # Spawn time NOT updated yet (still at reset value)
+
+    # Step through all removals and they all become timers (8 more steps)
+    for _ in range(8):
+        key, key_step = jax.random.split(key)
+        obs, state, _, _, _ = env.step(
+            key_step, state, Actions.DOWN, env.default_params
+        )
+
+    assert jnp.all(state.object_state.object_id == 0)  # All objects removed
+    assert jnp.all(state.object_state.respawn_timer > 0)  # All have timers
+    assert jnp.all(
+        state.object_state.respawn_object_id == 1
+    )  # Will respawn as object 1
+    assert jnp.all(
+        state.object_state.spawn_time == 0
+    )  # Spawn time NOT updated yet (still at reset value)
+
+    # Step once more - one object should respawn
+    key, key_step = jax.random.split(key)
+    obs, state, _, _, _ = env.step(key_step, state, Actions.DOWN, env.default_params)
+
+    assert (
+        jnp.count_nonzero(state.object_state.object_id == 0) == 8
+    )  # One objects present
+    assert (
+        jnp.count_nonzero(state.object_state.respawn_timer > 0) == 8
+    )  # One don't have timers
+    assert (
+        jnp.count_nonzero(state.object_state.respawn_object_id == 1) == 8
+    )  # Remaining will respawn as object 1
+    assert (
+        jnp.count_nonzero(state.object_state.spawn_time == 19) == 1
+    )  # One with updated spawn time
+
+    # Step through regen delay for remaining (8 more objects)
+    for _ in range(8):
+        key, key_step = jax.random.split(key)
+        obs, state, _, _, _ = env.step(
+            key_step, state, Actions.DOWN, env.default_params
+        )
+
+    assert jnp.all(state.object_state.object_id == 1)  # Objects respawned
+    assert jnp.all(state.object_state.respawn_timer == 0)  # Timers cleared
+    print(state.object_state.spawn_time)
+    assert jnp.all(
+        jnp.sort(state.object_state.spawn_time.flatten())
+        == jnp.sort(jnp.arange(19, 28))
+    )  # Spawn time = time at beginning of step when respawn occurred
+    assert state.time == 28  # Current time after step completes
+
+
 def test_no_expiry_backwards_compatibility():
     """Test that objects without expiry_time don't expire."""
     non_expiring_object = DefaultForagaxObject(
@@ -1412,6 +1519,7 @@ def test_expiry_with_normal_regen():
         objects=(normal_expiring,),
         biomes=(Biome(object_frequencies=(1.0,)),),
         observation_type="object",
+        max_expiries_per_step=16,  # Ensure all expiries processed each step
     )
 
     key = jax.random.key(456)
@@ -1506,6 +1614,7 @@ def test_expiry_spawn_time_tracking():
         objects=(expiring_obj,),
         biomes=(Biome(object_frequencies=(1.0,)),),
         observation_type="object",
+        max_expiries_per_step=4,  # Ensure all expiries processed each step
     )
 
     key = jax.random.key(111)
@@ -1878,6 +1987,7 @@ def test_object_color_grid_cleared_on_expiry():
         biomes=(Biome(object_frequencies=(0.5,)),),
         observation_type="color",
         dynamic_biomes=True,  # Enable dynamic biomes to test color grid
+        max_expiries_per_step=25,  # Ensure all expiries processed each step
     )
 
     key = jax.random.key(42)