continual-foragax 0.35.0__py3-none-any.whl → 0.36.1__py3-none-any.whl

{continual_foragax-0.35.0.dist-info → continual_foragax-0.36.1.dist-info}/METADATA

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

{continual_foragax-0.35.0.dist-info → continual_foragax-0.36.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
 foragax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 foragax/colors.py,sha256=rqNPiywP4Nvr0POhsGpasRk-nMMTS3DOwFRUgperlUk,2065
-foragax/env.py,sha256=K3noPwdYmQlnXVjslqVzX_FIB-CnOh37mWFArQXnf_Y,66324
+foragax/env.py,sha256=7IhLCEosM_IH19sMzRnGMvTNq9tlhVg4PU52hGz9XLE,66531
 foragax/objects.py,sha256=PPuLYjD7em7GL404eSpP6q8TxF8p7JtQ1kIwh7uD_tU,26860
-foragax/registry.py,sha256=Ph_Z3O5GpIjrgvbKL-8Iq-Kc6MqfZIsF9KDDDzm7N3o,18787
+foragax/registry.py,sha256=y_MEvM0K_E4fzdp4XokzLh1D94UAuCu16wGwJV4wa_A,19289
 foragax/rendering.py,sha256=bms7wvBZTofoR-K-2QD2Ggeed7Viw8uwAEiEpEM3eSo,2768
 foragax/weather.py,sha256=KNAiwuFz8V__6G75vZIWQKPocLzXqxXn-Vt4TbHIpcA,1258
 foragax/data/ECA_non-blended_custom/TG_SOUID100897.txt,sha256=N7URbX6VlCZvCboUogYjMzy1I-0cfNPOn0QTLSHHfQ0,1776751
@@ -128,8 +128,8 @@ foragax/data/ECA_non-blended_custom/TG_SOUID156887.txt,sha256=juzTPgJoJxfqmZkorL
 foragax/data/ECA_non-blended_custom/elements.txt,sha256=OtcUBoDAHxuln79BPKGu0tsQxG_5G2BfAX3Ck130kEA,4507
 foragax/data/ECA_non-blended_custom/metadata.txt,sha256=nudnmOCy5cPJfSXt_IjyX0S5-T7NkCZREICZSimqeqc,48260
 foragax/data/ECA_non-blended_custom/sources.txt,sha256=1j3lSmINAoCMqPqFrHfZJriOz6sTYZNOhXzUwvTLas0,20857
-continual_foragax-0.35.0.dist-info/METADATA,sha256=pZ1uSNXsaYkaFaz7xvO9X7bee8Jr9RjT7nbwW3tx5ps,4713
-continual_foragax-0.35.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.35.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.35.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.35.0.dist-info/RECORD,,
+continual_foragax-0.36.1.dist-info/METADATA,sha256=SgkncBLbu9_MerET0r6Zto2YSLqfvmF4VE9vYXXh87g,4713
+continual_foragax-0.36.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.36.1.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.36.1.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.36.1.dist-info/RECORD,,

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,
             )

foragax/registry.py

@@ -119,6 +119,15 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "dynamic_biomes": True,
         "biome_consumption_threshold": 200,
     },
+    "ForagaxDiwali-v3": {
+        "size": (15, 15),
+        "aperture_size": None,
+        "objects": None,
+        "biomes": (Biome(start=(0, 0), stop=(15, 15), object_frequencies=(0.5,)),),
+        "nowrap": False,
+        "deterministic_spawn": True,
+        "dynamic_biomes": True,
+    },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
         "aperture_size": None,
@@ -560,6 +569,12 @@ def make(
             reward_delay=reward_delay,
             regen_delay=(9, 11),
         )
+    if env_id == "ForagaxDiwali-v3":
+        config["objects"] = create_fourier_objects(
+            num_fourier_terms=10,
+            reward_delay=reward_delay,
+            regen_delay=(9, 11),
+        )[:1]
 
     if env_id == "ForagaxSineTwoBiome-v1":
         biome1_oyster, biome1_deathcap, biome2_oyster, biome2_deathcap = (