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

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

continual_foragax-0.56.0.tar.gz (7.7 MB view details)

Uploaded Aug 7, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

continual_foragax-0.56.0-py3-none-any.whl (8.3 MB view details)

Uploaded Aug 7, 2026 Python 3

File details

Details for the file continual_foragax-0.56.0.tar.gz.

File metadata

Download URL: continual_foragax-0.56.0.tar.gz
Upload date: Aug 7, 2026
Size: 7.7 MB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: uv/0.12.2 {"installer":{"name":"uv","version":"0.12.2","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}

File hashes

Hashes for continual_foragax-0.56.0.tar.gz
Algorithm	Hash digest
SHA256	`01b2c545dc458a3118bc87571e843c2d90de8e5980d68955f286db95d5c3eac4`
MD5	`6b532d58d9de0995858ddea4bf192e52`
BLAKE2b-256	`13e5bcd9f55a976a5da189848b47233463c1c151d78883f42cc7eb099254e2ee`

See more details on using hashes here.

File details

Details for the file continual_foragax-0.56.0-py3-none-any.whl.

File metadata

Download URL: continual_foragax-0.56.0-py3-none-any.whl
Upload date: Aug 7, 2026
Size: 8.3 MB
Tags: Python 3
Uploaded using Trusted Publishing? Yes
Uploaded via: uv/0.12.2 {"installer":{"name":"uv","version":"0.12.2","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}

File hashes

Hashes for continual_foragax-0.56.0-py3-none-any.whl
Algorithm	Hash digest
SHA256	`8af2b1efc506116dc4b697600b218dc6da7629310d342f4e5112d3f5c8e8507f`
MD5	`977d708d1dca6f51b11ec84bf64597d6`
BLAKE2b-256	`027adb4545f197d239ce12242bbf4b2a1fb6ca9f88be315b63275c5cb3b88516`