continual-foragax 0.39.0__tar.gz → 0.40.0__tar.gz

{continual_foragax-0.39.0 → continual_foragax-0.40.0}/PKG-INFO

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

{continual_foragax-0.39.0 → continual_foragax-0.40.0}/pyproject.toml

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

{continual_foragax-0.39.0 → continual_foragax-0.40.0}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.39.0 → continual_foragax-0.40.0}/src/foragax/env.py

@@ -6,7 +6,7 @@ Source: https://github.com/andnp/Forager
 from dataclasses import dataclass
 from enum import IntEnum
 from functools import partial
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Tuple, Union
 
 import jax
 import jax.numpy as jnp
@@ -128,7 +128,6 @@ class ForagaxEnv(environment.Environment):
         biomes: Tuple[Biome, ...] = (Biome(object_frequencies=()),),
         nowrap: bool = False,
         deterministic_spawn: bool = False,
-        teleport_interval: Optional[int] = None,
         observation_type: str = "object",
         dynamic_biomes: bool = False,
         biome_consumption_threshold: float = 0.9,
@@ -154,7 +153,6 @@ class ForagaxEnv(environment.Environment):
         self.observation_type = observation_type
         self.nowrap = nowrap
         self.deterministic_spawn = deterministic_spawn
-        self.teleport_interval = teleport_interval
         self.dynamic_biomes = dynamic_biomes
         self.biome_consumption_threshold = biome_consumption_threshold
         self.dynamic_biome_spawn_empty = dynamic_biome_spawn_empty
@@ -391,23 +389,6 @@ class ForagaxEnv(environment.Environment):
         is_blocking = self.object_blocking[obj_at_new_pos]
         pos = jax.lax.select(is_blocking, state.pos, new_pos)
 
-        # Check for automatic teleport
-        if self.teleport_interval is not None:
-            should_teleport = jnp.mod(state.time + 1, self.teleport_interval) == 0
-        else:
-            should_teleport = False
-
-        def teleport_fn():
-            # Calculate squared distances from current position to each biome center
-            diffs = self.biome_centers_jax - pos
-            distances = jnp.sum(diffs**2, axis=1)
-            # Find the index of the furthest biome center
-            furthest_idx = jnp.argmax(distances)
-            new_pos = self.biome_centers_jax[furthest_idx]
-            return new_pos
-
-        pos = jax.lax.cond(should_teleport, teleport_fn, lambda: pos)
-
         # 2. HANDLE COLLISIONS AND REWARDS
         obj_at_pos = current_objects[pos[1], pos[0]]
         is_collectable = self.object_collectable[obj_at_pos]

{continual_foragax-0.39.0 → continual_foragax-0.40.0}/src/foragax/registry.py

@@ -636,9 +636,6 @@ def make(
             biome2_deathcap,
         )
 
-    if env_id == "ForagaxTwoBiome-v16":
-        config["teleport_interval"] = 10000
-
     # Backward compatibility: map "world" to "object" with full world
     if observation_type == "world":
         # add deprecation warning

{continual_foragax-0.39.0 → continual_foragax-0.40.0}/tests/test_benchmark.py

@@ -2,6 +2,7 @@ import chex
 import jax
 import jax.numpy as jnp
 
+from foragax import registry
 from foragax.env import Actions, Biome, ForagaxEnv
 from foragax.objects import FLOWER, WALL
 
@@ -40,7 +41,7 @@ def test_benchmark_vision(benchmark):
     benchmark(benchmark_fn)
 
 
-def test_benchmark_creation(benchmark):
+def test_benchmark_reset(benchmark):
     env = ForagaxEnv(
         size=1_000,
         aperture_size=31,
@@ -63,6 +64,40 @@ def test_benchmark_creation(benchmark):
     benchmark(benchmark_fn)
 
 
+def test_benchmark_tiny_env(benchmark):
+    env = ForagaxEnv(
+        size=15,
+        aperture_size=11,
+        objects=(WALL, FLOWER),
+        biomes=(Biome(object_frequencies=(0.1, 0.1)),),
+        observation_type="color",
+    )
+    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=10)
+        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(benchmark):
     env = ForagaxEnv(
         size=1_000,
@@ -84,7 +119,7 @@ def test_benchmark_small_env(benchmark):
             _, 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)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
         return final_state
 
     key, run_key = jax.random.split(key)
@@ -120,7 +155,7 @@ def test_benchmark_big_env(benchmark):
             _, 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)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
         return final_state
 
     # warm-up compilation
@@ -136,7 +171,7 @@ def test_benchmark_big_env(benchmark):
 
 
 def test_benchmark_vmap_env(benchmark):
-    num_envs = 100
+    num_envs = 10
     env = ForagaxEnv(
         size=1_000,
         aperture_size=11,
@@ -163,7 +198,7 @@ def test_benchmark_vmap_env(benchmark):
             )
             return (new_states, key), None
 
-        (final_states, _), _ = jax.lax.scan(f, (states, key), None, length=1000)
+        (final_states, _), _ = jax.lax.scan(f, (states, key), None, length=10)
         return final_states
 
     # warm-up compilation
@@ -199,7 +234,7 @@ def test_benchmark_small_env_color(benchmark):
             _, 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)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
         return final_state
 
     key, run_key = jax.random.split(key)
@@ -233,7 +268,63 @@ def test_benchmark_small_env_world(benchmark):
             _, 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)
+        (final_state, _), _ = jax.lax.scan(f, (state, key), None, length=10)
+        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_diwali_v5(benchmark):
+    env = registry.make("ForagaxDiwali-v5")
+    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=10)
+        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_sine_two_biome_v1(benchmark):
+    env = registry.make("ForagaxSineTwoBiome-v1")
+    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=10)
         return final_state
 
     key, run_key = jax.random.split(key)

{continual_foragax-0.39.0 → continual_foragax-0.40.0}/tests/test_foragax.py

@@ -977,66 +977,6 @@ def test_single_color_all_objects():
         chex.assert_trees_all_equal(obj_obs, jnp.array([1.0]))
 
 
-def test_teleporting():
-    """Test automatic teleporting to the furthest biome center from current position."""
-    key = jax.random.key(0)
-
-    # Create environment with two biomes and teleport every 5 steps
-    env = ForagaxEnv(
-        size=(10, 10),
-        aperture_size=(3, 3),
-        objects=(WALL,),
-        biomes=(
-            Biome(
-                start=(1, 1), stop=(5, 5), object_frequencies=(0.0,)
-            ),  # Biome 0 center at (2,2)
-            Biome(
-                start=(6, 6), stop=(10, 10), object_frequencies=(0.0,)
-            ),  # Biome 1 center at (7,7)
-        ),
-        teleport_interval=5,
-        nowrap=True,
-        observation_type="color",
-    )
-    params = env.default_params
-
-    # Reset and get initial state
-    obs, state = env.reset_env(key, params)
-
-    # Agent should start at center (5, 5), which is not in either biome initially
-    # But let's manually place it in biome 0 for testing
-    state = state.replace(pos=jnp.array([2, 2]))  # In biome 0
-
-    # Step 4 times (time will be 0,1,2,3,4 after these steps)
-    for i in range(4):
-        key, step_key = jax.random.split(key)
-        obs, state, _, _, _ = env.step_env(step_key, state, Actions.LEFT, params)
-
-    # After 4 steps, time=4, next step should teleport (4+1) % 5 == 0
-    key, step_key = jax.random.split(key)
-    obs, state, _, _, _ = env.step_env(step_key, state, Actions.LEFT, params)
-
-    # Should have teleported to the furthest biome center (biome 1 center)
-    # From (2,2), (7,7) is further than (2,2)
-    expected_pos = jnp.array([7, 7])
-    chex.assert_trees_all_equal(state.pos, expected_pos)
-
-    # Step another 4 times to reach time=9, then teleport back
-    # From [7,7], move right to stay in biome: [8,7], [9,7], then stay
-    for i in range(4):
-        key, step_key = jax.random.split(key)
-        obs, state, _, _, _ = env.step(step_key, state, Actions.RIGHT, params)
-
-    # After another 4 steps, time=9, next step should teleport (9+1) % 5 == 0
-    key, step_key = jax.random.split(key)
-    obs, state, _, _, _ = env.step(step_key, state, Actions.RIGHT, params)
-
-    # Should teleport to the furthest biome center (biome 0 center)
-    # From (7,7), (2,2) is further than (7,7)
-    expected_pos = jnp.array([2, 2])
-    chex.assert_trees_all_equal(state.pos, expected_pos)
-
-
 def test_info_discount():
     """Test that info contains discount."""
     key = jax.random.key(0)