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

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

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

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

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

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

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

{continual_foragax-0.37.0 → continual_foragax-0.39.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,
+        dynamic_biome_spawn_empty: float = 0.0,
         max_expiries_per_step: int = 1,
     ):
         super().__init__()
@@ -156,6 +157,7 @@ class ForagaxEnv(environment.Environment):
         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
         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
@@ -772,20 +774,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

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

@@ -143,6 +143,17 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "dynamic_biomes": True,
         "biome_consumption_threshold": 1000,
     },
+    "ForagaxDiwali-v5": {
+        "size": (15, 15),
+        "aperture_size": None,
+        "objects": None,
+        "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.4,
+    },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
         "aperture_size": None,
@@ -597,6 +608,13 @@ def make(
             regen_delay=(9, 11),
             reward_repeat=100,
         )
+    if env_id == "ForagaxDiwali-v5":
+        config["objects"] = create_fourier_objects(
+            num_fourier_terms=10,
+            reward_delay=reward_delay,
+            regen_delay=(9, 11),
+            reward_repeat=100,
+        )[:1]
 
     if env_id == "ForagaxSineTwoBiome-v1":
         biome1_oyster, biome1_deathcap, biome2_oyster, biome2_deathcap = (

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

@@ -2811,13 +2811,14 @@ def test_fourier_reward_mathematical_properties():
     key = jax.random.key(42)
 
     # Create a simple single-term Fourier series: a1*cos(t) + b1*sin(t)
+    # We use period=8.0 so that clock=1 corresponds to phase 2π/8 = π/4
     simple_params = jnp.array(
         [
-            2 * jnp.pi,
+            8.0,  # period
             -1.0,
-            1.0,  # period=2π, min=-1, max=1
+            1.0,  # min=-1, max=1
             0.5,
-            0.5,  # a1=0.5, b1=0.5 (should give 0.5*(cos(t) + sin(t)))
+            0.5,  # a1=0.5, b1=0.5
             0.0,
             0.0,
             0.0,
@@ -2828,7 +2829,8 @@ def test_fourier_reward_mathematical_properties():
         dtype=jnp.float32,
     )
 
-    # At t=0: 0.5*(cos(0) + sin(0)) = 0.5*(1 + 0) = 0.5
+    # At t=0 (clock=0): phase=0
+    # 0.5*(cos(0) + sin(0)) = 0.5*(1 + 0) = 0.5
     reward_t0_simple = fourier_obj.reward(0, key, simple_params)
     # The actual reward is normalized: 2*(raw - min)/(max - min) - 1, then * base_magnitude
     # raw = 0.5, min=-1, max=1, so: 2*(0.5 - (-1))/(1 - (-1)) - 1 = 2*(1.5)/2 - 1 = 1.5 - 1 = 0.5
@@ -2837,9 +2839,10 @@ def test_fourier_reward_mathematical_properties():
         f"Expected ~1.0, got {reward_t0_simple}"
     )
 
-    # At t=π/4: 0.5*(cos(π/4) + sin(π/4)) = 0.5*(0.707 + 0.707) = 0.5*1.414 ≈ 0.707
+    # At t=1 (clock=1): phase = 2π * 1/8 = π/4
+    # 0.5*(cos(π/4) + sin(π/4)) = 0.5*(0.707 + 0.707) = 0.5*1.414 ≈ 0.707
+    reward_t_quarter = fourier_obj.reward(1, key, simple_params)
     t_quarter = jnp.pi / 4
-    reward_t_quarter = fourier_obj.reward(t_quarter, key, simple_params)
     raw_quarter = 0.5 * (jnp.cos(t_quarter) + jnp.sin(t_quarter))
     expected_quarter = (
         2 * (raw_quarter - (-1)) / (1 - (-1)) - 1