continual-foragax 0.28.1__tar.gz → 0.30.0__tar.gz

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/PKG-INFO

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

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.28.1"
+version = "0.30.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.28.1"
+version = "0.30.0"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
 

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/src/foragax/env.py

@@ -61,6 +61,7 @@ class EnvState(environment.EnvState):
     object_grid: jax.Array
     biome_grid: jax.Array
     time: int
+    digestion_buffer: jax.Array
 
 
 class ForagaxEnv(environment.Environment):
@@ -102,11 +103,6 @@ class ForagaxEnv(environment.Environment):
         if self.nowrap and not self.full_world:
             objects = objects + (PADDING,)
         self.objects = objects
-        self.weather_object = None
-        for o in objects:
-            if isinstance(o, WeatherObject):
-                self.weather_object = o
-                break
 
         # JIT-compatible versions of object and biome properties
         self.object_ids = jnp.arange(len(objects))
@@ -117,6 +113,13 @@ 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.reward_delay_fns = [o.reward_delay for o in objects]
+
+        # 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(
             [b.object_frequencies for b in biomes]
@@ -237,12 +240,36 @@ class ForagaxEnv(environment.Environment):
 
         # 2. HANDLE COLLISIONS AND REWARDS
         obj_at_pos = current_objects[pos[1], pos[0]]
-        key, subkey = jax.random.split(key)
-        reward = jax.lax.switch(obj_at_pos, self.reward_fns, state.time, subkey)
         is_collectable = self.object_collectable[obj_at_pos]
+        should_collect = is_collectable & (obj_at_pos > 0)
+
+        # Handle digestion: add reward to buffer if collected
+        digestion_buffer = state.digestion_buffer
+        key, reward_subkey = jax.random.split(key)
+        object_reward = jax.lax.switch(
+            obj_at_pos, self.reward_fns, state.time, reward_subkey
+        )
+        key, digestion_subkey = jax.random.split(key)
+        reward_delay = jax.lax.switch(
+            obj_at_pos, self.reward_delay_fns, state.time, digestion_subkey
+        )
+        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 & (reward_delay > 0),
+                lambda: digestion_buffer.at[
+                    (state.time + reward_delay) % self.max_reward_delay
+                ].add(object_reward),
+                lambda: digestion_buffer,
+            )
+            # Deliver current rewards
+            current_index = state.time % self.max_reward_delay
+            reward += digestion_buffer[current_index]
+            digestion_buffer = digestion_buffer.at[current_index].set(0.0)
 
         # 3. HANDLE OBJECT COLLECTION AND RESPAWNING
-        key, subkey, rand_key = jax.random.split(key, 3)
+        key, regen_subkey, rand_key = jax.random.split(key, 3)
 
         # Decrement timers (stored as negative values)
         is_timer = state.object_grid < 0
@@ -252,7 +279,7 @@ class ForagaxEnv(environment.Environment):
 
         # Collect object: set a timer
         regen_delay = jax.lax.switch(
-            obj_at_pos, self.regen_delay_fns, state.time, subkey
+            obj_at_pos, self.regen_delay_fns, state.time, regen_subkey
         )
         encoded_timer = obj_at_pos - ((regen_delay + 1) * num_obj_types)
 
@@ -301,10 +328,13 @@ class ForagaxEnv(environment.Environment):
         )
 
         info = {"discount": self.discount(state, params)}
-        if self.weather_object is not None:
-            info["temperature"] = get_temperature(
-                self.weather_object.rewards, state.time, self.weather_object.repeat
-            )
+        temperatures = jnp.zeros(len(self.objects))
+        for obj_index, obj in enumerate(self.objects):
+            if isinstance(obj, WeatherObject):
+                temperatures = temperatures.at[obj_index].set(
+                    get_temperature(obj.rewards, state.time, obj.repeat)
+                )
+        info["temperatures"] = temperatures
         info["biome_id"] = state.biome_grid[pos[1], pos[0]]
         info["object_collected_id"] = jax.lax.select(should_collect, obj_at_pos, -1)
 
@@ -314,6 +344,7 @@ class ForagaxEnv(environment.Environment):
             object_grid=object_grid,
             biome_grid=state.biome_grid,
             time=state.time + 1,
+            digestion_buffer=digestion_buffer,
         )
 
         done = self.is_terminal(state, params)
@@ -352,6 +383,7 @@ class ForagaxEnv(environment.Environment):
             object_grid=object_grid,
             biome_grid=biome_grid,
             time=0,
+            digestion_buffer=jnp.zeros((self.max_reward_delay,)),
         )
 
         return self.get_obs(state, params), state
@@ -412,6 +444,12 @@ class ForagaxEnv(environment.Environment):
                     int,
                 ),
                 "time": spaces.Discrete(params.max_steps_in_episode),
+                "digestion_buffer": spaces.Box(
+                    -jnp.inf,
+                    jnp.inf,
+                    (self.max_reward_delay,),
+                    float,
+                ),
             }
         )
 

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/src/foragax/objects.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Tuple
+from typing import Optional, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -17,12 +17,14 @@ class BaseForagaxObject:
         collectable: bool = False,
         color: Tuple[int, int, int] = (0, 0, 0),
         random_respawn: bool = False,
+        max_reward_delay: int = 0,
     ):
         self.name = name
         self.blocking = blocking
         self.collectable = collectable
         self.color = color
         self.random_respawn = random_respawn
+        self.max_reward_delay = max_reward_delay
 
     @abc.abstractmethod
     def reward(self, clock: int, rng: jax.Array) -> float:
@@ -30,8 +32,8 @@ class BaseForagaxObject:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def regen_delay(self, clock: int, rng: jax.Array) -> int:
-        """Regeneration delay function."""
+    def reward_delay(self, clock: int, rng: jax.Array) -> int:
+        """Reward delay function."""
         raise NotImplementedError
 
 
@@ -47,10 +49,17 @@ class DefaultForagaxObject(BaseForagaxObject):
         regen_delay: Tuple[int, int] = (10, 100),
         color: Tuple[int, int, int] = (255, 255, 255),
         random_respawn: bool = False,
+        reward_delay: int = 0,
+        max_reward_delay: Optional[int] = None,
     ):
-        super().__init__(name, blocking, collectable, color, random_respawn)
+        if max_reward_delay is None:
+            max_reward_delay = reward_delay
+        super().__init__(
+            name, blocking, collectable, color, random_respawn, max_reward_delay
+        )
         self.reward_val = reward
         self.regen_delay_range = regen_delay
+        self.reward_delay_val = reward_delay
 
     def reward(self, clock: int, rng: jax.Array) -> float:
         """Default reward function."""
@@ -61,6 +70,10 @@ class DefaultForagaxObject(BaseForagaxObject):
         min_delay, max_delay = self.regen_delay_range
         return jax.random.randint(rng, (), min_delay, max_delay)
 
+    def reward_delay(self, clock: int, rng: jax.Array) -> int:
+        """Default reward delay function."""
+        return self.reward_delay_val
+
 
 class NormalRegenForagaxObject(DefaultForagaxObject):
     """Object with regeneration delay from a normal distribution."""
@@ -74,6 +87,8 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
         std_regen_delay: int = 1,
         color: Tuple[int, int, int] = (0, 0, 0),
         random_respawn: bool = False,
+        reward_delay: int = 0,
+        max_reward_delay: Optional[int] = None,
     ):
         super().__init__(
             name=name,
@@ -82,6 +97,8 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
             regen_delay=(mean_regen_delay, mean_regen_delay),
             color=color,
             random_respawn=random_respawn,
+            reward_delay=reward_delay,
+            max_reward_delay=max_reward_delay,
         )
         self.mean_regen_delay = mean_regen_delay
         self.std_regen_delay = std_regen_delay
@@ -105,6 +122,8 @@ class WeatherObject(NormalRegenForagaxObject):
         std_regen_delay: int = 1,
         color: Tuple[int, int, int] = (0, 0, 0),
         random_respawn: bool = False,
+        reward_delay: int = 0,
+        max_reward_delay: Optional[int] = None,
     ):
         super().__init__(
             name=name,
@@ -113,6 +132,8 @@ class WeatherObject(NormalRegenForagaxObject):
             std_regen_delay=std_regen_delay,
             color=color,
             random_respawn=random_respawn,
+            reward_delay=reward_delay,
+            max_reward_delay=max_reward_delay,
         )
         self.rewards = rewards
         self.repeat = repeat
@@ -319,6 +340,7 @@ def create_weather_objects(
     multiplier: float = 1.0,
     same_color: bool = False,
     random_respawn: bool = False,
+    reward_delay: int = 0,
 ):
     """Create HOT and COLD WeatherObject instances using the specified file.
 
@@ -348,6 +370,7 @@ def create_weather_objects(
         multiplier=multiplier,
         color=hot_color,
         random_respawn=random_respawn,
+        reward_delay=reward_delay,
     )
 
     cold_color = hot_color if same_color else (0, 255, 255)
@@ -358,6 +381,7 @@ def create_weather_objects(
         multiplier=-multiplier,
         color=cold_color,
         random_respawn=random_respawn,
+        reward_delay=reward_delay,
     )
 
     return hot, cold

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/src/foragax/registry.py

@@ -348,7 +348,8 @@ def make(
     observation_type: str = "color",
     aperture_size: Optional[Tuple[int, int]] = (5, 5),
     file_index: int = 0,
-    nowrap: Optional[bool] = None,
+    repeat: int = 500,
+    reward_delay: int = 0,
     **kwargs: Any,
 ) -> ForagaxEnv:
     """Create a Foragax environment.
@@ -358,9 +359,9 @@ def make(
         observation_type: The type of observation to use. One of "object", "rgb", or "color".
         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. nowrap: If True, disables
-        wrapping around environment boundaries. If None, uses defaults per
-        environment.
+        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:
@@ -376,8 +377,6 @@ def make(
         else:
             aperture_size = (aperture_size, aperture_size)
     config["aperture_size"] = aperture_size
-    if nowrap is not None:
-        config["nowrap"] = nowrap
 
     # Handle special size and biome configurations
     if env_id in (
@@ -460,9 +459,16 @@ def make(
             "ForagaxWeather-v4",
             "ForagaxWeather-v5",
         )
-        random_respawn = env_id in ("ForagaxWeather-v4", "ForagaxWeather-v5")
+        random_respawn = env_id in (
+            "ForagaxWeather-v4",
+            "ForagaxWeather-v5",
+        )
         hot, cold = create_weather_objects(
-            file_index=file_index, same_color=same_color, random_respawn=random_respawn
+            file_index=file_index,
+            repeat=repeat,
+            same_color=same_color,
+            random_respawn=random_respawn,
+            reward_delay=reward_delay,
         )
         config["objects"] = (hot, cold)
 

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/tests/test_foragax.py

@@ -1034,13 +1034,57 @@ def test_info_temperature():
     key, step_key = jax.random.split(key)
     _, state, _, _, info = env.step(step_key, state, Actions.UP, params)
 
-    assert "temperature" in info
-    # Temperature should be based on time=0, so first temperature value
-    assert info["temperature"] == 10.0
+    assert "temperatures" in info
+    # Temperatures should be an array with temperature at index 1 (object ID 1 for the weather object)
+    assert len(info["temperatures"]) == 2  # EMPTY + 1 weather object
+    assert info["temperatures"][0] == 0.0  # EMPTY
+    assert info["temperatures"][1] == 10.0  # weather object at index 1
 
     key, step_key = jax.random.split(key)
     _, state, _, _, info = env.step(step_key, state, Actions.UP, params)
-    assert info["temperature"] == 20.0  # Next temperature value
+    assert info["temperatures"][1] == 20.0  # Next temperature value
+
+
+def test_info_multiple_weather_objects():
+    """Test that info contains temperatures for multiple weather objects."""
+    key = jax.random.key(0)
+    # Create two weather objects
+    hot_obj = WeatherObject(
+        name="hot",
+        rewards=jnp.array([10.0, 20.0]),
+        repeat=1,
+        multiplier=1.0,
+    )
+    cold_obj = WeatherObject(
+        name="cold",
+        rewards=jnp.array([5.0, 15.0]),
+        repeat=1,
+        multiplier=-1.0,
+    )
+
+    env = ForagaxEnv(
+        size=(5, 5),
+        objects=(hot_obj, cold_obj),
+        biomes=(Biome(object_frequencies=(0.1, 0.1)),),
+        observation_type="color",
+    )
+    params = env.default_params
+    obs, state = env.reset(key, params)
+
+    key, step_key = jax.random.split(key)
+    _, state, _, _, info = env.step(step_key, state, Actions.UP, params)
+
+    # Should have temperatures array
+    assert "temperatures" in info
+    assert len(info["temperatures"]) == 3  # EMPTY + 2 objects
+    assert info["temperatures"][0] == 0.0  # EMPTY
+    assert info["temperatures"][1] == 10.0  # hot object at index 1
+    assert info["temperatures"][2] == 5.0  # cold object at index 2 (raw temperature)
+
+    key, step_key = jax.random.split(key)
+    _, state, _, _, info = env.step(step_key, state, Actions.UP, params)
+    assert info["temperatures"][1] == 20.0  # hot next
+    assert info["temperatures"][2] == 15.0  # cold next
 
 
 def test_info_biome_id():
@@ -1108,3 +1152,54 @@ def test_info_object_collected_id():
 
     assert info["object_collected_id"] == -1  # No object collected
     assert reward == 0.0
+
+
+def test_reward_delay():
+    """Test that rewards are delayed by reward_delay for objects with k > 0."""
+    key = jax.random.key(0)
+
+    # Create an object with digestion delay of 2 steps
+    delayed_flower = DefaultForagaxObject(
+        name="delayed_flower",
+        reward=5.0,
+        collectable=True,
+        color=(0, 255, 0),
+        reward_delay=2,
+    )
+
+    env = ForagaxEnv(
+        size=(7, 7),
+        objects=(delayed_flower,),
+        observation_type="color",
+    )
+    params = env.default_params
+    _, state = env.reset(key, params)
+
+    delayed_flower_id = 1  # 0 is EMPTY
+
+    # Place the delayed flower and move the agent to it
+    grid = jnp.zeros((7, 7), dtype=int)
+    grid = grid.at[4, 3].set(delayed_flower_id)
+    state = state.replace(object_grid=grid, pos=jnp.array([3, 3]))
+
+    # Collect the flower by moving down - should get no immediate reward
+    key, step_key = jax.random.split(key)
+    _, state, reward, _, info = env.step(step_key, state, Actions.DOWN, params)
+
+    assert info["object_collected_id"] == delayed_flower_id
+    assert reward == 0.0  # No immediate reward due to digestion delay
+
+    # Step 1: Still no reward (time=1, reward should arrive at time=2)
+    key, step_key = jax.random.split(key)
+    _, state, reward, _, _ = env.step(step_key, state, Actions.UP, params)
+    assert reward == 0.0
+
+    # Step 2: Reward should arrive now (time=2)
+    key, step_key = jax.random.split(key)
+    _, state, reward, _, _ = env.step(step_key, state, Actions.UP, params)
+    assert reward == 5.0
+
+    # Step 3: No more rewards
+    key, step_key = jax.random.split(key)
+    _, state, reward, _, _ = env.step(step_key, state, Actions.UP, params)
+    assert reward == 0.0

{continual_foragax-0.28.1 → continual_foragax-0.30.0}/tests/test_registry.py

@@ -379,3 +379,75 @@ def test_foragax_twobiome_v13_random_respawn():
         "Deathcap objects should have random_respawn=True"
     )
     assert fake.random_respawn is True, "Fake objects should have random_respawn=True"
+
+
+def test_repeat_parameter_weather_environments():
+    """Test that the repeat parameter controls temperature cycling speed."""
+    # Create environments with different repeat values
+    env_repeat_100 = make("ForagaxWeather-v1", repeat=100, aperture_size=(5, 5))
+    env_repeat_200 = make("ForagaxWeather-v1", repeat=200, aperture_size=(5, 5))
+
+    # Get the weather objects
+    hot_100, cold_100 = env_repeat_100.objects[1], env_repeat_100.objects[2]
+    hot_200, cold_200 = env_repeat_200.objects[1], env_repeat_200.objects[2]
+
+    # Check that repeat values are set correctly
+    assert hot_100.repeat == 100
+    assert cold_100.repeat == 100
+    assert hot_200.repeat == 200
+    assert cold_200.repeat == 200
+
+    # Test temperature cycling by checking rewards at different clock times
+    key = jax.random.key(0)
+
+    # At clock=0, both should give the same temperature (first value)
+    temp_100_t0 = hot_100.reward(0, key)
+    temp_200_t0 = hot_200.reward(0, key)
+    assert temp_100_t0 == temp_200_t0, "Temperatures should be identical at clock=0"
+
+    # At clock=100, repeat_100 should move to next temperature (index 1), repeat_200 should stay at first (index 0)
+    temp_100_t100 = hot_100.reward(100, key)
+    temp_200_t100 = hot_200.reward(100, key)
+    assert temp_100_t0 != temp_100_t100, "repeat_100 should cycle at step 100"
+    assert temp_200_t0 == temp_200_t100, "repeat_200 should not cycle at step 100"
+
+    # At clock=200, repeat_100 should be at index 2, repeat_200 should be at index 1
+    temp_100_t200 = hot_100.reward(200, key)
+    temp_200_t200 = hot_200.reward(200, key)
+    assert temp_100_t200 != temp_200_t200, (
+        "Different repeat values should give different temperatures at step 200"
+    )
+
+    # At clock=400, repeat_100 should be at index 4, repeat_200 should be at index 2
+    temp_100_t400 = hot_100.reward(400, key)
+    temp_200_t400 = hot_200.reward(400, key)
+    assert temp_100_t400 != temp_100_t200, "repeat_100 should cycle again at step 400"
+    assert temp_200_t400 != temp_200_t200, "repeat_200 should cycle again at step 400"
+    assert temp_100_t400 != temp_200_t400, (
+        "Different repeat values should give different temperatures at step 400"
+    )
+
+
+def test_reward_delay_parameter_weather_environments():
+    """Test that the reward_delay parameter controls digestion delay."""
+    # Create environments with different reward_delay values
+    env_delays_0 = make("ForagaxWeather-v1", reward_delay=0, aperture_size=(5, 5))
+    env_delays_5 = make("ForagaxWeather-v1", reward_delay=5, aperture_size=(5, 5))
+
+    # Get the weather objects
+    hot_0, cold_0 = env_delays_0.objects[1], env_delays_0.objects[2]
+    hot_5, cold_5 = env_delays_5.objects[1], env_delays_5.objects[2]
+
+    # Check that reward_delay values are set correctly
+    assert hot_0.reward_delay_val == 0
+    assert cold_0.reward_delay_val == 0
+    assert hot_5.reward_delay_val == 5
+    assert cold_5.reward_delay_val == 5
+
+    # Test reward_delay function returns the correct values
+    key = jax.random.key(0)
+    delays_0 = hot_0.reward_delay(0, key)
+    delays_5 = hot_5.reward_delay(0, key)
+
+    assert delays_0 == 0, "reward_delay should return 0 for delays=0"
+    assert delays_5 == 5, "reward_delay should return 5 for delays=5"