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

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

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

{continual_foragax-0.33.0.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=kf0B0D5CMUnOMiZ8diBlmlt_vf2Wh-M4gSNoeo4jfHY,55147
+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.33.0.dist-info/METADATA,sha256=vEOZLNVNPhccIZDIrN-puYjVPQKWHtThbhoVdUjhF4A,4713
-continual_foragax-0.33.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.33.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.33.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.33.0.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

@@ -553,6 +553,15 @@ 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)
 
+        # 4. UPDATE STATE
+        state = EnvState(
+            pos=pos,
+            time=state.time + 1,
+            digestion_buffer=digestion_buffer,
+            object_state=object_state,
+            biome_state=biome_state,
+        )
+
         # Compute reward at each grid position
         fixed_key = jax.random.key(0)  # Fixed key for deterministic reward computation
 
@@ -570,15 +579,6 @@ class ForagaxEnv(environment.Environment):
         )
         info["rewards"] = reward_grid
 
-        # 4. UPDATE STATE
-        state = EnvState(
-            pos=pos,
-            time=state.time + 1,
-            digestion_buffer=digestion_buffer,
-            object_state=object_state,
-            biome_state=biome_state,
-        )
-
         done = self.is_terminal(state, params)
         return (
             jax.lax.stop_gradient(self.get_obs(state, params)),