continual-foragax 0.10.0__tar.gz → 0.10.1__tar.gz

{continual_foragax-0.10.0 → continual_foragax-0.10.1}/PKG-INFO

@@ -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 → continual_foragax-0.10.1}/pyproject.toml

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

{continual_foragax-0.10.0 → continual_foragax-0.10.1}/src/continual_foragax.egg-info/PKG-INFO

@@ -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 → continual_foragax-0.10.1}/src/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
 

{continual_foragax-0.10.0 → continual_foragax-0.10.1}/tests/test_foragax.py

@@ -143,12 +143,25 @@ def test_rgb_observation_mode_large_aperture():
 
 def test_world_observation_mode():
     # can use world observation mode
-    env = ForagaxWorldEnv(size=(10, 10))
+    env = ForagaxWorldEnv(
+        size=(10, 10),
+        aperture_size=5,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.1, 0.1)),),
+    )
     params = env.default_params
     key = jax.random.key(0)
     obs, state = env.reset(key, params)
 
-    assert obs.shape == (10, 10, 1)
+    assert obs.shape == (10, 10, 3)
+    # Check agent channel (last channel) is 0 except 1 at agent position
+    agent_channel = obs[:, :, 2]
+    assert jnp.sum(agent_channel) == 1
+    assert agent_channel[state.pos[1], state.pos[0]] == 1
+
+    # check non-zero values in other channels
+    assert jnp.sum(obs[:, :, 0]) > 0
+    assert jnp.sum(obs[:, :, 1]) > 0
 
 
 def test_basic_movement():