continual-foragax 0.20.1__tar.gz → 0.22.0__tar.gz

{continual_foragax-0.20.1 → continual_foragax-0.22.0}/PKG-INFO

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

{continual_foragax-0.20.1 → continual_foragax-0.22.0}/pyproject.toml

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

{continual_foragax-0.20.1 → continual_foragax-0.22.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.20.1 → continual_foragax-0.22.0}/src/continual_foragax.egg-info/SOURCES.txt

@@ -136,4 +136,5 @@ src/foragax/data/ECA_non-blended_custom/TG_SOUID156887.txt
 src/foragax/data/ECA_non-blended_custom/elements.txt
 src/foragax/data/ECA_non-blended_custom/metadata.txt
 src/foragax/data/ECA_non-blended_custom/sources.txt
+tests/test_benchmark.py
 tests/test_foragax.py

{continual_foragax-0.20.1 → continual_foragax-0.22.0}/src/foragax/env.py

@@ -10,6 +10,7 @@ from typing import Any, Dict, Tuple, Union
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 from flax import struct
 from gymnax.environments import environment, spaces
 
@@ -66,13 +67,16 @@ class ForagaxEnv(environment.Environment):
 
     def __init__(
         self,
+        name: str = "Foragax-v0",
         size: Union[Tuple[int, int], int] = (10, 10),
         aperture_size: Union[Tuple[int, int], int] = (5, 5),
         objects: Tuple[BaseForagaxObject, ...] = (),
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
+        deterministic_spawn: bool = False,
     ):
         super().__init__()
+        self._name = name
         if isinstance(size, int):
             size = (size, size)
         self.size = size
@@ -81,6 +85,7 @@ class ForagaxEnv(environment.Environment):
             aperture_size = (aperture_size, aperture_size)
         self.aperture_size = aperture_size
         self.nowrap = nowrap
+        self.deterministic_spawn = deterministic_spawn
         objects = (EMPTY,) + objects
         if self.nowrap:
             objects = objects + (PADDING,)
@@ -103,12 +108,35 @@ class ForagaxEnv(environment.Environment):
         self.biome_object_frequencies = jnp.array(
             [b.object_frequencies for b in biomes]
         )
-        self.biome_starts = jnp.array(
+        self.biome_starts = np.array(
             [b.start if b.start is not None else (-1, -1) for b in biomes]
         )
-        self.biome_stops = jnp.array(
+        self.biome_stops = np.array(
             [b.stop if b.stop is not None else (-1, -1) for b in biomes]
         )
+        self.biome_sizes = np.prod(self.biome_stops - self.biome_starts, axis=1)
+        self.biome_masks = []
+        for i in range(self.biome_object_frequencies.shape[0]):
+            # Create mask for the biome
+            start = jax.lax.select(
+                self.biome_starts[i, 0] == -1,
+                jnp.array([0, 0]),
+                self.biome_starts[i],
+            )
+            stop = jax.lax.select(
+                self.biome_stops[i, 0] == -1,
+                jnp.array(self.size),
+                self.biome_stops[i],
+            )
+            rows = jnp.arange(self.size[1])[:, None]
+            cols = jnp.arange(self.size[0])
+            mask = (
+                (rows >= start[1])
+                & (rows < stop[1])
+                & (cols >= start[0])
+                & (cols < stop[0])
+            )
+            self.biome_masks.append(mask)
 
     @property
     def default_params(self) -> EnvParams:
@@ -196,57 +224,18 @@ class ForagaxEnv(environment.Environment):
         self, key: jax.Array, params: EnvParams
     ) -> Tuple[jax.Array, EnvState]:
         """Reset environment state."""
-        key, subkey = jax.random.split(key)
-
         object_grid = jnp.zeros((self.size[1], self.size[0]), dtype=int)
-
-        iter_key = subkey
+        key, iter_key = jax.random.split(key)
         for i in range(self.biome_object_frequencies.shape[0]):
             iter_key, biome_key = jax.random.split(iter_key)
-            # Generate random layout
-            grid_rand = jax.random.uniform(biome_key, (self.size[1], self.size[0]))
-
-            # Create mask for the biome
-            start = jax.lax.select(
-                self.biome_starts[i, 0] == -1,
-                jnp.array([0, 0]),
-                self.biome_starts[i],
-            )
-            stop = jax.lax.select(
-                self.biome_stops[i, 0] == -1,
-                jnp.array(self.size),
-                self.biome_stops[i],
-            )
-
-            rows = jnp.arange(self.size[1])[:, None]
-            cols = jnp.arange(self.size[0])
+            mask = self.biome_masks[i]
 
-            mask = (
-                (rows >= start[1])
-                & (rows < stop[1])
-                & (cols >= start[0])
-                & (cols < stop[0])
-            )
-
-            # Generate objects for this biome and update the main grid
-            biome_freqs = self.biome_object_frequencies[i]
-            empty_freq = 1.0 - jnp.sum(biome_freqs)
-            all_freqs = jnp.concatenate([jnp.array([empty_freq]), biome_freqs])
-
-            cumulative_freqs = jnp.cumsum(
-                jnp.concatenate([jnp.array([0.0]), all_freqs])
-            )
-
-            # Determine which object to place in each cell
-            # The last object ID will be used for any value of grid_rand >= cumulative_freqs[-1]
-            # so we don't need to cap grid_rand
-            obj_ids_for_biome = jnp.arange(len(all_freqs))
-            cell_obj_ids = (
-                jnp.searchsorted(cumulative_freqs, grid_rand, side="right") - 1
-            )
-            biome_objects = obj_ids_for_biome[cell_obj_ids]
-
-            object_grid = jnp.where(mask, biome_objects, object_grid)
+            if self.deterministic_spawn:
+                biome_objects = self.generate_biome_new(i, biome_key)
+                object_grid = object_grid.at[mask].set(biome_objects)
+            else:
+                biome_objects = self.generate_biome_old(i, biome_key)
+                object_grid = jnp.where(mask, biome_objects, object_grid)
 
         # Place agent in the center of the world and ensure the cell is empty.
         agent_pos = jnp.array([self.size[0] // 2, self.size[1] // 2])
@@ -260,6 +249,25 @@ class ForagaxEnv(environment.Environment):
 
         return self.get_obs(state, params), state
 
+    def generate_biome_old(self, i: int, biome_key: jax.Array):
+        biome_freqs = self.biome_object_frequencies[i]
+        grid_rand = jax.random.uniform(biome_key, (self.size[1], self.size[0]))
+        empty_freq = 1.0 - jnp.sum(biome_freqs)
+        all_freqs = jnp.concatenate([jnp.array([empty_freq]), biome_freqs])
+        cumulative_freqs = jnp.cumsum(jnp.concatenate([jnp.array([0.0]), all_freqs]))
+        biome_objects = jnp.searchsorted(cumulative_freqs, grid_rand, side="right") - 1
+        return biome_objects
+
+    def generate_biome_new(self, i: int, biome_key: jax.Array):
+        biome_freqs = self.biome_object_frequencies[i]
+        grid = jnp.linspace(0, 1, self.biome_sizes[i], endpoint=False)
+        biome_objects = len(biome_freqs) - jnp.searchsorted(
+            jnp.cumsum(biome_freqs[::-1]), grid, side="right"
+        )
+        flat_biome_objects = biome_objects.flatten()
+        shuffled_objects = jax.random.permutation(biome_key, flat_biome_objects)
+        return shuffled_objects
+
     def is_terminal(self, state: EnvState, params: EnvParams) -> jax.Array:
         """Foragax is a continuing environment."""
         return False
@@ -267,7 +275,7 @@ class ForagaxEnv(environment.Environment):
     @property
     def name(self) -> str:
         """Environment name."""
-        return "Foragax-v0"
+        return self._name
 
     @property
     def num_actions(self) -> int:
@@ -438,13 +446,17 @@ class ForagaxObjectEnv(ForagaxEnv):
 
     def __init__(
         self,
+        name: str = "Foragax-v0",
         size: Union[Tuple[int, int], int] = (10, 10),
         aperture_size: Union[Tuple[int, int], int] = (5, 5),
         objects: Tuple[BaseForagaxObject, ...] = (),
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
+        deterministic_spawn: bool = False,
     ):
-        super().__init__(size, aperture_size, objects, biomes, nowrap)
+        super().__init__(
+            name, size, aperture_size, objects, biomes, nowrap, deterministic_spawn
+        )
 
         # Compute unique colors and mapping for partial observability
         # Exclude EMPTY (index 0) from color channels

{continual_foragax-0.20.1 → continual_foragax-0.22.0}/src/foragax/objects.py

@@ -242,6 +242,34 @@ GREEN_FAKE_2 = NormalRegenForagaxObject(
     mean_regen_delay=10,
     std_regen_delay=1,
 )
+BROWN_MOREL_UNIFORM = DefaultForagaxObject(
+    name="brown_morel",
+    reward=10.0,
+    collectable=True,
+    color=(63, 30, 25),
+    regen_delay=(90, 110),
+)
+BROWN_OYSTER_UNIFORM = DefaultForagaxObject(
+    name="brown_oyster",
+    reward=1.0,
+    collectable=True,
+    color=(63, 30, 25),
+    regen_delay=(9, 11),
+)
+GREEN_DEATHCAP_UNIFORM = DefaultForagaxObject(
+    name="green_deathcap",
+    reward=-5.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(9, 11),
+)
+GREEN_FAKE_UNIFORM = DefaultForagaxObject(
+    name="green_fake",
+    reward=0.0,
+    collectable=True,
+    color=(0, 255, 0),
+    regen_delay=(9, 11),
+)
 
 
 def create_weather_objects(

{continual_foragax-0.20.1 → continual_foragax-0.22.0}/src/foragax/registry.py

@@ -12,12 +12,16 @@ from foragax.env import (
 from foragax.objects import (
     BROWN_MOREL,
     BROWN_MOREL_2,
+    BROWN_MOREL_UNIFORM,
     BROWN_OYSTER,
+    BROWN_OYSTER_UNIFORM,
     GREEN_DEATHCAP,
     GREEN_DEATHCAP_2,
     GREEN_DEATHCAP_3,
+    GREEN_DEATHCAP_UNIFORM,
     GREEN_FAKE,
     GREEN_FAKE_2,
+    GREEN_FAKE_UNIFORM,
     LARGE_MOREL,
     LARGE_OYSTER,
     MEDIUM_MOREL,
@@ -147,6 +151,31 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
         "biomes": None,
         "nowrap": True,
     },
+    "ForagaxTwoBiome-v8": {
+        "size": None,
+        "aperture_size": None,
+        "objects": (
+            BROWN_MOREL_UNIFORM,
+            BROWN_OYSTER_UNIFORM,
+            GREEN_DEATHCAP_UNIFORM,
+            GREEN_FAKE_UNIFORM,
+        ),
+        "biomes": None,
+        "nowrap": True,
+    },
+    "ForagaxTwoBiome-v9": {
+        "size": None,
+        "aperture_size": None,
+        "objects": (
+            BROWN_MOREL_UNIFORM,
+            BROWN_OYSTER_UNIFORM,
+            GREEN_DEATHCAP_UNIFORM,
+            GREEN_FAKE_UNIFORM,
+        ),
+        "biomes": None,
+        "nowrap": True,
+        "deterministic_spawn": True,
+    },
     "ForagaxTwoBiomeSmall-v1": {
         "size": (16, 8),
         "aperture_size": None,
@@ -217,7 +246,7 @@ def make(
     if nowrap is not None:
         config["nowrap"] = nowrap
 
-    if env_id == "ForagaxTwoBiome-v7":
+    if env_id in ("ForagaxTwoBiome-v7", "ForagaxTwoBiome-v8", "ForagaxTwoBiome-v9"):
         margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
         config["size"] = (width, 15)
@@ -270,5 +299,6 @@ def make(
         raise ValueError(f"Unknown observation type: {observation_type}")
 
     env_class = env_class_map[observation_type]
+    config["name"] = env_id
 
     return env_class(**config)

continual_foragax-0.22.0/tests/test_benchmark.py

@@ -0,0 +1,239 @@
+import chex
+import jax
+import jax.numpy as jnp
+
+from foragax.env import Actions, Biome, ForagaxObjectEnv, ForagaxRGBEnv, ForagaxWorldEnv
+from foragax.objects import FLOWER, WALL
+
+
+def test_benchmark_vision(benchmark):
+    env = ForagaxObjectEnv(size=7, aperture_size=3, objects=(WALL,))
+    params = env.default_params
+    key = jax.random.key(0)
+    _, state = env.reset(key, params)
+
+    grid = jnp.zeros((7, 7), dtype=int)
+    grid = grid.at[4, 3].set(1)
+    grid = grid.at[5, 3].set(1)
+    grid = grid.at[2, 0].set(1)
+    state = state.replace(object_grid=grid)
+
+    @jax.jit
+    def _run(state, key):
+        key, step_key = jax.random.split(key)
+        obs, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+        return obs, new_state
+
+    # warm-up
+    obs, new_state = _run(state, key)
+
+    expected = jnp.zeros((3, 3, 1), dtype=int)
+    expected = expected.at[2, 1, 0].set(1)
+
+    chex.assert_trees_all_equal(new_state.pos, jnp.array([3, 3]))
+    chex.assert_trees_all_equal(obs, expected)
+
+    def benchmark_fn():
+        # use a fixed key for benchmark consistency
+        _run(state, jax.random.key(1))[0].block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_creation(benchmark):
+    env = ForagaxObjectEnv(
+        size=1_000,
+        aperture_size=31,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+    )
+    params = env.default_params
+
+    @jax.jit
+    def _build(key):
+        _, state = env.reset(key, params)
+        return state
+
+    # no warm-up
+
+    def benchmark_fn():
+        _build(jax.random.key(1)).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_small_env(benchmark):
+    env = ForagaxObjectEnv(
+        size=1_000,
+        aperture_size=11,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.1, 0.1)),),
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+    key, reset_key = jax.random.split(key)
+    _, state = env.reset(reset_key, params)
+
+    @jax.jit
+    def _run(state, key):
+        def f(carry, _):
+            state, key = carry
+            key, step_key = jax.random.split(key, 2)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+            return (new_state, key), None
+
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=1000)
+        return final_state
+
+    key, run_key = jax.random.split(key)
+    _run(state, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(state, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_big_env(benchmark):
+    env = ForagaxObjectEnv(
+        size=10_000,
+        aperture_size=61,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+
+    # Reset is part of the setup, not benchmarked
+    key, reset_key = jax.random.split(key)
+    _, state = env.reset(reset_key, params)
+
+    @jax.jit
+    def _run(state, key):
+        def f(carry, _):
+            state, key = carry
+            key, step_key = jax.random.split(key, 2)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+            return (new_state, key), None
+
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
+        return final_state
+
+    # warm-up compilation
+    key, run_key = jax.random.split(key)
+    _run(state, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        # use a fixed key for benchmark consistency
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(state, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_vmap_env(benchmark):
+    num_envs = 100
+    env = ForagaxObjectEnv(
+        size=1_000,
+        aperture_size=11,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.1, 0.1)),),
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+
+    # Reset is part of the setup, not benchmarked
+    key, reset_key = jax.random.split(key)
+    reset_keys = jax.random.split(reset_key, num_envs)
+    states = jax.vmap(env.reset, in_axes=(0, None))(reset_keys, params)[1]
+
+    @jax.jit
+    def _run(states, key):
+        def f(carry, _):
+            states, key = carry
+            key, step_key = jax.random.split(key, 2)
+            step_keys = jax.random.split(step_key, num_envs)
+            _, new_states, _, _, _ = jax.vmap(env.step, in_axes=(0, 0, None, None))(
+                step_keys, states, Actions.DOWN, params
+            )
+            return (new_states, key), None
+
+        (final_states, _), _ = jax.lax.scan(f, (states, key), None, length=1000)
+        return final_states
+
+    # warm-up compilation
+    key, run_key = jax.random.split(key)
+    _run(states, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        # use a fixed key for benchmark consistency
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(states, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_small_env_color(benchmark):
+    env = ForagaxRGBEnv(
+        size=1_000,
+        aperture_size=15,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+    key, reset_key = jax.random.split(key)
+    _, state = env.reset(reset_key, params)
+
+    @jax.jit
+    def _run(state, key):
+        def f(carry, _):
+            state, key = carry
+            key, step_key = jax.random.split(key, 2)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+            return (new_state, key), None
+
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
+        return final_state
+
+    key, run_key = jax.random.split(key)
+    _run(state, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(state, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)
+
+
+def test_benchmark_small_env_world(benchmark):
+    env = ForagaxWorldEnv(
+        size=1_000,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+    )
+    params = env.default_params
+    key = jax.random.key(0)
+    key, reset_key = jax.random.split(key)
+    _, state = env.reset(reset_key, params)
+
+    @jax.jit
+    def _run(state, key):
+        def f(carry, _):
+            state, key = carry
+            key, step_key = jax.random.split(key, 2)
+            _, new_state, _, _, _ = env.step(step_key, state, Actions.DOWN, params)
+            return (new_state, key), None
+
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=100)
+        return final_state
+
+    key, run_key = jax.random.split(key)
+    _run(state, run_key).pos.block_until_ready()
+
+    def benchmark_fn():
+        key, run_key = jax.random.split(jax.random.key(1))
+        _run(state, run_key).pos.block_until_ready()
+
+    benchmark(benchmark_fn)