continual-foragax 0.38.0__py3-none-any.whl → 0.40.0__py3-none-any.whl

{continual_foragax-0.38.0.dist-info → continual_foragax-0.40.0.dist-info}/METADATA

@@ -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.dist-info → continual_foragax-0.40.0.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=9EB02GgSc80KRjj35z_UyM4VXTGuyHBAHX3d9M5AvY0,66966
+foragax/env.py,sha256=-lBBh70JRInOFR6mmU8uyA9nzgIIPapIdY20y4-vwQU,67848
 foragax/objects.py,sha256=aVc7lD3CTyRP9wm_Vs93qo4l_B1kbiYGKPtkd_SVXjs,27061
-foragax/registry.py,sha256=u8mv5mPTHKRsIE7FTi36J5tXp3FuVJ-AmJRb7HWW9hI,20636
+foragax/registry.py,sha256=2qi7Dq96RZVyBQwAnejdBztQWJOQHNAw26UUPBTzIEY,20551
 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.38.0.dist-info/METADATA,sha256=ODZtbDZHv2GRj0yRC4rR34Hu0PaN_8e4Oo8XPmhWGxU,4713
-continual_foragax-0.38.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.38.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.38.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.38.0.dist-info/RECORD,,
+continual_foragax-0.40.0.dist-info/METADATA,sha256=DN8INomefM2v426Q75KtXL-VqD_fsFKAjQMaphc9Rl4,4713
+continual_foragax-0.40.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.40.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.40.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.40.0.dist-info/RECORD,,

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)
 

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