continual-foragax 0.27.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

{continual_foragax-0.27.0.dist-info → continual_foragax-0.28.0.dist-info}/METADATA

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

{continual_foragax-0.27.0.dist-info → continual_foragax-0.28.0.dist-info}/RECORD

@@ -1,8 +1,8 @@
 foragax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 foragax/colors.py,sha256=rqNPiywP4Nvr0POhsGpasRk-nMMTS3DOwFRUgperlUk,2065
-foragax/env.py,sha256=2Z2-3wjL_gHCqehMFsZA761VlUWQvsTU6PawoTtzcCY,25103
+foragax/env.py,sha256=lg6nRT2goO84FNGdL0CbUTGj1iBiQXUYOWf2qavMMvE,25658
 foragax/objects.py,sha256=FCLZ-8d7qq9VMTG6G-TaRt842-sjgB0-DH0IoHwwngI,9503
-foragax/registry.py,sha256=7-6VDN1MKVEvX_1u5G8NkSpv9BccEmtjJa77-OTNg3A,14324
+foragax/registry.py,sha256=HysNaZs1tcbAcr53l8Cb2NeZ-_FmE6OpUe_zIks-ObM,15089
 foragax/rendering.py,sha256=bms7wvBZTofoR-K-2QD2Ggeed7Viw8uwAEiEpEM3eSo,2768
 foragax/weather.py,sha256=KNAiwuFz8V__6G75vZIWQKPocLzXqxXn-Vt4TbHIpcA,1258
 foragax/data/ECA_non-blended_custom/TG_SOUID100897.txt,sha256=N7URbX6VlCZvCboUogYjMzy1I-0cfNPOn0QTLSHHfQ0,1776751
@@ -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.27.0.dist-info/METADATA,sha256=Axd__jkh5OIh1T9CmzB9IarAguoWNvsECjllZ20mG3w,4897
-continual_foragax-0.27.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-continual_foragax-0.27.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
-continual_foragax-0.27.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
-continual_foragax-0.27.0.dist-info/RECORD,,
+continual_foragax-0.28.0.dist-info/METADATA,sha256=ijxVOxZXSpbQ3ORIHKrUxgciSiuBTqxpA-nE2TWvitQ,4897
+continual_foragax-0.28.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+continual_foragax-0.28.0.dist-info/entry_points.txt,sha256=Qiu6iE_XudrDO_bVAMeA435h4PO9ourt8huvSHiuMPc,41
+continual_foragax-0.28.0.dist-info/top_level.txt,sha256=-z3SDK6RfLIcLI24n8rdbeFzlVY3hunChzlu-v1Fncs,8
+continual_foragax-0.28.0.dist-info/RECORD,,

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,9 +84,17 @@ 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
@@ -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)

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)