continual-foragax 0.34.0__tar.gz → 0.35.0__tar.gz

{continual_foragax-0.34.0 → continual_foragax-0.35.0}/PKG-INFO

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

{continual_foragax-0.34.0 → continual_foragax-0.35.0}/pyproject.toml

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

{continual_foragax-0.34.0 → continual_foragax-0.35.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.34.0 → continual_foragax-0.35.0}/src/foragax/env.py

@@ -690,11 +690,16 @@ class ForagaxEnv(environment.Environment):
 
         num_biomes = self.biome_object_frequencies.shape[0]
 
-        # Compute consumption rates for all biomes
-        consumption_rates = biome_state.consumption_count / jnp.maximum(
-            1.0, biome_state.total_objects.astype(float)
-        )
-        should_respawn = consumption_rates >= self.biome_consumption_threshold
+        if isinstance(self.biome_consumption_threshold, float):
+            # Compute consumption rates for all biomes
+            consumption_rates = biome_state.consumption_count / jnp.maximum(
+                1.0, biome_state.total_objects.astype(float)
+            )
+            should_respawn = consumption_rates >= self.biome_consumption_threshold
+        else:
+            should_respawn = (
+                biome_state.consumption_count >= self.biome_consumption_threshold
+            )
 
         # Split key for all biomes in parallel
         key, subkey = jax.random.split(key)

{continual_foragax-0.34.0 → continual_foragax-0.35.0}/src/foragax/objects.py

@@ -240,6 +240,7 @@ class FourierObject(BaseForagaxObject):
         color: Tuple[int, int, int] = (0, 0, 0),
         reward_delay: int = 0,
         max_reward_delay: Optional[int] = None,
+        regen_delay: Optional[Tuple[int, int]] = None,
     ):
         if max_reward_delay is None:
             max_reward_delay = reward_delay
@@ -248,13 +249,14 @@ class FourierObject(BaseForagaxObject):
             blocking=False,
             collectable=True,
             color=color,
-            random_respawn=False,  # Objects don't respawn individually
+            random_respawn=True,
             max_reward_delay=max_reward_delay,
             expiry_time=None,
         )
         self.num_fourier_terms = num_fourier_terms
         self.base_magnitude = base_magnitude
         self.reward_delay_val = reward_delay
+        self.regen_delay_range = regen_delay
 
     def get_state(self, key: jax.Array) -> jax.Array:
         """Generate random Fourier series parameters.
@@ -353,6 +355,9 @@ class FourierObject(BaseForagaxObject):
 
     def regen_delay(self, clock: int, rng: jax.Array) -> int:
         """No individual regeneration - returns infinity."""
+        if self.regen_delay_range is not None:
+            min_delay, max_delay = self.regen_delay_range
+            return jax.random.randint(rng, (), min_delay, max_delay)
         return jnp.iinfo(jnp.int32).max
 
     def expiry_regen_delay(self, clock: int, rng: jax.Array) -> int:
@@ -709,6 +714,7 @@ def create_fourier_objects(
     num_fourier_terms: int = 10,
     base_magnitude: float = 1.0,
     reward_delay: int = 0,
+    regen_delay: Optional[Tuple[int, int]] = None,
 ):
     """Create HOT and COLD FourierObject instances.
 
@@ -726,6 +732,7 @@ def create_fourier_objects(
         base_magnitude=base_magnitude,
         color=(0, 0, 0),
         reward_delay=reward_delay,
+        regen_delay=regen_delay,
     )
 
     cold = FourierObject(
@@ -734,6 +741,7 @@ def create_fourier_objects(
         base_magnitude=base_magnitude,
         color=(0, 0, 0),
         reward_delay=reward_delay,
+        regen_delay=regen_delay,
     )
 
     return hot, cold

{continual_foragax-0.34.0 → continual_foragax-0.35.0}/src/foragax/registry.py

@@ -104,6 +104,21 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "dynamic_biomes": True,
         "biome_consumption_threshold": 0.9,
     },
+    "ForagaxDiwali-v2": {
+        "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": 200,
+    },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
         "aperture_size": None,
@@ -539,6 +554,12 @@ def make(
             num_fourier_terms=10,
             reward_delay=reward_delay,
         )
+    if env_id == "ForagaxDiwali-v2":
+        config["objects"] = create_fourier_objects(
+            num_fourier_terms=10,
+            reward_delay=reward_delay,
+            regen_delay=(9, 11),
+        )
 
     if env_id == "ForagaxSineTwoBiome-v1":
         biome1_oyster, biome1_deathcap, biome2_oyster, biome2_deathcap = (

{continual_foragax-0.34.0 → continual_foragax-0.35.0}/tests/test_registry.py

@@ -586,6 +586,63 @@ def test_foragax_diwali_v1_creation():
     assert jnp.all(state.biome_state.generation == 0), "Initial generation should be 0"
 
 
+def test_foragax_diwali_v2_creation():
+    """Test that ForagaxDiwali-v2 can be created and initialized."""
+    env = make("ForagaxDiwali-v2", observation_type="object", aperture_size=(5, 5))
+
+    # Check that dynamic biomes are enabled
+    assert env.dynamic_biomes, "Dynamic biomes should be enabled for ForagaxDiwali-v2"
+    assert env.biome_consumption_threshold == 200, "Threshold should be 200"
+    assert env.num_fourier_terms == 10, "Should have 10 Fourier terms"
+
+    # Check that objects are FourierObjects
+    assert isinstance(env.objects[1], FourierObject), (
+        "First object should be FourierObject"
+    )
+    assert isinstance(env.objects[2], FourierObject), (
+        "Second object should be FourierObject"
+    )
+
+    # Check regen_delay for objects
+    assert env.objects[1].regen_delay_range == (9, 11), (
+        "Object 1 should have regen_delay_range (9, 11)"
+    )
+    assert env.objects[2].regen_delay_range == (9, 11), (
+        "Object 2 should have regen_delay_range (9, 11)"
+    )
+
+    # Initialize environment
+    key = jax.random.key(0)
+    obs, state = env.reset(key, env.default_params)
+
+    # Check state has new fields
+    assert hasattr(state, "object_state"), "State should have object_state"
+    assert hasattr(state.object_state, "color"), "ObjectState should have color"
+    assert hasattr(state.object_state, "state_params"), (
+        "ObjectState should have state_params"
+    )
+    assert hasattr(state, "biome_state"), "State should have biome_state"
+    assert hasattr(state.biome_state, "consumption_count"), (
+        "BiomeState should have consumption_count"
+    )
+    assert hasattr(state.biome_state, "total_objects"), (
+        "BiomeState should have total_objects"
+    )
+    assert hasattr(state.biome_state, "generation"), "BiomeState should have generation"
+
+    # Check shapes
+    chex.assert_shape(state.object_state.color, (15, 15, 3))
+    chex.assert_shape(
+        state.object_state.state_params, (15, 15, 22)
+    )  # 2 + 2*10 = 22 params
+    chex.assert_shape(state.biome_state.consumption_count, (2,))  # 2 biomes
+    chex.assert_shape(state.biome_state.total_objects, (2,))
+    chex.assert_shape(state.biome_state.generation, (2,))
+
+    # Check initial generation is 0
+    assert jnp.all(state.biome_state.generation == 0), "Initial generation should be 0"
+
+
 def test_foragax_sine_twobiome_v1_creation():
     """Test that ForagaxSineTwoBiome-v1 can be created and initialized."""
     env = make("ForagaxSineTwoBiome-v1", observation_type="color", aperture_size=(5, 5))