continual-foragax 0.7.0__tar.gz → 0.8.0__tar.gz

continual_foragax-0.8.0/PKG-INFO

@@ -0,0 +1,121 @@
+Metadata-Version: 2.4
+Name: continual-foragax
+Version: 0.8.0
+Summary: A continual reinforcement learning benchmark
+Author-email: Steven Tang <stang5@ualberta.ca>
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: gymnax
+Requires-Dist: six; python_version < "3.10"
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-benchmark; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+
+# foragax
+
+Foragax is a lightweight, JAX-first grid-world environment suite for continual / procedural
+experiments. It provides a small collection of environment variants (weather, multi-biome,
+etc.), a registry factory for easy construction, and simple example scripts for plotting and
+visualization.
+
+This version is a [Gymnax](https://github.com/RobertTLange/gymnax) environment implemented in JAX. The original implementation of Forager is implemented in Numba is available at [andnp/forager](https://github.com/andnp/forager). In addition to the original features, this implementation includes: biomes, visualization, and a weather environment.
+
+Key ideas:
+
+- Functional, JAX-friendly API (explicit PRNG keys, immutable env state objects).
+- Multiple observation modalities: Object and RGB, as well as aperture based or
+	full-world observations.
+- Customizable biomes
+- Customizable object placement, respawning, and rewards.
+- Visualization via RGB rendering and plotting.
+
+## Quickstart
+
+We recommend installing with pip from https://pypi.org/project/continual-foragax/.
+
+```bash
+pip install continual-foragax
+```
+
+We support Python 3.8 through Python 3.13.
+
+The codebase expects JAX and other numeric dependencies. If you don't have JAX installed, see
+the JAX install instructions for your platform; the project `uv.lock` pins compatible versions.
+
+## Minimal example (from examples)
+
+Use the registry factory to create an environment and run it with JAX-style RNG keys and an
+explicit environment state.
+
+```python
+from foragax.registry import make
+import jax
+
+# create env (observation_type is one of: 'object', 'rgb', 'world')
+env = make(
+		"ForagaxWeather-v1",
+		aperture_size=5,
+		observation_type="object",
+)
+
+# environment parameters and RNG
+env_params = env.default_params
+key = jax.random.key(0)
+key, key_reset = jax.random.split(key)
+
+# reset returns (obs, env_state)
+_, env_state = env.reset(key_reset, env_params)
+
+# sampling an action and stepping (functional-style)
+key, key_act, key_step = jax.random.split(key, 3)
+action = env.action_space(env_params).sample(key_act)
+_, next_env_state, reward, done, info = env.step(key_step, env_state, action, env_params)
+
+# rendering supports multiple modes: 'world' and 'aperture'
+frame = env.render(env_state, env_params, render_mode="aperture")
+```
+
+See `examples/plot.py` and `examples/visualize.py` for runnable scripts that produce a sample
+plot and saved videos using Gym/Gymnasium helpers.
+
+## Registry and included environments
+
+Use `foragax.registry.make` to construct environments by id. Example environment ids include:
+
+- `ForagaxTwoBiomeSmall-v1` / `-v2` — hand-crafted small multi-biome layouts
+- `ForagaxWeather-v1` — small weather-driven two-biome environment used by examples
+
+The `make` factory accepts the following notable kwargs:
+
+- `observation_type`: one of `"object"`, `"rgb"`, or `"world"`.
+- `aperture_size`: integer or tuple controlling the agent's local observation aperture.
+- `file_index`: used to pick weather locations.
+
+## Custom objects and extensions
+
+The codebase includes an object system for placing items into biomes and controlling
+behaviour (rewards, respawn / regen behavior, blocking/collectable flags). See
+`foragax.objects` for the canonical object definitions and helpers like
+`create_weather_objects` used by the registry.
+
+If you want to add new object classes, follow the examples in `foragax.objects` and add the
+class into registry configs or construct environments programmatically.
+
+## Design notes
+
+- JAX-first: RNG keys and immutable env state are passed explicitly so environments can be
+	stepped inside JIT/pmapped loops if desired.
+- Small, composable environment variants are provided through the registry (easy to add more).
+
+## Examples
+
+- `examples/plot.py` — runs a short random policy in `ForagaxWeather-v1` and produces a
+	temperature vs reward plot (saves to `plots/sample_plot.png`).
+- `examples/visualize.py` — runs environments at multiple aperture sizes and saves short
+	videos under `videos/` using `save_video`.
+
+## Development
+
+Run unit tests via pytest.

continual_foragax-0.8.0/README.md

@@ -0,0 +1,106 @@
+# foragax
+
+Foragax is a lightweight, JAX-first grid-world environment suite for continual / procedural
+experiments. It provides a small collection of environment variants (weather, multi-biome,
+etc.), a registry factory for easy construction, and simple example scripts for plotting and
+visualization.
+
+This version is a [Gymnax](https://github.com/RobertTLange/gymnax) environment implemented in JAX. The original implementation of Forager is implemented in Numba is available at [andnp/forager](https://github.com/andnp/forager). In addition to the original features, this implementation includes: biomes, visualization, and a weather environment.
+
+Key ideas:
+
+- Functional, JAX-friendly API (explicit PRNG keys, immutable env state objects).
+- Multiple observation modalities: Object and RGB, as well as aperture based or
+	full-world observations.
+- Customizable biomes
+- Customizable object placement, respawning, and rewards.
+- Visualization via RGB rendering and plotting.
+
+## Quickstart
+
+We recommend installing with pip from https://pypi.org/project/continual-foragax/.
+
+```bash
+pip install continual-foragax
+```
+
+We support Python 3.8 through Python 3.13.
+
+The codebase expects JAX and other numeric dependencies. If you don't have JAX installed, see
+the JAX install instructions for your platform; the project `uv.lock` pins compatible versions.
+
+## Minimal example (from examples)
+
+Use the registry factory to create an environment and run it with JAX-style RNG keys and an
+explicit environment state.
+
+```python
+from foragax.registry import make
+import jax
+
+# create env (observation_type is one of: 'object', 'rgb', 'world')
+env = make(
+		"ForagaxWeather-v1",
+		aperture_size=5,
+		observation_type="object",
+)
+
+# environment parameters and RNG
+env_params = env.default_params
+key = jax.random.key(0)
+key, key_reset = jax.random.split(key)
+
+# reset returns (obs, env_state)
+_, env_state = env.reset(key_reset, env_params)
+
+# sampling an action and stepping (functional-style)
+key, key_act, key_step = jax.random.split(key, 3)
+action = env.action_space(env_params).sample(key_act)
+_, next_env_state, reward, done, info = env.step(key_step, env_state, action, env_params)
+
+# rendering supports multiple modes: 'world' and 'aperture'
+frame = env.render(env_state, env_params, render_mode="aperture")
+```
+
+See `examples/plot.py` and `examples/visualize.py` for runnable scripts that produce a sample
+plot and saved videos using Gym/Gymnasium helpers.
+
+## Registry and included environments
+
+Use `foragax.registry.make` to construct environments by id. Example environment ids include:
+
+- `ForagaxTwoBiomeSmall-v1` / `-v2` — hand-crafted small multi-biome layouts
+- `ForagaxWeather-v1` — small weather-driven two-biome environment used by examples
+
+The `make` factory accepts the following notable kwargs:
+
+- `observation_type`: one of `"object"`, `"rgb"`, or `"world"`.
+- `aperture_size`: integer or tuple controlling the agent's local observation aperture.
+- `file_index`: used to pick weather locations.
+
+## Custom objects and extensions
+
+The codebase includes an object system for placing items into biomes and controlling
+behaviour (rewards, respawn / regen behavior, blocking/collectable flags). See
+`foragax.objects` for the canonical object definitions and helpers like
+`create_weather_objects` used by the registry.
+
+If you want to add new object classes, follow the examples in `foragax.objects` and add the
+class into registry configs or construct environments programmatically.
+
+## Design notes
+
+- JAX-first: RNG keys and immutable env state are passed explicitly so environments can be
+	stepped inside JIT/pmapped loops if desired.
+- Small, composable environment variants are provided through the registry (easy to add more).
+
+## Examples
+
+- `examples/plot.py` — runs a short random policy in `ForagaxWeather-v1` and produces a
+	temperature vs reward plot (saves to `plots/sample_plot.png`).
+- `examples/visualize.py` — runs environments at multiple aperture sizes and saves short
+	videos under `videos/` using `save_video`.
+
+## Development
+
+Run unit tests via pytest.

{continual_foragax-0.7.0 → continual_foragax-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "continual-foragax"
-version = "0.7.0"
+version = "0.8.0"
 description = "A continual reinforcement learning benchmark"
 readme = "README.md"
 authors = [
@@ -30,6 +30,12 @@ build-backend = "setuptools.build_meta"
 [tool]
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.7.0"
+version = "0.8.0"
 tag_format = "$version"
 version_files = ["pyproject.toml"]
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+"foragax" = ["../data/ECA_non-blended_custom/*"]

continual_foragax-0.8.0/src/continual_foragax.egg-info/PKG-INFO

@@ -0,0 +1,121 @@
+Metadata-Version: 2.4
+Name: continual-foragax
+Version: 0.8.0
+Summary: A continual reinforcement learning benchmark
+Author-email: Steven Tang <stang5@ualberta.ca>
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: gymnax
+Requires-Dist: six; python_version < "3.10"
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-benchmark; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+
+# foragax
+
+Foragax is a lightweight, JAX-first grid-world environment suite for continual / procedural
+experiments. It provides a small collection of environment variants (weather, multi-biome,
+etc.), a registry factory for easy construction, and simple example scripts for plotting and
+visualization.
+
+This version is a [Gymnax](https://github.com/RobertTLange/gymnax) environment implemented in JAX. The original implementation of Forager is implemented in Numba is available at [andnp/forager](https://github.com/andnp/forager). In addition to the original features, this implementation includes: biomes, visualization, and a weather environment.
+
+Key ideas:
+
+- Functional, JAX-friendly API (explicit PRNG keys, immutable env state objects).
+- Multiple observation modalities: Object and RGB, as well as aperture based or
+	full-world observations.
+- Customizable biomes
+- Customizable object placement, respawning, and rewards.
+- Visualization via RGB rendering and plotting.
+
+## Quickstart
+
+We recommend installing with pip from https://pypi.org/project/continual-foragax/.
+
+```bash
+pip install continual-foragax
+```
+
+We support Python 3.8 through Python 3.13.
+
+The codebase expects JAX and other numeric dependencies. If you don't have JAX installed, see
+the JAX install instructions for your platform; the project `uv.lock` pins compatible versions.
+
+## Minimal example (from examples)
+
+Use the registry factory to create an environment and run it with JAX-style RNG keys and an
+explicit environment state.
+
+```python
+from foragax.registry import make
+import jax
+
+# create env (observation_type is one of: 'object', 'rgb', 'world')
+env = make(
+		"ForagaxWeather-v1",
+		aperture_size=5,
+		observation_type="object",
+)
+
+# environment parameters and RNG
+env_params = env.default_params
+key = jax.random.key(0)
+key, key_reset = jax.random.split(key)
+
+# reset returns (obs, env_state)
+_, env_state = env.reset(key_reset, env_params)
+
+# sampling an action and stepping (functional-style)
+key, key_act, key_step = jax.random.split(key, 3)
+action = env.action_space(env_params).sample(key_act)
+_, next_env_state, reward, done, info = env.step(key_step, env_state, action, env_params)
+
+# rendering supports multiple modes: 'world' and 'aperture'
+frame = env.render(env_state, env_params, render_mode="aperture")
+```
+
+See `examples/plot.py` and `examples/visualize.py` for runnable scripts that produce a sample
+plot and saved videos using Gym/Gymnasium helpers.
+
+## Registry and included environments
+
+Use `foragax.registry.make` to construct environments by id. Example environment ids include:
+
+- `ForagaxTwoBiomeSmall-v1` / `-v2` — hand-crafted small multi-biome layouts
+- `ForagaxWeather-v1` — small weather-driven two-biome environment used by examples
+
+The `make` factory accepts the following notable kwargs:
+
+- `observation_type`: one of `"object"`, `"rgb"`, or `"world"`.
+- `aperture_size`: integer or tuple controlling the agent's local observation aperture.
+- `file_index`: used to pick weather locations.
+
+## Custom objects and extensions
+
+The codebase includes an object system for placing items into biomes and controlling
+behaviour (rewards, respawn / regen behavior, blocking/collectable flags). See
+`foragax.objects` for the canonical object definitions and helpers like
+`create_weather_objects` used by the registry.
+
+If you want to add new object classes, follow the examples in `foragax.objects` and add the
+class into registry configs or construct environments programmatically.
+
+## Design notes
+
+- JAX-first: RNG keys and immutable env state are passed explicitly so environments can be
+	stepped inside JIT/pmapped loops if desired.
+- Small, composable environment variants are provided through the registry (easy to add more).
+
+## Examples
+
+- `examples/plot.py` — runs a short random policy in `ForagaxWeather-v1` and produces a
+	temperature vs reward plot (saves to `plots/sample_plot.png`).
+- `examples/visualize.py` — runs environments at multiple aperture sizes and saves short
+	videos under `videos/` using `save_video`.
+
+## Development
+
+Run unit tests via pytest.

{continual_foragax-0.7.0 → continual_foragax-0.8.0}/src/continual_foragax.egg-info/SOURCES.txt

@@ -10,4 +10,5 @@ src/foragax/__init__.py
 src/foragax/env.py
 src/foragax/objects.py
 src/foragax/registry.py
+src/foragax/weather.py
 tests/test_foragax.py

{continual_foragax-0.7.0 → continual_foragax-0.8.0}/src/foragax/env.py

@@ -13,7 +13,8 @@ import jax.numpy as jnp
 from flax import struct
 from gymnax.environments import environment, spaces
 
-from foragax.objects import AGENT, EMPTY, BaseForagaxObject
+from foragax.objects import AGENT, EMPTY, BaseForagaxObject, WeatherObject
+from foragax.weather import get_temperature
 
 
 class Actions(IntEnum):
@@ -71,8 +72,13 @@ class ForagaxEnv(environment.Environment):
         if isinstance(aperture_size, int):
             aperture_size = (aperture_size, aperture_size)
         self.aperture_size = aperture_size
-
         objects = (EMPTY,) + objects
+        self.objects = objects
+        self.weather_object = None
+        for o in objects:
+            if isinstance(o, WeatherObject):
+                self.weather_object = o
+                break
 
         # JIT-compatible versions of object and biome properties
         self.object_ids = jnp.arange(len(objects))
@@ -149,6 +155,12 @@ class ForagaxEnv(environment.Environment):
         new_val_at_pos = jax.lax.select(is_collectable, encoded_timer, val_at_pos)
         object_grid = object_grid.at[pos[1], pos[0]].set(new_val_at_pos)
 
+        info = {"discount": self.discount(state, params)}
+        if self.weather_object is not None:
+            info["temperature"] = get_temperature(
+                self.weather_object.rewards, state.time, self.weather_object.repeat
+            )
+
         # 4. UPDATE STATE
         state = EnvState(
             pos=pos,
@@ -162,7 +174,7 @@ class ForagaxEnv(environment.Environment):
             jax.lax.stop_gradient(state),
             reward,
             done,
-            {"discount": self.discount(state, params)},
+            info,
         )
 
     def reset_env(

{continual_foragax-0.7.0 → continual_foragax-0.8.0}/src/foragax/objects.py

@@ -4,6 +4,8 @@ from typing import Tuple
 import jax
 import jax.numpy as jnp
 
+from foragax.weather import FILE_PATHS, get_temperature, load_data
+
 
 class BaseForagaxObject:
     """Base class for objects in the Foragax environment."""
@@ -85,6 +87,35 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
         return jnp.maximum(0, delay).astype(int)
 
 
+class WeatherObject(NormalRegenForagaxObject):
+    """Object with reward based on temperature data."""
+
+    def __init__(
+        self,
+        name: str,
+        rewards: jnp.ndarray,
+        repeat: int,
+        multiplier: float = 1.0,
+        mean_regen_delay: int = 10,
+        std_regen_delay: int = 1,
+        color: Tuple[int, int, int] = (0, 0, 0),
+    ):
+        super().__init__(
+            name=name,
+            collectable=True,
+            mean_regen_delay=mean_regen_delay,
+            std_regen_delay=std_regen_delay,
+            color=color,
+        )
+        self.rewards = rewards
+        self.repeat = repeat
+        self.multiplier = multiplier
+
+    def reward(self, clock: int, rng: jax.Array) -> float:
+        """Reward is based on temperature."""
+        return get_temperature(self.rewards, clock, self.repeat) * self.multiplier
+
+
 EMPTY = DefaultForagaxObject()
 WALL = DefaultForagaxObject(name="wall", blocking=True, color=(127, 127, 127))
 FLOWER = DefaultForagaxObject(
@@ -147,3 +178,43 @@ DEATHCAP = DefaultForagaxObject(
     color=(193, 178, 30),
 )
 AGENT = DefaultForagaxObject(name="agent", blocking=True, color=(0, 0, 255))
+
+
+def create_weather_objects(
+    file_index: int = 0, repeat: int = 100, multiplier: float = 1.0
+):
+    """Create HOT and COLD WeatherObject instances using the specified file.
+
+    Args:
+        file_index: Index into `FILE_PATHS` to select the temperature file.
+        repeat: How many steps each temperature value repeats for.
+        multiplier: Base multiplier applied to HOT; COLD will use -multiplier.
+
+    Returns:
+        A tuple (HOT, COLD) of WeatherObject instances.
+    """
+    # Clamp file_index
+    if file_index < 0 or file_index >= len(FILE_PATHS):
+        raise IndexError(
+            f"file_index {file_index} out of range (0..{len(FILE_PATHS) - 1})"
+        )
+
+    rewards = load_data(FILE_PATHS[file_index])
+
+    hot = WeatherObject(
+        name="hot",
+        rewards=rewards,
+        repeat=repeat,
+        multiplier=multiplier,
+        color=(255, 0, 255),
+    )
+
+    cold = WeatherObject(
+        name="cold",
+        rewards=rewards,
+        repeat=repeat,
+        multiplier=-multiplier,
+        color=(0, 255, 255),
+    )
+
+    return hot, cold

{continual_foragax-0.7.0 → continual_foragax-0.8.0}/src/foragax/registry.py

@@ -9,12 +9,28 @@ from foragax.env import (
     ForagaxRGBEnv,
     ForagaxWorldEnv,
 )
-from foragax.objects import LARGE_MOREL, LARGE_OYSTER, MEDIUM_MOREL
+from foragax.objects import (
+    LARGE_MOREL,
+    LARGE_OYSTER,
+    MEDIUM_MOREL,
+    create_weather_objects,
+)
 
 ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
+    "ForagaxWeather-v1": {
+        "size": (15, 15),
+        "aperture_size": None,
+        "objects": None,
+        "biomes": (
+            # Hot biome
+            Biome(start=(0, 3), stop=(15, 5), object_frequencies=(0.5, 0.0)),
+            # Cold biome
+            Biome(start=(0, 10), stop=(15, 12), object_frequencies=(0.0, 0.5)),
+        ),
+    },
     "ForagaxTwoBiomeSmall-v1": {
         "size": (16, 8),
-        "aperture_size": (5, 5),
+        "aperture_size": None,
         "objects": (LARGE_MOREL, LARGE_OYSTER),
         "biomes": (
             # Morel biome
@@ -25,7 +41,7 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
     },
     "ForagaxTwoBiomeSmall-v2": {
         "size": (16, 8),
-        "aperture_size": (5, 5),
+        "aperture_size": None,
         "objects": (MEDIUM_MOREL, LARGE_OYSTER),
         "biomes": (
             # Morel biome
@@ -40,7 +56,8 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
 def make(
     env_id: str,
     observation_type: str = "object",
-    aperture_size: Optional[Tuple[int, int]] = None,
+    aperture_size: Optional[Tuple[int, int]] = (5, 5),
+    file_index: int = 0,
 ) -> ForagaxEnv:
     """Create a Foragax environment.
 
@@ -58,8 +75,11 @@ def make(
 
     config = ENV_CONFIGS[env_id].copy()
 
-    if aperture_size is not None:
-        config["aperture_size"] = aperture_size
+    config["aperture_size"] = aperture_size
+
+    if env_id.startswith("ForagaxWeather"):
+        hot, cold = create_weather_objects(file_index=file_index)
+        config["objects"] = (hot, cold)
 
     env_class_map = {
         "object": ForagaxObjectEnv,

continual_foragax-0.8.0/src/foragax/weather.py

@@ -0,0 +1,37 @@
+"""Helper functions for loading temperature data for the ForagerWeather env."""
+
+import csv
+from glob import glob
+
+import jax.numpy as jnp
+import numpy as np
+
+DATA_PATH = "data/ECA_non-blended_custom"
+FILE_PATHS = sorted(glob(f"{DATA_PATH}/TG_*.txt"))
+
+
+def load_data(file_path: str):
+    """Load and normalize temperature data from a file"""
+    with open(file_path, "r") as f:
+        # Skip header lines
+        for _ in range(21):
+            next(f)
+        reader = csv.DictReader(f)
+        if reader.fieldnames:
+            reader.fieldnames = [fn.strip() for fn in reader.fieldnames]
+        tg = []
+        for row in reader:
+            if not row or row["Q_TG"].strip() != "0":
+                continue
+            tg.append(float(row["TG"]))
+        tg = np.array(tg)
+        mean_temperature = tg / 10
+        min_temp = mean_temperature.min()
+        max_temp = mean_temperature.max()
+        normalized = (mean_temperature - min_temp) / (max_temp - min_temp) * 2 - 1
+        return jnp.array(normalized)
+
+
+def get_temperature(rewards: jnp.ndarray, clock: int, repeat: int) -> float:
+    """Get the temperature for a given clock time."""
+    return rewards[clock // repeat % len(rewards)]

continual_foragax-0.7.0/PKG-INFO

@@ -1,16 +0,0 @@
-Metadata-Version: 2.4
-Name: continual-foragax
-Version: 0.7.0
-Summary: A continual reinforcement learning benchmark
-Author-email: Steven Tang <stang5@ualberta.ca>
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-Requires-Dist: gymnax
-Requires-Dist: six; python_version < "3.10"
-Provides-Extra: dev
-Requires-Dist: pre-commit; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: pytest-benchmark; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
-
-# foragax

continual_foragax-0.7.0/README.md

@@ -1 +0,0 @@
-# foragax

continual_foragax-0.7.0/src/continual_foragax.egg-info/PKG-INFO

@@ -1,16 +0,0 @@
-Metadata-Version: 2.4
-Name: continual-foragax
-Version: 0.7.0
-Summary: A continual reinforcement learning benchmark
-Author-email: Steven Tang <stang5@ualberta.ca>
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-Requires-Dist: gymnax
-Requires-Dist: six; python_version < "3.10"
-Provides-Extra: dev
-Requires-Dist: pre-commit; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: pytest-benchmark; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
-
-# foragax