continual-foragax 0.10.0__py3-none-any.whl → 0.10.1__py3-none-any.whl

{continual_foragax-0.10.0.dist-info → continual_foragax-0.10.1.dist-info}/METADATA

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

{continual_foragax-0.10.0.dist-info → continual_foragax-0.10.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=OtpcyqzBOQLdTvvRegD3SYm4mi4Ga2WE5eJ7OQmQOaw,18294
+foragax/env.py,sha256=X9oc60xNL4uTFnbt_BynN_c3XmVa9MYbSclW-g4qQoc,18628
 foragax/objects.py,sha256=CyBxrykTxpHCI_2hE9jE8mG4TU8R7VxzKdQ5mtxkEqU,6004
 foragax/registry.py,sha256=7_RDXvm_3RNO7culBLGkE0jH8Wk_q6jbMv72dZx4JO8,2722
 foragax/rendering.py,sha256=KAoQpdndy5JDQlyG0c5QDHuH-_Tfy5RuVlDtndnHVjc,2765
@@ -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.10.0.dist-info/METADATA,sha256=Oo3pJnjoU7VLruPXrhNT-WfzF7cGvkGN8crW-BclMsk,4897
-continual_foragax-0.10.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.10.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.10.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.10.0.dist-info/RECORD,,
+continual_foragax-0.10.1.dist-info/METADATA,sha256=tgSSa5FtfA6qRBCZ_L3r3XnsfDpjhKMTRFJ9TL_OfCo,4897
+continual_foragax-0.10.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.10.1.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.10.1.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.10.1.dist-info/RECORD,,

foragax/env.py

@@ -492,7 +492,14 @@ class ForagaxWorldEnv(ForagaxEnv):
         num_obj_types = len(self.object_ids)
         # Decode grid for observation
         obs_grid = jnp.maximum(0, state.object_grid)
-        obs = jax.nn.one_hot(obs_grid, num_obj_types)
+        obs = jnp.zeros((self.size[1], self.size[0], num_obj_types), dtype=jnp.float32)
+
+        num_object_channels = num_obj_types - 1
+        # create masks for all objects at once
+        object_ids = jnp.arange(1, num_obj_types)
+        object_masks = obs_grid[..., None] == object_ids[None, None, :]
+        obs = obs.at[:, :, :num_object_channels].set(object_masks.astype(float))
+
         obs = obs.at[state.pos[1], state.pos[0], -1].set(1)
         return obs