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

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

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

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

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

{continual_foragax-0.29.0 → 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.29.0
+Version: 0.30.0
 Summary: A continual reinforcement learning benchmark
 Author-email: Steven Tang <stang5@ualberta.ca>
 Requires-Python: >=3.8

{continual_foragax-0.29.0 → continual_foragax-0.30.0}/src/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,
                 ),
             }

{continual_foragax-0.29.0 → continual_foragax-0.30.0}/src/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,8 +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,
+            reward_delay=reward_delay,
+            max_reward_delay=max_reward_delay,
         )
         self.rewards = rewards
         self.repeat = repeat
@@ -340,7 +340,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 +370,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 +381,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

{continual_foragax-0.29.0 → continual_foragax-0.30.0}/src/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)
 

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

@@ -1154,8 +1154,8 @@ def test_info_object_collected_id():
     assert reward == 0.0
 
 
-def test_digestion_delay():
-    """Test that rewards are delayed by digestion_steps for objects with k > 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
@@ -1164,7 +1164,7 @@ def test_digestion_delay():
         reward=5.0,
         collectable=True,
         color=(0, 255, 0),
-        digestion_steps=2,
+        reward_delay=2,
     )
 
     env = ForagaxEnv(

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

@@ -181,107 +181,6 @@ 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))
@@ -480,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"