continual-foragax 0.27.0__tar.gz → 0.28.1__tar.gz

{continual_foragax-0.27.0 → continual_foragax-0.28.1}/PKG-INFO

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

{continual_foragax-0.27.0 → continual_foragax-0.28.1}/pyproject.toml

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

{continual_foragax-0.27.0 → continual_foragax-0.28.1}/src/continual_foragax.egg-info/PKG-INFO

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

{continual_foragax-0.27.0 → continual_foragax-0.28.1}/src/foragax/env.py

@@ -76,6 +76,7 @@ class ForagaxEnv(environment.Environment):
         nowrap: bool = False,
         deterministic_spawn: bool = False,
         teleport_interval: Optional[int] = None,
+        observation_type: str = "object",
     ):
         super().__init__()
         self._name = name
@@ -83,14 +84,22 @@ class ForagaxEnv(environment.Environment):
             size = (size, size)
         self.size = size
 
-        if isinstance(aperture_size, int):
-            aperture_size = (aperture_size, aperture_size)
-        self.aperture_size = aperture_size
+        # Handle aperture_size = -1 for world view
+        if isinstance(aperture_size, int) and aperture_size == -1:
+            self.full_world = True
+            self.aperture_size = self.size  # Use full size for consistency
+        else:
+            self.full_world = False
+            if isinstance(aperture_size, int):
+                aperture_size = (aperture_size, aperture_size)
+            self.aperture_size = aperture_size
+
+        self.observation_type = observation_type
         self.nowrap = nowrap
         self.deterministic_spawn = deterministic_spawn
         self.teleport_interval = teleport_interval
         objects = (EMPTY,) + objects
-        if self.nowrap:
+        if self.nowrap and not self.full_world:
             objects = objects + (PADDING,)
         self.objects = objects
         self.weather_object = None
@@ -152,6 +161,29 @@ class ForagaxEnv(environment.Environment):
             )
             self.biome_masks.append(mask)
 
+        # Compute unique colors and mapping for partial observability (for 'color' observation_type)
+        # Exclude EMPTY (index 0) from color channels
+        object_colors_no_empty = self.object_colors[1:]
+
+        # Find unique colors in order of first appearance
+        unique_colors = []
+        color_indices = jnp.zeros(len(object_colors_no_empty), dtype=int)
+        color_map = {}
+        next_channel = 0
+
+        for i, color in enumerate(object_colors_no_empty):
+            color_tuple = tuple(color.tolist())
+            if color_tuple not in color_map:
+                color_map[color_tuple] = next_channel
+                unique_colors.append(color)
+                next_channel += 1
+            color_indices = color_indices.at[i].set(color_map[color_tuple])
+
+        self.unique_colors = jnp.array(unique_colors)
+        self.num_color_channels = len(unique_colors)
+        # color_indices maps from object_id-1 to color_channel_index
+        self.object_to_color_map = color_indices
+
     @property
     def default_params(self) -> EnvParams:
         return EnvParams(
@@ -412,6 +444,102 @@ class ForagaxEnv(environment.Environment):
 
         return aperture
 
+    def get_obs(self, state: EnvState, params: EnvParams, key=None) -> jax.Array:
+        """Get observation based on observation_type and full_world."""
+        # Decode grid for observation
+        obs_grid = jnp.maximum(0, state.object_grid)
+
+        if self.full_world:
+            return self._get_world_obs(obs_grid, state)
+        else:
+            grid = self._get_aperture(obs_grid, state.pos)
+            return self._get_aperture_obs(grid, state)
+
+    def _get_world_obs(self, obs_grid: jax.Array, state: EnvState) -> jax.Array:
+        """Get world observation."""
+        num_obj_types = len(self.object_ids)
+        if self.observation_type == "object":
+            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
+        elif self.observation_type == "rgb":
+            obs = jax.nn.one_hot(obs_grid, num_obj_types)
+            # Agent position
+            obs = obs.at[state.pos[1], state.pos[0], :].set(0)
+            obs = obs.at[state.pos[1], state.pos[0], -1].set(1)
+            colors = self.object_colors / 255.0
+            obs = jnp.tensordot(obs, colors, axes=1)
+            return obs
+        elif self.observation_type == "color":
+            # Handle case with no objects (only EMPTY)
+            if self.num_color_channels == 0:
+                return jnp.zeros(obs_grid.shape + (0,), dtype=jnp.float32)
+            # Map object IDs to color channel indices
+            color_channels = jnp.where(
+                obs_grid == 0,
+                -1,
+                jnp.take(self.object_to_color_map, obs_grid - 1, axis=0),
+            )
+            obs = jax.nn.one_hot(color_channels, self.num_color_channels, axis=-1)
+            return obs
+        else:
+            raise ValueError(f"Unknown observation_type: {self.observation_type}")
+
+    def _get_aperture_obs(self, aperture: jax.Array, state: EnvState) -> jax.Array:
+        """Get aperture observation."""
+        if self.observation_type == "object":
+            num_obj_types = len(self.object_ids)
+            obs = jax.nn.one_hot(aperture, num_obj_types, axis=-1)
+            return obs
+        elif self.observation_type == "rgb":
+            num_obj_types = len(self.object_ids)
+            aperture_one_hot = jax.nn.one_hot(aperture, num_obj_types)
+            colors = self.object_colors / 255.0
+            obs = jnp.tensordot(aperture_one_hot, colors, axes=1)
+            return obs
+        elif self.observation_type == "color":
+            # Handle case with no objects (only EMPTY)
+            if self.num_color_channels == 0:
+                return jnp.zeros(aperture.shape + (0,), dtype=jnp.float32)
+            # Map object IDs to color channel indices
+            color_channels = jnp.where(
+                aperture == 0,
+                -1,  # Special value for EMPTY
+                jnp.take(self.object_to_color_map, aperture - 1, axis=0),
+            )
+            obs = jax.nn.one_hot(color_channels, self.num_color_channels, axis=-1)
+            return obs
+        else:
+            raise ValueError(f"Unknown observation_type: {self.observation_type}")
+
+    def observation_space(self, params: EnvParams) -> spaces.Box:
+        """Observation space based on observation_type and full_world."""
+        if self.full_world:
+            size = tuple(reversed(self.size))
+        else:
+            size = self.aperture_size
+
+        if self.observation_type == "rgb":
+            channels = 3
+        elif self.observation_type == "object":
+            num_obj_types = len(self.objects)
+            channels = num_obj_types
+        elif self.observation_type == "color":
+            channels = self.num_color_channels
+        else:
+            raise ValueError(f"Unknown observation_type: {self.observation_type}")
+
+        obs_shape = (*size, channels)
+
+        return spaces.Box(0, 1, obs_shape, float)
+
     @partial(jax.jit, static_argnames=("self", "render_mode"))
     def render(self, state: EnvState, params: EnvParams, render_mode: str = "world"):
         """Render the environment state."""
@@ -520,133 +648,3 @@ class ForagaxEnv(environment.Environment):
             raise ValueError(f"Unknown render_mode: {render_mode}")
 
         return img
-
-
-class ForagaxObjectEnv(ForagaxEnv):
-    """Foragax environment with object-based aperture observation."""
-
-    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,
-        teleport_interval: Optional[int] = None,
-    ):
-        super().__init__(
-            name,
-            size,
-            aperture_size,
-            objects,
-            biomes,
-            nowrap,
-            deterministic_spawn,
-            teleport_interval,
-        )
-
-        # Compute unique colors and mapping for partial observability
-        # Exclude EMPTY (index 0) from color channels
-        object_colors_no_empty = self.object_colors[1:]
-
-        # Find unique colors in order of first appearance
-        unique_colors = []
-        color_indices = jnp.zeros(len(object_colors_no_empty), dtype=int)
-        color_map = {}
-        next_channel = 0
-
-        for i, color in enumerate(object_colors_no_empty):
-            color_tuple = tuple(color.tolist())
-            if color_tuple not in color_map:
-                color_map[color_tuple] = next_channel
-                unique_colors.append(color)
-                next_channel += 1
-            color_indices = color_indices.at[i].set(color_map[color_tuple])
-
-        self.unique_colors = jnp.array(unique_colors)
-        self.num_color_channels = len(unique_colors)
-        # color_indices maps from object_id-1 to color_channel_index
-        self.object_to_color_map = color_indices
-
-    def get_obs(self, state: EnvState, params: EnvParams, key=None) -> jax.Array:
-        # Decode grid for observation
-        obs_grid = jnp.maximum(0, state.object_grid)
-        aperture = self._get_aperture(obs_grid, state.pos)
-
-        # Handle case with no objects (only EMPTY)
-        if self.num_color_channels == 0:
-            return jnp.zeros(aperture.shape + (0,), dtype=jnp.float32)
-
-        # Map object IDs to color channel indices
-        # aperture contains object IDs (0 = EMPTY, 1+ = objects)
-        # For EMPTY (0), we want no color channel activated
-        # For objects (1+), map to color channel using object_to_color_map
-        color_channels = jnp.where(
-            aperture == 0,
-            -1,  # Special value for EMPTY
-            jnp.take(self.object_to_color_map, aperture - 1, axis=0),
-        )
-
-        # Create one-hot encoding for color channels
-        # jax.nn.one_hot produces all zeros for -1 (EMPTY positions)
-        obs = jax.nn.one_hot(color_channels, self.num_color_channels, axis=-1)
-
-        return obs
-
-    def observation_space(self, params: EnvParams) -> spaces.Box:
-        obs_shape = (
-            self.aperture_size[0],
-            self.aperture_size[1],
-            self.num_color_channels,
-        )
-        return spaces.Box(0, 1, obs_shape, float)
-
-
-class ForagaxRGBEnv(ForagaxEnv):
-    """Foragax environment with color-based aperture observation."""
-
-    def get_obs(self, state: EnvState, params: EnvParams, key=None) -> jax.Array:
-        num_obj_types = len(self.object_ids)
-        # Decode grid for observation
-        obs_grid = jnp.maximum(0, state.object_grid)
-        aperture = self._get_aperture(obs_grid, state.pos)
-        aperture_one_hot = jax.nn.one_hot(aperture, num_obj_types)
-
-        # Agent position is always at the center of the aperture
-        center = (self.aperture_size[1] // 2, self.aperture_size[0] // 2)
-        aperture_one_hot = aperture_one_hot.at[center[0], center[1], :].set(0)
-        aperture_one_hot = aperture_one_hot.at[center[0], center[1], -1].set(1)
-
-        colors = self.object_colors / 255.0
-        obs = jnp.tensordot(aperture_one_hot, colors, axes=1)
-        return obs
-
-    def observation_space(self, params: EnvParams) -> spaces.Box:
-        obs_shape = (self.aperture_size[0], self.aperture_size[1], 3)
-        return spaces.Box(0, 1, obs_shape, float)
-
-
-class ForagaxWorldEnv(ForagaxEnv):
-    """Foragax environment with world observation."""
-
-    def get_obs(self, state: EnvState, params: EnvParams, key=None) -> jax.Array:
-        num_obj_types = len(self.object_ids)
-        # Decode grid for observation
-        obs_grid = jnp.maximum(0, state.object_grid)
-        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
-
-    def observation_space(self, params: EnvParams) -> spaces.Box:
-        num_obj_types = len(self.object_ids)
-        obs_shape = (self.size[1], self.size[0], num_obj_types)
-        return spaces.Box(0, 1, obs_shape, float)

{continual_foragax-0.27.0 → continual_foragax-0.28.1}/src/foragax/registry.py

@@ -1,13 +1,11 @@
 """Factory functions for creating Foragax environment variants."""
 
+import warnings
 from typing import Any, Dict, Optional, Tuple
 
 from foragax.env import (
     Biome,
     ForagaxEnv,
-    ForagaxObjectEnv,
-    ForagaxRGBEnv,
-    ForagaxWorldEnv,
 )
 from foragax.objects import (
     BROWN_MOREL,
@@ -347,21 +345,23 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
 
 def make(
     env_id: str,
-    observation_type: str = "object",
+    observation_type: str = "color",
     aperture_size: Optional[Tuple[int, int]] = (5, 5),
     file_index: int = 0,
     nowrap: Optional[bool] = None,
+    **kwargs: Any,
 ) -> ForagaxEnv:
     """Create a Foragax environment.
 
     Args:
         env_id: The ID of the environment to create.
-        observation_type: The type of observation to use. One of "object", "rgb", or "world".
-        aperture_size: The size of the agent's observation aperture. If None, the default
-            for the environment is used.
+        observation_type: The type of observation to use. One of "object", "rgb", or "color".
+        aperture_size: The size of the agent's observation aperture. If -1, full world observation.
+            If None, the default for the environment is used.
         file_index: File index for weather objects. nowrap: If True, disables
         wrapping around environment boundaries. If None, uses defaults per
         environment.
+        **kwargs: Additional keyword arguments to pass to the ForagaxEnv constructor.
 
     Returns:
         A Foragax environment instance.
@@ -371,11 +371,15 @@ def make(
 
     config = ENV_CONFIGS[env_id].copy()
     if isinstance(aperture_size, int):
-        aperture_size = (aperture_size, aperture_size)
+        if aperture_size == -1:
+            aperture_size = -1  # keep as -1
+        else:
+            aperture_size = (aperture_size, aperture_size)
     config["aperture_size"] = aperture_size
     if nowrap is not None:
         config["nowrap"] = nowrap
 
+    # Handle special size and biome configurations
     if env_id in (
         "ForagaxTwoBiome-v7",
         "ForagaxTwoBiome-v8",
@@ -384,7 +388,10 @@ def make(
         "ForagaxTwoBiome-v15",
         "ForagaxTwoBiome-v16",
     ):
-        margin = aperture_size[1] // 2 + 1
+        if aperture_size == -1:
+            margin = 0  # for world view, no margin needed
+        else:
+            margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
         config["size"] = (width, 15)
         config["biomes"] = (
@@ -403,7 +410,10 @@ def make(
         )
 
     if env_id in ("ForagaxTwoBiome-v11", "ForagaxTwoBiome-v12"):
-        margin = aperture_size[1] // 2 + 1
+        if aperture_size == -1:
+            margin = 0
+        else:
+            margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
         config["size"] = (width, 15)
         config["biomes"] = (
@@ -422,7 +432,10 @@ def make(
         )
 
     if env_id in ("ForagaxWeather-v3", "ForagaxWeather-v4"):
-        margin = aperture_size[1] // 2 + 1
+        if aperture_size == -1:
+            margin = 0
+        else:
+            margin = aperture_size[1] // 2 + 1
         width = 2 * margin + 9
         config["size"] = (15, width)
         config["biomes"] = (
@@ -456,16 +469,20 @@ def make(
     if env_id == "ForagaxTwoBiome-v16":
         config["teleport_interval"] = 10000
 
-    env_class_map = {
-        "object": ForagaxObjectEnv,
-        "rgb": ForagaxRGBEnv,
-        "world": ForagaxWorldEnv,
-    }
+    # Backward compatibility: map "world" to "object" with full world
+    if observation_type == "world":
+        # add deprecation warning
+        warnings.warn(
+            "'world' observation_type is deprecated. Use 'object' with aperture_size=-1 instead.",
+            DeprecationWarning,
+        )
+        observation_type = "object"
+        config["aperture_size"] = -1
 
-    if observation_type not in env_class_map:
-        raise ValueError(f"Unknown observation type: {observation_type}")
+    if observation_type not in ("object", "rgb", "color"):
+        raise ValueError(f"Unknown observation_type: {observation_type}")
 
-    env_class = env_class_map[observation_type]
     config["name"] = env_id
+    config["observation_type"] = observation_type
 
-    return env_class(**config)
+    return ForagaxEnv(**config, **kwargs)

{continual_foragax-0.27.0 → continual_foragax-0.28.1}/tests/test_benchmark.py

@@ -2,12 +2,12 @@ import chex
 import jax
 import jax.numpy as jnp
 
-from foragax.env import Actions, Biome, ForagaxObjectEnv, ForagaxRGBEnv, ForagaxWorldEnv
+from foragax.env import Actions, Biome, ForagaxEnv
 from foragax.objects import FLOWER, WALL
 
 
 def test_benchmark_vision(benchmark):
-    env = ForagaxObjectEnv(size=7, aperture_size=3, objects=(WALL,))
+    env = ForagaxEnv(size=7, aperture_size=3, objects=(WALL,), observation_type="color")
     params = env.default_params
     key = jax.random.key(0)
     _, state = env.reset(key, params)
@@ -41,11 +41,12 @@ def test_benchmark_vision(benchmark):
 
 
 def test_benchmark_creation(benchmark):
-    env = ForagaxObjectEnv(
+    env = ForagaxEnv(
         size=1_000,
         aperture_size=31,
         objects=(WALL, FLOWER),
         biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+        observation_type="color",
     )
     params = env.default_params
 
@@ -63,11 +64,12 @@ def test_benchmark_creation(benchmark):
 
 
 def test_benchmark_small_env(benchmark):
-    env = ForagaxObjectEnv(
+    env = ForagaxEnv(
         size=1_000,
         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)
@@ -96,11 +98,12 @@ def test_benchmark_small_env(benchmark):
 
 
 def test_benchmark_big_env(benchmark):
-    env = ForagaxObjectEnv(
+    env = ForagaxEnv(
         size=10_000,
         aperture_size=61,
         objects=(WALL, FLOWER),
         biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+        observation_type="color",
     )
     params = env.default_params
     key = jax.random.key(0)
@@ -134,11 +137,12 @@ def test_benchmark_big_env(benchmark):
 
 def test_benchmark_vmap_env(benchmark):
     num_envs = 100
-    env = ForagaxObjectEnv(
+    env = ForagaxEnv(
         size=1_000,
         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)
@@ -175,11 +179,12 @@ def test_benchmark_vmap_env(benchmark):
 
 
 def test_benchmark_small_env_color(benchmark):
-    env = ForagaxRGBEnv(
+    env = ForagaxEnv(
         size=1_000,
         aperture_size=15,
         objects=(WALL, FLOWER),
         biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+        observation_type="rgb",
     )
     params = env.default_params
     key = jax.random.key(0)
@@ -208,10 +213,12 @@ def test_benchmark_small_env_color(benchmark):
 
 
 def test_benchmark_small_env_world(benchmark):
-    env = ForagaxWorldEnv(
+    env = ForagaxEnv(
         size=1_000,
         objects=(WALL, FLOWER),
         biomes=(Biome(object_frequencies=(0.05, 0.05)),),
+        observation_type="object",
+        aperture_size=-1,
     )
     params = env.default_params
     key = jax.random.key(0)