continual-foragax 0.29.0__py3-none-any.whl → 0.30.1__py3-none-any.whl

{continual_foragax-0.29.0.dist-info → continual_foragax-0.30.1.dist-info}/METADATA

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

{continual_foragax-0.29.0.dist-info → continual_foragax-0.30.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=OFtDT8c5nskflnXmMNbRUvu5pVs9vhIezDgZaICXSyE,27535
-foragax/objects.py,sha256=j7FivgT4uz6N4FkOTmpM0t-YjTUkYUBLAznWpVqqjrU,10509
-foragax/registry.py,sha256=bFSTDo7XU4G0njHjLTgPdzWyStkAGckimDYcnGYLIIg,15529
+foragax/env.py,sha256=4NZ5JsUGjAepmzw2uxu5_ikyVZnZ7vazy062Xzx22Zg,27481
+foragax/objects.py,sha256=M0nECANGfUvvBRMKSS7akGtoO2Suv5eroI-9Aj326sw,10368
+foragax/registry.py,sha256=Dxg6cWIPwg91fNrCPxADJv35u6jFg_8dI5iTpCMFEFA,15229
 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.29.0.dist-info/METADATA,sha256=vaiiCWr06OczH8LgsRotJXRt7q4KZMtnvuFj6up5v3U,4897
-continual_foragax-0.29.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.29.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.29.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.29.0.dist-info/RECORD,,
+continual_foragax-0.30.1.dist-info/METADATA,sha256=9iwHDGT1ZbvjL_CNRFQRQsPBPWKTFBetxrJPk_OXKug,4897
+continual_foragax-0.30.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.30.1.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.30.1.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.30.1.dist-info/RECORD,,

foragax/env.py

@@ -113,12 +113,12 @@ class ForagaxEnv(environment.Environment):
 
         self.reward_fns = [o.reward for o in objects]
         self.regen_delay_fns = [o.regen_delay for o in objects]
-        self.digestion_steps_fns = [o.digestion_steps for o in objects]
+        self.reward_delay_fns = [o.reward_delay for o in objects]
 
-        # Compute digestion steps per object (using max_digestion_steps attribute)
-        object_max_digestion_steps = jnp.array([o.max_digestion_steps for o in objects])
-        self.max_digestion_steps = (
-            int(jnp.max(object_max_digestion_steps)) + 1 if len(objects) > 0 else 0
+        # Compute reward steps per object (using max_reward_delay attribute)
+        object_max_reward_delay = jnp.array([o.max_reward_delay for o in objects])
+        self.max_reward_delay = (
+            int(jnp.max(object_max_reward_delay)) + 1 if len(objects) > 0 else 0
         )
 
         self.biome_object_frequencies = jnp.array(
@@ -250,21 +250,21 @@ class ForagaxEnv(environment.Environment):
             obj_at_pos, self.reward_fns, state.time, reward_subkey
         )
         key, digestion_subkey = jax.random.split(key)
-        digestion_steps = jax.lax.switch(
-            obj_at_pos, self.digestion_steps_fns, state.time, digestion_subkey
+        reward_delay = jax.lax.switch(
+            obj_at_pos, self.reward_delay_fns, state.time, digestion_subkey
         )
-        reward = jnp.where(should_collect & (digestion_steps == 0), object_reward, 0.0)
-        if self.max_digestion_steps > 0:
+        reward = jnp.where(should_collect & (reward_delay == 0), object_reward, 0.0)
+        if self.max_reward_delay > 0:
             # Add delayed rewards to buffer
             digestion_buffer = jax.lax.cond(
-                should_collect & (digestion_steps > 0),
+                should_collect & (reward_delay > 0),
                 lambda: digestion_buffer.at[
-                    (state.time + digestion_steps) % self.max_digestion_steps
+                    (state.time + reward_delay) % self.max_reward_delay
                 ].add(object_reward),
                 lambda: digestion_buffer,
             )
             # Deliver current rewards
-            current_index = state.time % self.max_digestion_steps
+            current_index = state.time % self.max_reward_delay
             reward += digestion_buffer[current_index]
             digestion_buffer = digestion_buffer.at[current_index].set(0.0)
 
@@ -383,7 +383,7 @@ class ForagaxEnv(environment.Environment):
             object_grid=object_grid,
             biome_grid=biome_grid,
             time=0,
-            digestion_buffer=jnp.zeros((self.max_digestion_steps,)),
+            digestion_buffer=jnp.zeros((self.max_reward_delay,)),
         )
 
         return self.get_obs(state, params), state
@@ -447,7 +447,7 @@ class ForagaxEnv(environment.Environment):
                 "digestion_buffer": spaces.Box(
                     -jnp.inf,
                     jnp.inf,
-                    (self.max_digestion_steps,),
+                    (self.max_reward_delay,),
                     float,
                 ),
             }

foragax/objects.py

@@ -17,14 +17,14 @@ class BaseForagaxObject:
         collectable: bool = False,
         color: Tuple[int, int, int] = (0, 0, 0),
         random_respawn: bool = False,
-        max_digestion_steps: int = 0,
+        max_reward_delay: int = 0,
     ):
         self.name = name
         self.blocking = blocking
         self.collectable = collectable
         self.color = color
         self.random_respawn = random_respawn
-        self.max_digestion_steps = max_digestion_steps
+        self.max_reward_delay = max_reward_delay
 
     @abc.abstractmethod
     def reward(self, clock: int, rng: jax.Array) -> float:
@@ -32,8 +32,8 @@ class BaseForagaxObject:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def digestion_steps(self, clock: int, rng: jax.Array) -> int:
-        """Digestion steps function."""
+    def reward_delay(self, clock: int, rng: jax.Array) -> int:
+        """Reward delay function."""
         raise NotImplementedError
 
 
@@ -49,17 +49,17 @@ class DefaultForagaxObject(BaseForagaxObject):
         regen_delay: Tuple[int, int] = (10, 100),
         color: Tuple[int, int, int] = (255, 255, 255),
         random_respawn: bool = False,
-        digestion_steps: int = 0,
-        max_digestion_steps: Optional[int] = None,
+        reward_delay: int = 0,
+        max_reward_delay: Optional[int] = None,
     ):
-        if max_digestion_steps is None:
-            max_digestion_steps = digestion_steps
+        if max_reward_delay is None:
+            max_reward_delay = reward_delay
         super().__init__(
-            name, blocking, collectable, color, random_respawn, max_digestion_steps
+            name, blocking, collectable, color, random_respawn, max_reward_delay
         )
         self.reward_val = reward
         self.regen_delay_range = regen_delay
-        self.digestion_steps_val = digestion_steps
+        self.reward_delay_val = reward_delay
 
     def reward(self, clock: int, rng: jax.Array) -> float:
         """Default reward function."""
@@ -70,9 +70,9 @@ class DefaultForagaxObject(BaseForagaxObject):
         min_delay, max_delay = self.regen_delay_range
         return jax.random.randint(rng, (), min_delay, max_delay)
 
-    def digestion_steps(self, clock: int, rng: jax.Array) -> int:
-        """Default digestion steps function."""
-        return self.digestion_steps_val
+    def reward_delay(self, clock: int, rng: jax.Array) -> int:
+        """Default reward delay function."""
+        return self.reward_delay_val
 
 
 class NormalRegenForagaxObject(DefaultForagaxObject):
@@ -87,8 +87,8 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
         std_regen_delay: int = 1,
         color: Tuple[int, int, int] = (0, 0, 0),
         random_respawn: bool = False,
-        digestion_steps: int = 0,
-        max_digestion_steps: Optional[int] = None,
+        reward_delay: int = 0,
+        max_reward_delay: Optional[int] = None,
     ):
         super().__init__(
             name=name,
@@ -97,8 +97,8 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
             regen_delay=(mean_regen_delay, mean_regen_delay),
             color=color,
             random_respawn=random_respawn,
-            digestion_steps=digestion_steps,
-            max_digestion_steps=max_digestion_steps,
+            reward_delay=reward_delay,
+            max_reward_delay=max_reward_delay,
         )
         self.mean_regen_delay = mean_regen_delay
         self.std_regen_delay = std_regen_delay
@@ -122,8 +122,8 @@ class WeatherObject(NormalRegenForagaxObject):
         std_regen_delay: int = 1,
         color: Tuple[int, int, int] = (0, 0, 0),
         random_respawn: bool = False,
-        digestion_steps: int = 0,
-        max_digestion_steps: Optional[int] = None,
+        reward_delay: int = 0,
+        max_reward_delay: Optional[int] = None,
     ):
         super().__init__(
             name=name,
@@ -132,16 +132,15 @@ class WeatherObject(NormalRegenForagaxObject):
             std_regen_delay=std_regen_delay,
             color=color,
             random_respawn=random_respawn,
-            digestion_steps=digestion_steps,
-            max_digestion_steps=max_digestion_steps,
+            reward_delay=reward_delay,
+            max_reward_delay=max_reward_delay,
         )
-        self.rewards = rewards
+        self.rewards = rewards * multiplier
         self.repeat = repeat
-        self.multiplier = multiplier
 
     def reward(self, clock: int, rng: jax.Array) -> float:
         """Reward is based on temperature."""
-        return get_temperature(self.rewards, clock, self.repeat) * self.multiplier
+        return get_temperature(self.rewards, clock, self.repeat)
 
 
 EMPTY = DefaultForagaxObject()
@@ -340,7 +339,7 @@ def create_weather_objects(
     multiplier: float = 1.0,
     same_color: bool = False,
     random_respawn: bool = False,
-    digestion_steps: int = 0,
+    reward_delay: int = 0,
 ):
     """Create HOT and COLD WeatherObject instances using the specified file.
 
@@ -370,7 +369,7 @@ def create_weather_objects(
         multiplier=multiplier,
         color=hot_color,
         random_respawn=random_respawn,
-        digestion_steps=digestion_steps,
+        reward_delay=reward_delay,
     )
 
     cold_color = hot_color if same_color else (0, 255, 255)
@@ -381,7 +380,7 @@ def create_weather_objects(
         multiplier=-multiplier,
         color=cold_color,
         random_respawn=random_respawn,
-        digestion_steps=digestion_steps,
+        reward_delay=reward_delay,
     )
 
     return hot, cold

foragax/registry.py

@@ -83,19 +83,6 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "nowrap": False,
         "deterministic_spawn": True,
     },
-    "ForagaxWeather-v6": {
-        "size": (15, 15),
-        "aperture_size": None,
-        "objects": None,
-        "biomes": (
-            # Hot biome
-            Biome(start=(0, 3), stop=(15, 5), object_frequencies=(0.5, 0.0)),
-            # Cold biome
-            Biome(start=(0, 10), stop=(15, 12), object_frequencies=(0.0, 0.5)),
-        ),
-        "nowrap": False,
-        "deterministic_spawn": True,
-    },
     "ForagaxTwoBiome-v1": {
         "size": (15, 15),
         "aperture_size": None,
@@ -361,6 +348,8 @@ def make(
     observation_type: str = "color",
     aperture_size: Optional[Tuple[int, int]] = (5, 5),
     file_index: int = 0,
+    repeat: int = 500,
+    reward_delay: int = 0,
     **kwargs: Any,
 ) -> ForagaxEnv:
     """Create a Foragax environment.
@@ -371,6 +360,8 @@ def make(
         aperture_size: The size of the agent's observation aperture. If -1, full world observation.
             If None, the default for the environment is used.
         file_index: File index for weather objects.
+        repeat: How many steps each temperature value repeats for (weather environments).
+        reward_delay: Number of steps required to digest food items (weather environments).
         **kwargs: Additional keyword arguments to pass to the ForagaxEnv constructor.
 
     Returns:
@@ -467,19 +458,17 @@ def make(
             "ForagaxWeather-v3",
             "ForagaxWeather-v4",
             "ForagaxWeather-v5",
-            "ForagaxWeather-v6",
         )
         random_respawn = env_id in (
             "ForagaxWeather-v4",
             "ForagaxWeather-v5",
-            "ForagaxWeather-v6",
         )
-        digestion_steps = 10 if env_id in ("ForagaxWeather-v6") else 0
         hot, cold = create_weather_objects(
             file_index=file_index,
+            repeat=repeat,
             same_color=same_color,
             random_respawn=random_respawn,
-            digestion_steps=digestion_steps,
+            reward_delay=reward_delay,
         )
         config["objects"] = (hot, cold)