continual-foragax 0.32.1__py3-none-any.whl → 0.33.1__py3-none-any.whl

{continual_foragax-0.32.1.dist-info → continual_foragax-0.33.1.dist-info}/METADATA

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

{continual_foragax-0.32.1.dist-info → continual_foragax-0.33.1.dist-info}/RECORD

@@ -1,6 +1,6 @@
 foragax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 foragax/colors.py,sha256=rqNPiywP4Nvr0POhsGpasRk-nMMTS3DOwFRUgperlUk,2065
-foragax/env.py,sha256=Jjo7XypfQf6ePoKYTV3xvolK4qpacuaifWLZB0ke5y8,54559
+foragax/env.py,sha256=bDhNQTaqcoOwm9Csb1LHoduuNdE1j1RAhGnVV7cAEPI,55147
 foragax/objects.py,sha256=9wv0ZKT89dDkaeVwUwkVo4dwhRVeUxvsTyhoyYKfOEw,26508
 foragax/registry.py,sha256=hfzQHNgX6uoOdbf4_21iH25abQVQZIjBWn7h5bdrSBg,17981
 foragax/rendering.py,sha256=bms7wvBZTofoR-K-2QD2Ggeed7Viw8uwAEiEpEM3eSo,2768
@@ -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.32.1.dist-info/METADATA,sha256=ZG39JPQKbUW7ag-vTZtcDfL8Wvt-nCfO-KOCOZMgOIo,4713
-continual_foragax-0.32.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.32.1.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.32.1.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.32.1.dist-info/RECORD,,
+continual_foragax-0.33.1.dist-info/METADATA,sha256=n4TeSbm2f0oBPHWnRPozN-Y76qn3RlF9ey2QI8jtVU0,4713
+continual_foragax-0.33.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.33.1.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.33.1.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.33.1.dist-info/RECORD,,

foragax/env.py

@@ -562,6 +562,23 @@ class ForagaxEnv(environment.Environment):
             biome_state=biome_state,
         )
 
+        # 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
+
         done = self.is_terminal(state, params)
         return (
             jax.lax.stop_gradient(self.get_obs(state, params)),