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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: continual-foragax
-Version: 0.28.1
+Version: 0.29.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.29.0}/pyproject.toml

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

{continual_foragax-0.28.1 → continual_foragax-0.29.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.29.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.29.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.digestion_steps_fns = [o.digestion_steps 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
+        )
 
         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)
+        digestion_steps = jax.lax.switch(
+            obj_at_pos, self.digestion_steps_fns, state.time, digestion_subkey
+        )
+        reward = jnp.where(should_collect & (digestion_steps == 0), object_reward, 0.0)
+        if self.max_digestion_steps > 0:
+            # Add delayed rewards to buffer
+            digestion_buffer = jax.lax.cond(
+                should_collect & (digestion_steps > 0),
+                lambda: digestion_buffer.at[
+                    (state.time + digestion_steps) % self.max_digestion_steps
+                ].add(object_reward),
+                lambda: digestion_buffer,
+            )
+            # Deliver current rewards
+            current_index = state.time % self.max_digestion_steps
+            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_digestion_steps,)),
         )
 
         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_digestion_steps,),
+                    float,
+                ),
             }
         )
 

{continual_foragax-0.28.1 → continual_foragax-0.29.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_digestion_steps: 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
 
     @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 digestion_steps(self, clock: int, rng: jax.Array) -> int:
+        """Digestion steps 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,
+        digestion_steps: int = 0,
+        max_digestion_steps: Optional[int] = None,
     ):
-        super().__init__(name, blocking, collectable, color, random_respawn)
+        if max_digestion_steps is None:
+            max_digestion_steps = digestion_steps
+        super().__init__(
+            name, blocking, collectable, color, random_respawn, max_digestion_steps
+        )
         self.reward_val = reward
         self.regen_delay_range = regen_delay
+        self.digestion_steps_val = digestion_steps
 
     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 digestion_steps(self, clock: int, rng: jax.Array) -> int:
+        """Default digestion steps function."""
+        return self.digestion_steps_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,
+        digestion_steps: int = 0,
+        max_digestion_steps: 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,
+            digestion_steps=digestion_steps,
+            max_digestion_steps=max_digestion_steps,
         )
         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,
+        digestion_steps: int = 0,
+        max_digestion_steps: 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,
+            digestion_steps=digestion_steps,
+            max_digestion_steps=max_digestion_steps,
         )
         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,
+    digestion_steps: 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,
+        digestion_steps=digestion_steps,
     )
 
     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,
+        digestion_steps=digestion_steps,
     )
 
     return hot, cold

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

@@ -83,6 +83,19 @@ 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,
@@ -348,7 +361,6 @@ def make(
     observation_type: str = "color",
     aperture_size: Optional[Tuple[int, int]] = (5, 5),
     file_index: int = 0,
-    nowrap: Optional[bool] = None,
     **kwargs: Any,
 ) -> ForagaxEnv:
     """Create a Foragax environment.
@@ -358,9 +370,7 @@ 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.
         **kwargs: Additional keyword arguments to pass to the ForagaxEnv constructor.
 
     Returns:
@@ -376,8 +386,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 (
@@ -459,10 +467,19 @@ def make(
             "ForagaxWeather-v3",
             "ForagaxWeather-v4",
             "ForagaxWeather-v5",
+            "ForagaxWeather-v6",
+        )
+        random_respawn = env_id in (
+            "ForagaxWeather-v4",
+            "ForagaxWeather-v5",
+            "ForagaxWeather-v6",
         )
-        random_respawn = env_id in ("ForagaxWeather-v4", "ForagaxWeather-v5")
+        digestion_steps = 10 if env_id in ("ForagaxWeather-v6") else 0
         hot, cold = create_weather_objects(
-            file_index=file_index, same_color=same_color, random_respawn=random_respawn
+            file_index=file_index,
+            same_color=same_color,
+            random_respawn=random_respawn,
+            digestion_steps=digestion_steps,
         )
         config["objects"] = (hot, cold)
 

{continual_foragax-0.28.1 → continual_foragax-0.29.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_digestion_delay():
+    """Test that rewards are delayed by digestion_steps 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),
+        digestion_steps=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.29.0}/tests/test_registry.py

@@ -181,6 +181,107 @@ def test_foragax_weather_v5_color_configuration():
     assert hot_v5.color == cold_v5.color, "v5 should use same color for hot and cold"
 
 
+def test_foragax_weather_v6_registry():
+    """Test that ForagaxWeather-v6 can be created via registry and has correct config."""
+    env = make("ForagaxWeather-v6", aperture_size=(5, 5))
+
+    # Check basic configuration
+    assert env.name == "ForagaxWeather-v6"
+    assert env.deterministic_spawn is True
+    assert env.nowrap is False  # v6 enables wrapping
+
+    # Check that weather objects have random_respawn=True
+    hot, cold = env.objects[1], env.objects[2]  # Skip EMPTY (index 0)
+    assert hot.name == "hot"
+    assert cold.name == "cold"
+    assert hot.random_respawn is True
+    assert cold.random_respawn is True
+
+    # Check digestion steps
+    assert hot.max_digestion_steps == 10  # digestion_steps=10
+    assert cold.max_digestion_steps == 10  # digestion_steps=10
+
+    # Test basic functionality
+    key = jax.random.key(0)
+    obs, state = env.reset(key, env.default_params)
+    assert obs.shape == (5, 5, 1)  # 1 color channel (hot/cold same color, no padding)
+
+    # Test stepping
+    key, step_key = jax.random.split(key)
+    action = env.action_space(env.default_params).sample(step_key)
+    obs2, state2, reward, done, info = env.step(
+        step_key, state, action, env.default_params
+    )
+    assert obs2.shape == (5, 5, 1)
+    assert not done
+
+
+def test_foragax_weather_v6_deterministic_spawn():
+    """Test that ForagaxWeather-v6 uses deterministic spawning."""
+    env = make("ForagaxWeather-v6", aperture_size=(5, 5))
+    params = env.default_params
+
+    # Test that multiple resets with same key produce same object placement
+    key = jax.random.key(42)
+    _, state1 = env.reset(key, params)
+
+    key = jax.random.key(42)  # Same key
+    _, state2 = env.reset(key, params)
+
+    # Object grids should be identical (deterministic spawn)
+    chex.assert_trees_all_equal(state1.object_grid, state2.object_grid)
+
+    # But different keys should produce different placements
+    key1 = jax.random.key(42)
+    key2 = jax.random.key(43)
+    _, state1 = env.reset(key1, params)
+    _, state2 = env.reset(key2, params)
+
+    # Should be different (shuffled deterministically)
+    assert not jnp.array_equal(state1.object_grid, state2.object_grid)
+
+    # Test that number of objects is the same
+    num_hot_1 = jnp.sum(state1.object_grid == 1)
+    num_cold_1 = jnp.sum(state1.object_grid == 2)
+    num_hot_2 = jnp.sum(state2.object_grid == 1)
+    num_cold_2 = jnp.sum(state2.object_grid == 2)
+    assert num_hot_1 == num_hot_2
+    assert num_cold_1 == num_cold_2
+
+
+def test_foragax_weather_v6_random_respawn():
+    """Test that ForagaxWeather-v6 weather objects have random_respawn=True."""
+    env = make("ForagaxWeather-v6", aperture_size=(5, 5))
+
+    # Check that weather objects have random_respawn=True
+    hot, cold = env.objects[1], env.objects[2]  # Skip EMPTY
+    assert hot.random_respawn is True, "Hot objects should have random_respawn=True"
+    assert cold.random_respawn is True, "Cold objects should have random_respawn=True"
+
+
+def test_foragax_weather_v6_color_configuration():
+    """Test that ForagaxWeather-v6 weather objects use the same color."""
+    env = make("ForagaxWeather-v6", aperture_size=(5, 5))
+    hot_v6, cold_v6 = env.objects[1], env.objects[2]
+
+    # Same color configuration: v6 uses same color
+    assert hot_v6.color == cold_v6.color, "v6 should use same color for hot and cold"
+
+
+def test_foragax_weather_v6_digestion_steps():
+    """Test that ForagaxWeather-v6 weather objects have digestion_steps=10."""
+    env = make("ForagaxWeather-v6", aperture_size=(5, 5))
+
+    # Check that weather objects have digestion_steps=10
+    hot, cold = env.objects[1], env.objects[2]  # Skip EMPTY
+    assert hot.digestion_steps(0, jax.random.key(0)) == 10, (
+        "Hot objects should have digestion_steps=10"
+    )
+    assert cold.digestion_steps(0, jax.random.key(0)) == 10, (
+        "Cold objects should have digestion_steps=10"
+    )
+
+
 def test_foragax_twobiome_v10_registry():
     """Test that ForagaxTwoBiome-v10 can be created via registry and has correct config."""
     env = make("ForagaxTwoBiome-v10", aperture_size=(5, 5))