continual-foragax 0.32.1__tar.gz → 0.33.0__tar.gz

{continual_foragax-0.32.1 → continual_foragax-0.33.0}/PKG-INFO

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

{continual_foragax-0.32.1 → continual_foragax-0.33.0}/pyproject.toml

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

{continual_foragax-0.32.1 → continual_foragax-0.33.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.32.1 → continual_foragax-0.33.0}/src/foragax/env.py

@@ -553,6 +553,23 @@ class ForagaxEnv(environment.Environment):
         info["biome_id"] = object_state.biome_id[pos[1], pos[0]]
         info["object_collected_id"] = jax.lax.select(should_collect, obj_at_pos, -1)
 
+        # Compute reward at each grid position
+        fixed_key = jax.random.key(0)  # Fixed key for deterministic reward computation
+
+        def compute_reward(obj_id, params):
+            return jax.lax.cond(
+                obj_id > 0,
+                lambda: jax.lax.switch(
+                    obj_id, self.reward_fns, state.time, fixed_key, params
+                ),
+                lambda: 0.0,
+            )
+
+        reward_grid = jax.vmap(jax.vmap(compute_reward))(
+            object_state.object_id, object_state.state_params
+        )
+        info["rewards"] = reward_grid
+
         # 4. UPDATE STATE
         state = EnvState(
             pos=pos,

{continual_foragax-0.32.1 → continual_foragax-0.33.0}/tests/test_foragax.py

@@ -2797,3 +2797,35 @@ def test_sine_object():
         r_orig = sine_obj.reward(t, key, None)
         r_shift = sine_obj_shifted.reward(t, key, None)
         assert jnp.allclose(r_orig + r_shift, 0.0, atol=0.01)
+
+
+def test_info_rewards():
+    """Test that info contains rewards with reward values for each grid position."""
+    key = jax.random.key(0)
+    env = ForagaxEnv(
+        size=(5, 5),
+        objects=(FLOWER,),
+        biomes=(Biome(object_frequencies=(0.5,)),),
+        observation_type="object",
+    )
+    params = env.default_params
+    obs, state = env.reset(key, params)
+
+    key, step_key = jax.random.split(key)
+    _, _, _, _, info = env.step(step_key, state, Actions.UP, params)
+
+    # Check that rewards is in info
+    assert "rewards" in info
+
+    # Check shape matches environment size
+    assert info["rewards"].shape == (5, 5)
+
+    # Check that positions with objects have non-zero rewards
+    object_mask = state.object_state.object_id > 0
+    rewards_at_objects = info["rewards"][object_mask]
+    assert jnp.all(rewards_at_objects == FLOWER.reward_val)
+
+    # Check that empty positions have zero rewards
+    empty_mask = state.object_state.object_id == 0
+    rewards_at_empty = info["rewards"][empty_mask]
+    assert jnp.all(rewards_at_empty == 0.0)