foragax

Foragax is a lightweight, JAX-first grid-world environment suite for continual / procedural experiments. It provides a small collection of environment variants, a registry factory for easy construction, and example scripts for visualization.

This version is a Gymnax environment implemented in JAX. The original implementation of Forager (in Numba) is available at andnp/forager. In addition to the original features, this implementation includes biomes and visualization.

Key ideas:

Functional, JAX-friendly API (explicit PRNG keys, immutable env state objects).
Multiple observation modalities: object, RGB, and color, as well as aperture-based or full-world observations.
Customizable biomes.
Customizable object placement, respawning, and rewards.
Visualization via RGB rendering.

Quickstart

We recommend installing with pip from https://pypi.org/project/continual-foragax/.

pip install continual-foragax

Requires Python 3.8 or newer.

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

Use the registry factory to create an environment and run it with JAX-style RNG keys and an explicit environment state.

from foragax.registry import make
import jax

env = make(
    "ForagaxSquareWaveTwoBiome-v11",
    aperture_size=9,
    observation_type="color",
)

env_params = env.default_params
key = jax.random.key(0)
key, key_reset = jax.random.split(key)
obs, env_state = env.reset(key_reset, env_params)

key, key_act, key_step = jax.random.split(key, 3)
action = env.action_space(env_params).sample(key_act)
obs, env_state, reward, done, info = env.step(key_step, env_state, action, env_params)

frame = env.render(env_state, env_params, render_mode="world")

See examples/observation.py and examples/visualize.py for runnable scripts that save short videos under videos/ using Gymnasium helpers.

Registry and included environments

Use foragax.registry.make to construct environments by id. The registered ids are:

ForagaxBig-v5 — large multi-biome layout with Fourier-modulated rewards (used by examples/visualize.py).
ForagaxSquareWaveTwoBiome-v11 — two-biome layout with square-wave reward shifts (used by examples/observation.py).
ForagaxTwoBiomeLarge-v1 — 15×15 two-biome layout with a Morel biome and an Oyster biome (containing Deathcaps), built from LARGE_MOREL, LARGE_OYSTER, and LARGE_DEATHCAP in foragax.objects.

The make factory accepts the following kwargs:

observation_type: one of "object", "rgb", or "color" (default "color").
aperture_size: int, (int, int), or -1 for full-world observation. Defaults to (5, 5); pass None to use the environment's own default.
reward_delay: steps required to digest food items (default 0).
random_shift_max_steps: random initial offset on the underlying time signal (default 0).
Additional **kwargs are forwarded to the ForagaxEnv constructor and override config defaults.

Custom objects and extensions

Object classes in foragax.objects define rewards, respawn / regen behavior, and blocking/collectable flags. The registry presets above are built using two helpers from this module:

create_fourier_objects — Fourier-modulated reward objects (used by ForagaxBig-v5).
create_shift_square_wave_biome_objects — square-wave biome objects (used by ForagaxSquareWaveTwoBiome-v11).

Weather-driven environments are also supported even though no weather preset is currently registered: compose WeatherObject or WeatherWaveObject from foragax.objects with foragax.weather.get_temperature (which reads ECA&D temperature data shipped under foragax/data/) to build one programmatically.

To add new object classes, follow the patterns in foragax.objects and either register a new entry in foragax.registry.ENV_CONFIGS or construct ForagaxEnv directly.

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/observation.py — runs a random policy in ForagaxSquareWaveTwoBiome-v11 and saves a video of the color observations to videos/.
examples/visualize.py — runs a random policy in ForagaxBig-v5 and saves periodic world_reward render videos to videos/.

Development

Sync dev dependencies and run the test suite:

uv sync --dev
uv run pytest tests

The project uses uv for package management and ruff for formatting and linting.

Citation

If you use Foragax in your research, please cite:

@misc{tang2026forager,
    title={Forager: a lightweight testbed for continual learning with partial observability in RL},
    author={Steven Tang and Xinze Xiong and Anna Hakhverdyan and Andrew Patterson and Jacob Adkins and Jiamin He and Esraa Elelimy and Parham Mohammad Panahi and Martha White and Adam White},
    year={2026},
    eprint={2605.01131},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2605.01131},
}

Acknowledgments

We acknowledge the data providers in the ECA&D project. Klein Tank, A.M.G. and Coauthors, 2002. Daily dataset of 20th-century surface air temperature and precipitation series for the European Climate Assessment. Int. J. of Climatol., 22, 1441-1453.

Data and metadata available at https://www.ecad.eu

Metadata

Release files for continual-foragax 0.57.1

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for continual-foragax 0.57.1
File	Size	Uploaded
continual_foragax-0.57.1.tar.gz	7.7 MB	Aug 20, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for continual-foragax 0.57.1
File	Interpreter	ABI	Platform	Reset
continual_foragax-0.57.1-py3-none-any.whl 8.3 MB Aug 20, 2026	Python 3	none	any	Details

Total release size: 16.0 MB

Release files / continual_foragax-0.57.1.tar.gz

Download URL	continual_foragax-0.57.1.tar.gz
Size	7.7 MB
Tags	Source
SHA-256 checksum How to use checksums	`9014301b2eace78e08e16e6cb893ad5d647ecabbca85aa3e776769d08e5c1ca1`
BLAKE2b-256 checksum How to use checksums	`6978e78c997c5f4106f031bc1e60fef00b68f84f718694566ae8d902e356d27b`
Upload date	Aug 20, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`uv/0.12.5 {"installer":{"name":"uv","version":"0.12.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}`

Release files / continual_foragax-0.57.1-py3-none-any.whl

Download URL	continual_foragax-0.57.1-py3-none-any.whl
Size	8.3 MB
Tags	Python 3
SHA-256 checksum How to use checksums	`1562a782e5e9b79d1b56f5c9bbd776a051df42ea0640314120a32f1d4a3a4946`
BLAKE2b-256 checksum How to use checksums	`3e334378395acdcb232101308e8812269224e681e9aa1095cc7c5df107b8295a`
Upload date	Aug 20, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`uv/0.12.5 {"installer":{"name":"uv","version":"0.12.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}`