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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (20) hide show
  1. continual_foragax-0.8.0/PKG-INFO +121 -0
  2. continual_foragax-0.8.0/README.md +106 -0
  3. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/pyproject.toml +14 -8
  4. continual_foragax-0.8.0/src/continual_foragax.egg-info/PKG-INFO +121 -0
  5. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/continual_foragax.egg-info/SOURCES.txt +1 -0
  6. continual_foragax-0.8.0/src/continual_foragax.egg-info/requires.txt +10 -0
  7. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/foragax/env.py +18 -6
  8. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/foragax/objects.py +71 -0
  9. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/foragax/registry.py +26 -6
  10. continual_foragax-0.8.0/src/foragax/weather.py +37 -0
  11. continual_foragax-0.6.0/PKG-INFO +0 -16
  12. continual_foragax-0.6.0/README.md +0 -1
  13. continual_foragax-0.6.0/src/continual_foragax.egg-info/PKG-INFO +0 -16
  14. continual_foragax-0.6.0/src/continual_foragax.egg-info/requires.txt +0 -10
  15. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/setup.cfg +0 -0
  16. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/continual_foragax.egg-info/dependency_links.txt +0 -0
  17. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/continual_foragax.egg-info/entry_points.txt +0 -0
  18. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/continual_foragax.egg-info/top_level.txt +0 -0
  19. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/src/foragax/__init__.py +0 -0
  20. {continual_foragax-0.6.0 → continual_foragax-0.8.0}/tests/test_foragax.py +0 -0
@@ -0,0 +1,121 @@
1
+ Metadata-Version: 2.4
2
+ Name: continual-foragax
3
+ Version: 0.8.0
4
+ Summary: A continual reinforcement learning benchmark
5
+ Author-email: Steven Tang <stang5@ualberta.ca>
6
+ Requires-Python: >=3.8
7
+ Description-Content-Type: text/markdown
8
+ Requires-Dist: gymnax
9
+ Requires-Dist: six; python_version < "3.10"
10
+ Provides-Extra: dev
11
+ Requires-Dist: pre-commit; extra == "dev"
12
+ Requires-Dist: pytest; extra == "dev"
13
+ Requires-Dist: pytest-benchmark; extra == "dev"
14
+ Requires-Dist: ruff; extra == "dev"
15
+
16
+ # foragax
17
+
18
+ Foragax is a lightweight, JAX-first grid-world environment suite for continual / procedural
19
+ experiments. It provides a small collection of environment variants (weather, multi-biome,
20
+ etc.), a registry factory for easy construction, and simple example scripts for plotting and
21
+ visualization.
22
+
23
+ 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.
24
+
25
+ Key ideas:
26
+
27
+ - Functional, JAX-friendly API (explicit PRNG keys, immutable env state objects).
28
+ - Multiple observation modalities: Object and RGB, as well as aperture based or
29
+ full-world observations.
30
+ - Customizable biomes
31
+ - Customizable object placement, respawning, and rewards.
32
+ - Visualization via RGB rendering and plotting.
33
+
34
+ ## Quickstart
35
+
36
+ We recommend installing with pip from https://pypi.org/project/continual-foragax/.
37
+
38
+ ```bash
39
+ pip install continual-foragax
40
+ ```
41
+
42
+ We support Python 3.8 through Python 3.13.
43
+
44
+ The codebase expects JAX and other numeric dependencies. If you don't have JAX installed, see
45
+ the JAX install instructions for your platform; the project `uv.lock` pins compatible versions.
46
+
47
+ ## Minimal example (from examples)
48
+
49
+ Use the registry factory to create an environment and run it with JAX-style RNG keys and an
50
+ explicit environment state.
51
+
52
+ ```python
53
+ from foragax.registry import make
54
+ import jax
55
+
56
+ # create env (observation_type is one of: 'object', 'rgb', 'world')
57
+ env = make(
58
+ "ForagaxWeather-v1",
59
+ aperture_size=5,
60
+ observation_type="object",
61
+ )
62
+
63
+ # environment parameters and RNG
64
+ env_params = env.default_params
65
+ key = jax.random.key(0)
66
+ key, key_reset = jax.random.split(key)
67
+
68
+ # reset returns (obs, env_state)
69
+ _, env_state = env.reset(key_reset, env_params)
70
+
71
+ # sampling an action and stepping (functional-style)
72
+ key, key_act, key_step = jax.random.split(key, 3)
73
+ action = env.action_space(env_params).sample(key_act)
74
+ _, next_env_state, reward, done, info = env.step(key_step, env_state, action, env_params)
75
+
76
+ # rendering supports multiple modes: 'world' and 'aperture'
77
+ frame = env.render(env_state, env_params, render_mode="aperture")
78
+ ```
79
+
80
+ See `examples/plot.py` and `examples/visualize.py` for runnable scripts that produce a sample
81
+ plot and saved videos using Gym/Gymnasium helpers.
82
+
83
+ ## Registry and included environments
84
+
85
+ Use `foragax.registry.make` to construct environments by id. Example environment ids include:
86
+
87
+ - `ForagaxTwoBiomeSmall-v1` / `-v2` — hand-crafted small multi-biome layouts
88
+ - `ForagaxWeather-v1` — small weather-driven two-biome environment used by examples
89
+
90
+ The `make` factory accepts the following notable kwargs:
91
+
92
+ - `observation_type`: one of `"object"`, `"rgb"`, or `"world"`.
93
+ - `aperture_size`: integer or tuple controlling the agent's local observation aperture.
94
+ - `file_index`: used to pick weather locations.
95
+
96
+ ## Custom objects and extensions
97
+
98
+ The codebase includes an object system for placing items into biomes and controlling
99
+ behaviour (rewards, respawn / regen behavior, blocking/collectable flags). See
100
+ `foragax.objects` for the canonical object definitions and helpers like
101
+ `create_weather_objects` used by the registry.
102
+
103
+ If you want to add new object classes, follow the examples in `foragax.objects` and add the
104
+ class into registry configs or construct environments programmatically.
105
+
106
+ ## Design notes
107
+
108
+ - JAX-first: RNG keys and immutable env state are passed explicitly so environments can be
109
+ stepped inside JIT/pmapped loops if desired.
110
+ - Small, composable environment variants are provided through the registry (easy to add more).
111
+
112
+ ## Examples
113
+
114
+ - `examples/plot.py` — runs a short random policy in `ForagaxWeather-v1` and produces a
115
+ temperature vs reward plot (saves to `plots/sample_plot.png`).
116
+ - `examples/visualize.py` — runs environments at multiple aperture sizes and saves short
117
+ videos under `videos/` using `save_video`.
118
+
119
+ ## Development
120
+
121
+ Run unit tests via pytest.
@@ -0,0 +1,106 @@
1
+ # foragax
2
+
3
+ Foragax is a lightweight, JAX-first grid-world environment suite for continual / procedural
4
+ experiments. It provides a small collection of environment variants (weather, multi-biome,
5
+ etc.), a registry factory for easy construction, and simple example scripts for plotting and
6
+ visualization.
7
+
8
+ 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.
9
+
10
+ Key ideas:
11
+
12
+ - Functional, JAX-friendly API (explicit PRNG keys, immutable env state objects).
13
+ - Multiple observation modalities: Object and RGB, as well as aperture based or
14
+ full-world observations.
15
+ - Customizable biomes
16
+ - Customizable object placement, respawning, and rewards.
17
+ - Visualization via RGB rendering and plotting.
18
+
19
+ ## Quickstart
20
+
21
+ We recommend installing with pip from https://pypi.org/project/continual-foragax/.
22
+
23
+ ```bash
24
+ pip install continual-foragax
25
+ ```
26
+
27
+ We support Python 3.8 through Python 3.13.
28
+
29
+ The codebase expects JAX and other numeric dependencies. If you don't have JAX installed, see
30
+ the JAX install instructions for your platform; the project `uv.lock` pins compatible versions.
31
+
32
+ ## Minimal example (from examples)
33
+
34
+ Use the registry factory to create an environment and run it with JAX-style RNG keys and an
35
+ explicit environment state.
36
+
37
+ ```python
38
+ from foragax.registry import make
39
+ import jax
40
+
41
+ # create env (observation_type is one of: 'object', 'rgb', 'world')
42
+ env = make(
43
+ "ForagaxWeather-v1",
44
+ aperture_size=5,
45
+ observation_type="object",
46
+ )
47
+
48
+ # environment parameters and RNG
49
+ env_params = env.default_params
50
+ key = jax.random.key(0)
51
+ key, key_reset = jax.random.split(key)
52
+
53
+ # reset returns (obs, env_state)
54
+ _, env_state = env.reset(key_reset, env_params)
55
+
56
+ # sampling an action and stepping (functional-style)
57
+ key, key_act, key_step = jax.random.split(key, 3)
58
+ action = env.action_space(env_params).sample(key_act)
59
+ _, next_env_state, reward, done, info = env.step(key_step, env_state, action, env_params)
60
+
61
+ # rendering supports multiple modes: 'world' and 'aperture'
62
+ frame = env.render(env_state, env_params, render_mode="aperture")
63
+ ```
64
+
65
+ See `examples/plot.py` and `examples/visualize.py` for runnable scripts that produce a sample
66
+ plot and saved videos using Gym/Gymnasium helpers.
67
+
68
+ ## Registry and included environments
69
+
70
+ Use `foragax.registry.make` to construct environments by id. Example environment ids include:
71
+
72
+ - `ForagaxTwoBiomeSmall-v1` / `-v2` — hand-crafted small multi-biome layouts
73
+ - `ForagaxWeather-v1` — small weather-driven two-biome environment used by examples
74
+
75
+ The `make` factory accepts the following notable kwargs:
76
+
77
+ - `observation_type`: one of `"object"`, `"rgb"`, or `"world"`.
78
+ - `aperture_size`: integer or tuple controlling the agent's local observation aperture.
79
+ - `file_index`: used to pick weather locations.
80
+
81
+ ## Custom objects and extensions
82
+
83
+ The codebase includes an object system for placing items into biomes and controlling
84
+ behaviour (rewards, respawn / regen behavior, blocking/collectable flags). See
85
+ `foragax.objects` for the canonical object definitions and helpers like
86
+ `create_weather_objects` used by the registry.
87
+
88
+ If you want to add new object classes, follow the examples in `foragax.objects` and add the
89
+ class into registry configs or construct environments programmatically.
90
+
91
+ ## Design notes
92
+
93
+ - JAX-first: RNG keys and immutable env state are passed explicitly so environments can be
94
+ stepped inside JIT/pmapped loops if desired.
95
+ - Small, composable environment variants are provided through the registry (easy to add more).
96
+
97
+ ## Examples
98
+
99
+ - `examples/plot.py` — runs a short random policy in `ForagaxWeather-v1` and produces a
100
+ temperature vs reward plot (saves to `plots/sample_plot.png`).
101
+ - `examples/visualize.py` — runs environments at multiple aperture sizes and saves short
102
+ videos under `videos/` using `save_video`.
103
+
104
+ ## Development
105
+
106
+ Run unit tests via pytest.
@@ -1,23 +1,23 @@
1
1
  [project]
2
2
  name = "continual-foragax"
3
- version = "0.6.0"
3
+ version = "0.8.0"
4
4
  description = "A continual reinforcement learning benchmark"
5
5
  readme = "README.md"
6
6
  authors = [
7
7
  { name = "Steven Tang", email = "stang5@ualberta.ca" }
8
8
  ]
9
- requires-python = ">=3.9"
9
+ requires-python = ">=3.8"
10
10
  dependencies = [
11
11
  "gymnax",
12
- "six; python_version == '3.9'",
12
+ "six; python_version < '3.10'",
13
13
  ]
14
14
 
15
15
  [project.optional-dependencies]
16
16
  dev = [
17
- "pre-commit>=4.3.0",
18
- "pytest>=8.4.1",
19
- "pytest-benchmark>=5.1.0",
20
- "ruff>=0.12.9",
17
+ "pre-commit",
18
+ "pytest",
19
+ "pytest-benchmark",
20
+ "ruff",
21
21
  ]
22
22
 
23
23
  [project.scripts]
@@ -30,6 +30,12 @@ build-backend = "setuptools.build_meta"
30
30
  [tool]
31
31
  [tool.commitizen]
32
32
  name = "cz_conventional_commits"
33
- version = "0.6.0"
33
+ version = "0.8.0"
34
34
  tag_format = "$version"
35
35
  version_files = ["pyproject.toml"]
36
+
37
+ [tool.setuptools]
38
+ include-package-data = true
39
+
40
+ [tool.setuptools.package-data]
41
+ "foragax" = ["../data/ECA_non-blended_custom/*"]
@@ -0,0 +1,121 @@
1
+ Metadata-Version: 2.4
2
+ Name: continual-foragax
3
+ Version: 0.8.0
4
+ Summary: A continual reinforcement learning benchmark
5
+ Author-email: Steven Tang <stang5@ualberta.ca>
6
+ Requires-Python: >=3.8
7
+ Description-Content-Type: text/markdown
8
+ Requires-Dist: gymnax
9
+ Requires-Dist: six; python_version < "3.10"
10
+ Provides-Extra: dev
11
+ Requires-Dist: pre-commit; extra == "dev"
12
+ Requires-Dist: pytest; extra == "dev"
13
+ Requires-Dist: pytest-benchmark; extra == "dev"
14
+ Requires-Dist: ruff; extra == "dev"
15
+
16
+ # foragax
17
+
18
+ Foragax is a lightweight, JAX-first grid-world environment suite for continual / procedural
19
+ experiments. It provides a small collection of environment variants (weather, multi-biome,
20
+ etc.), a registry factory for easy construction, and simple example scripts for plotting and
21
+ visualization.
22
+
23
+ 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.
24
+
25
+ Key ideas:
26
+
27
+ - Functional, JAX-friendly API (explicit PRNG keys, immutable env state objects).
28
+ - Multiple observation modalities: Object and RGB, as well as aperture based or
29
+ full-world observations.
30
+ - Customizable biomes
31
+ - Customizable object placement, respawning, and rewards.
32
+ - Visualization via RGB rendering and plotting.
33
+
34
+ ## Quickstart
35
+
36
+ We recommend installing with pip from https://pypi.org/project/continual-foragax/.
37
+
38
+ ```bash
39
+ pip install continual-foragax
40
+ ```
41
+
42
+ We support Python 3.8 through Python 3.13.
43
+
44
+ The codebase expects JAX and other numeric dependencies. If you don't have JAX installed, see
45
+ the JAX install instructions for your platform; the project `uv.lock` pins compatible versions.
46
+
47
+ ## Minimal example (from examples)
48
+
49
+ Use the registry factory to create an environment and run it with JAX-style RNG keys and an
50
+ explicit environment state.
51
+
52
+ ```python
53
+ from foragax.registry import make
54
+ import jax
55
+
56
+ # create env (observation_type is one of: 'object', 'rgb', 'world')
57
+ env = make(
58
+ "ForagaxWeather-v1",
59
+ aperture_size=5,
60
+ observation_type="object",
61
+ )
62
+
63
+ # environment parameters and RNG
64
+ env_params = env.default_params
65
+ key = jax.random.key(0)
66
+ key, key_reset = jax.random.split(key)
67
+
68
+ # reset returns (obs, env_state)
69
+ _, env_state = env.reset(key_reset, env_params)
70
+
71
+ # sampling an action and stepping (functional-style)
72
+ key, key_act, key_step = jax.random.split(key, 3)
73
+ action = env.action_space(env_params).sample(key_act)
74
+ _, next_env_state, reward, done, info = env.step(key_step, env_state, action, env_params)
75
+
76
+ # rendering supports multiple modes: 'world' and 'aperture'
77
+ frame = env.render(env_state, env_params, render_mode="aperture")
78
+ ```
79
+
80
+ See `examples/plot.py` and `examples/visualize.py` for runnable scripts that produce a sample
81
+ plot and saved videos using Gym/Gymnasium helpers.
82
+
83
+ ## Registry and included environments
84
+
85
+ Use `foragax.registry.make` to construct environments by id. Example environment ids include:
86
+
87
+ - `ForagaxTwoBiomeSmall-v1` / `-v2` — hand-crafted small multi-biome layouts
88
+ - `ForagaxWeather-v1` — small weather-driven two-biome environment used by examples
89
+
90
+ The `make` factory accepts the following notable kwargs:
91
+
92
+ - `observation_type`: one of `"object"`, `"rgb"`, or `"world"`.
93
+ - `aperture_size`: integer or tuple controlling the agent's local observation aperture.
94
+ - `file_index`: used to pick weather locations.
95
+
96
+ ## Custom objects and extensions
97
+
98
+ The codebase includes an object system for placing items into biomes and controlling
99
+ behaviour (rewards, respawn / regen behavior, blocking/collectable flags). See
100
+ `foragax.objects` for the canonical object definitions and helpers like
101
+ `create_weather_objects` used by the registry.
102
+
103
+ If you want to add new object classes, follow the examples in `foragax.objects` and add the
104
+ class into registry configs or construct environments programmatically.
105
+
106
+ ## Design notes
107
+
108
+ - JAX-first: RNG keys and immutable env state are passed explicitly so environments can be
109
+ stepped inside JIT/pmapped loops if desired.
110
+ - Small, composable environment variants are provided through the registry (easy to add more).
111
+
112
+ ## Examples
113
+
114
+ - `examples/plot.py` — runs a short random policy in `ForagaxWeather-v1` and produces a
115
+ temperature vs reward plot (saves to `plots/sample_plot.png`).
116
+ - `examples/visualize.py` — runs environments at multiple aperture sizes and saves short
117
+ videos under `videos/` using `save_video`.
118
+
119
+ ## Development
120
+
121
+ Run unit tests via pytest.
@@ -10,4 +10,5 @@ src/foragax/__init__.py
10
10
  src/foragax/env.py
11
11
  src/foragax/objects.py
12
12
  src/foragax/registry.py
13
+ src/foragax/weather.py
13
14
  tests/test_foragax.py
@@ -0,0 +1,10 @@
1
+ gymnax
2
+
3
+ [:python_version < "3.10"]
4
+ six
5
+
6
+ [dev]
7
+ pre-commit
8
+ pytest
9
+ pytest-benchmark
10
+ ruff
@@ -6,14 +6,15 @@ Source: https://github.com/andnp/Foragax
6
6
  from dataclasses import dataclass
7
7
  from enum import IntEnum
8
8
  from functools import partial
9
- from typing import Any, Tuple, Union
9
+ from typing import Any, Dict, Tuple, Union
10
10
 
11
11
  import jax
12
12
  import jax.numpy as jnp
13
13
  from flax import struct
14
14
  from gymnax.environments import environment, spaces
15
15
 
16
- from foragax.objects import AGENT, EMPTY, BaseForagaxObject
16
+ from foragax.objects import AGENT, EMPTY, BaseForagaxObject, WeatherObject
17
+ from foragax.weather import get_temperature
17
18
 
18
19
 
19
20
  class Actions(IntEnum):
@@ -71,8 +72,13 @@ class ForagaxEnv(environment.Environment):
71
72
  if isinstance(aperture_size, int):
72
73
  aperture_size = (aperture_size, aperture_size)
73
74
  self.aperture_size = aperture_size
74
-
75
75
  objects = (EMPTY,) + objects
76
+ self.objects = objects
77
+ self.weather_object = None
78
+ for o in objects:
79
+ if isinstance(o, WeatherObject):
80
+ self.weather_object = o
81
+ break
76
82
 
77
83
  # JIT-compatible versions of object and biome properties
78
84
  self.object_ids = jnp.arange(len(objects))
@@ -105,7 +111,7 @@ class ForagaxEnv(environment.Environment):
105
111
  state: EnvState,
106
112
  action: Union[int, float, jax.Array],
107
113
  params: EnvParams,
108
- ) -> tuple[jax.Array, EnvState, jax.Array, jax.Array, dict[Any, Any]]:
114
+ ) -> Tuple[jax.Array, EnvState, jax.Array, jax.Array, Dict[Any, Any]]:
109
115
  """Perform single timestep state transition."""
110
116
  num_obj_types = len(self.object_ids)
111
117
  # Decode the object grid: positive values are objects, negative are timers (treat as empty)
@@ -149,6 +155,12 @@ class ForagaxEnv(environment.Environment):
149
155
  new_val_at_pos = jax.lax.select(is_collectable, encoded_timer, val_at_pos)
150
156
  object_grid = object_grid.at[pos[1], pos[0]].set(new_val_at_pos)
151
157
 
158
+ info = {"discount": self.discount(state, params)}
159
+ if self.weather_object is not None:
160
+ info["temperature"] = get_temperature(
161
+ self.weather_object.rewards, state.time, self.weather_object.repeat
162
+ )
163
+
152
164
  # 4. UPDATE STATE
153
165
  state = EnvState(
154
166
  pos=pos,
@@ -162,12 +174,12 @@ class ForagaxEnv(environment.Environment):
162
174
  jax.lax.stop_gradient(state),
163
175
  reward,
164
176
  done,
165
- {"discount": self.discount(state, params)},
177
+ info,
166
178
  )
167
179
 
168
180
  def reset_env(
169
181
  self, key: jax.Array, params: EnvParams
170
- ) -> tuple[jax.Array, EnvState]:
182
+ ) -> Tuple[jax.Array, EnvState]:
171
183
  """Reset environment state."""
172
184
  key, subkey = jax.random.split(key)
173
185
 
@@ -4,6 +4,8 @@ from typing import Tuple
4
4
  import jax
5
5
  import jax.numpy as jnp
6
6
 
7
+ from foragax.weather import FILE_PATHS, get_temperature, load_data
8
+
7
9
 
8
10
  class BaseForagaxObject:
9
11
  """Base class for objects in the Foragax environment."""
@@ -85,6 +87,35 @@ class NormalRegenForagaxObject(DefaultForagaxObject):
85
87
  return jnp.maximum(0, delay).astype(int)
86
88
 
87
89
 
90
+ class WeatherObject(NormalRegenForagaxObject):
91
+ """Object with reward based on temperature data."""
92
+
93
+ def __init__(
94
+ self,
95
+ name: str,
96
+ rewards: jnp.ndarray,
97
+ repeat: int,
98
+ multiplier: float = 1.0,
99
+ mean_regen_delay: int = 10,
100
+ std_regen_delay: int = 1,
101
+ color: Tuple[int, int, int] = (0, 0, 0),
102
+ ):
103
+ super().__init__(
104
+ name=name,
105
+ collectable=True,
106
+ mean_regen_delay=mean_regen_delay,
107
+ std_regen_delay=std_regen_delay,
108
+ color=color,
109
+ )
110
+ self.rewards = rewards
111
+ self.repeat = repeat
112
+ self.multiplier = multiplier
113
+
114
+ def reward(self, clock: int, rng: jax.Array) -> float:
115
+ """Reward is based on temperature."""
116
+ return get_temperature(self.rewards, clock, self.repeat) * self.multiplier
117
+
118
+
88
119
  EMPTY = DefaultForagaxObject()
89
120
  WALL = DefaultForagaxObject(name="wall", blocking=True, color=(127, 127, 127))
90
121
  FLOWER = DefaultForagaxObject(
@@ -147,3 +178,43 @@ DEATHCAP = DefaultForagaxObject(
147
178
  color=(193, 178, 30),
148
179
  )
149
180
  AGENT = DefaultForagaxObject(name="agent", blocking=True, color=(0, 0, 255))
181
+
182
+
183
+ def create_weather_objects(
184
+ file_index: int = 0, repeat: int = 100, multiplier: float = 1.0
185
+ ):
186
+ """Create HOT and COLD WeatherObject instances using the specified file.
187
+
188
+ Args:
189
+ file_index: Index into `FILE_PATHS` to select the temperature file.
190
+ repeat: How many steps each temperature value repeats for.
191
+ multiplier: Base multiplier applied to HOT; COLD will use -multiplier.
192
+
193
+ Returns:
194
+ A tuple (HOT, COLD) of WeatherObject instances.
195
+ """
196
+ # Clamp file_index
197
+ if file_index < 0 or file_index >= len(FILE_PATHS):
198
+ raise IndexError(
199
+ f"file_index {file_index} out of range (0..{len(FILE_PATHS) - 1})"
200
+ )
201
+
202
+ rewards = load_data(FILE_PATHS[file_index])
203
+
204
+ hot = WeatherObject(
205
+ name="hot",
206
+ rewards=rewards,
207
+ repeat=repeat,
208
+ multiplier=multiplier,
209
+ color=(255, 0, 255),
210
+ )
211
+
212
+ cold = WeatherObject(
213
+ name="cold",
214
+ rewards=rewards,
215
+ repeat=repeat,
216
+ multiplier=-multiplier,
217
+ color=(0, 255, 255),
218
+ )
219
+
220
+ return hot, cold
@@ -9,12 +9,28 @@ from foragax.env import (
9
9
  ForagaxRGBEnv,
10
10
  ForagaxWorldEnv,
11
11
  )
12
- from foragax.objects import LARGE_MOREL, LARGE_OYSTER, MEDIUM_MOREL
12
+ from foragax.objects import (
13
+ LARGE_MOREL,
14
+ LARGE_OYSTER,
15
+ MEDIUM_MOREL,
16
+ create_weather_objects,
17
+ )
13
18
 
14
19
  ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
20
+ "ForagaxWeather-v1": {
21
+ "size": (15, 15),
22
+ "aperture_size": None,
23
+ "objects": None,
24
+ "biomes": (
25
+ # Hot biome
26
+ Biome(start=(0, 3), stop=(15, 5), object_frequencies=(0.5, 0.0)),
27
+ # Cold biome
28
+ Biome(start=(0, 10), stop=(15, 12), object_frequencies=(0.0, 0.5)),
29
+ ),
30
+ },
15
31
  "ForagaxTwoBiomeSmall-v1": {
16
32
  "size": (16, 8),
17
- "aperture_size": (5, 5),
33
+ "aperture_size": None,
18
34
  "objects": (LARGE_MOREL, LARGE_OYSTER),
19
35
  "biomes": (
20
36
  # Morel biome
@@ -25,7 +41,7 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
25
41
  },
26
42
  "ForagaxTwoBiomeSmall-v2": {
27
43
  "size": (16, 8),
28
- "aperture_size": (5, 5),
44
+ "aperture_size": None,
29
45
  "objects": (MEDIUM_MOREL, LARGE_OYSTER),
30
46
  "biomes": (
31
47
  # Morel biome
@@ -40,7 +56,8 @@ ENV_CONFIGS: Dict[str, Dict[str, Any]] = {
40
56
  def make(
41
57
  env_id: str,
42
58
  observation_type: str = "object",
43
- aperture_size: Optional[Tuple[int, int]] = None,
59
+ aperture_size: Optional[Tuple[int, int]] = (5, 5),
60
+ file_index: int = 0,
44
61
  ) -> ForagaxEnv:
45
62
  """Create a Foragax environment.
46
63
 
@@ -58,8 +75,11 @@ def make(
58
75
 
59
76
  config = ENV_CONFIGS[env_id].copy()
60
77
 
61
- if aperture_size is not None:
62
- config["aperture_size"] = aperture_size
78
+ config["aperture_size"] = aperture_size
79
+
80
+ if env_id.startswith("ForagaxWeather"):
81
+ hot, cold = create_weather_objects(file_index=file_index)
82
+ config["objects"] = (hot, cold)
63
83
 
64
84
  env_class_map = {
65
85
  "object": ForagaxObjectEnv,
@@ -0,0 +1,37 @@
1
+ """Helper functions for loading temperature data for the ForagerWeather env."""
2
+
3
+ import csv
4
+ from glob import glob
5
+
6
+ import jax.numpy as jnp
7
+ import numpy as np
8
+
9
+ DATA_PATH = "data/ECA_non-blended_custom"
10
+ FILE_PATHS = sorted(glob(f"{DATA_PATH}/TG_*.txt"))
11
+
12
+
13
+ def load_data(file_path: str):
14
+ """Load and normalize temperature data from a file"""
15
+ with open(file_path, "r") as f:
16
+ # Skip header lines
17
+ for _ in range(21):
18
+ next(f)
19
+ reader = csv.DictReader(f)
20
+ if reader.fieldnames:
21
+ reader.fieldnames = [fn.strip() for fn in reader.fieldnames]
22
+ tg = []
23
+ for row in reader:
24
+ if not row or row["Q_TG"].strip() != "0":
25
+ continue
26
+ tg.append(float(row["TG"]))
27
+ tg = np.array(tg)
28
+ mean_temperature = tg / 10
29
+ min_temp = mean_temperature.min()
30
+ max_temp = mean_temperature.max()
31
+ normalized = (mean_temperature - min_temp) / (max_temp - min_temp) * 2 - 1
32
+ return jnp.array(normalized)
33
+
34
+
35
+ def get_temperature(rewards: jnp.ndarray, clock: int, repeat: int) -> float:
36
+ """Get the temperature for a given clock time."""
37
+ return rewards[clock // repeat % len(rewards)]
@@ -1,16 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: continual-foragax
3
- Version: 0.6.0
4
- Summary: A continual reinforcement learning benchmark
5
- Author-email: Steven Tang <stang5@ualberta.ca>
6
- Requires-Python: >=3.9
7
- Description-Content-Type: text/markdown
8
- Requires-Dist: gymnax
9
- Requires-Dist: six; python_version == "3.9"
10
- Provides-Extra: dev
11
- Requires-Dist: pre-commit>=4.3.0; extra == "dev"
12
- Requires-Dist: pytest>=8.4.1; extra == "dev"
13
- Requires-Dist: pytest-benchmark>=5.1.0; extra == "dev"
14
- Requires-Dist: ruff>=0.12.9; extra == "dev"
15
-
16
- # foragax
@@ -1 +0,0 @@
1
- # foragax
@@ -1,16 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: continual-foragax
3
- Version: 0.6.0
4
- Summary: A continual reinforcement learning benchmark
5
- Author-email: Steven Tang <stang5@ualberta.ca>
6
- Requires-Python: >=3.9
7
- Description-Content-Type: text/markdown
8
- Requires-Dist: gymnax
9
- Requires-Dist: six; python_version == "3.9"
10
- Provides-Extra: dev
11
- Requires-Dist: pre-commit>=4.3.0; extra == "dev"
12
- Requires-Dist: pytest>=8.4.1; extra == "dev"
13
- Requires-Dist: pytest-benchmark>=5.1.0; extra == "dev"
14
- Requires-Dist: ruff>=0.12.9; extra == "dev"
15
-
16
- # foragax
@@ -1,10 +0,0 @@
1
- gymnax
2
-
3
- [:python_version == "3.9"]
4
- six
5
-
6
- [dev]
7
- pre-commit>=4.3.0
8
- pytest>=8.4.1
9
- pytest-benchmark>=5.1.0
10
- ruff>=0.12.9