A reinforcement learning environment written from scratch in JAX. Source on GitHub.

I built foragers world as a way to explore simulated environments in JAX, evolutionary processes and population control RL. The world is designed to be simple by nature (only one control lever, small environment state) to allow a researcher to reason with changes made by the agent. The goal of the RL agent is to control the population of foragers within a given band. Critically, it is important where this band sits, and I found that controlling the population closer to the extinction threshold was much harder for the agent.

Foragers are simple creatures. They breed asexually and randomly walk around their grid environment. They have a measure of energy which is depleted when they take a step and when they breed. Breeding splits their energy between them and their offspring and requires a minimum amount of energy to do.

The RL agent watches that population and gets one lever - a global movement cost. Raising the movement cost causes foragers to burn energy more rapidly and die off. Lowering it allows them to move about and breed more easily. The goal is to hold the population inside a target band for as long as possible, scoring +1 for every step it stays there over a 500–step episode.

The whole simulation environment is pure functional JAX. I utilised fixed size arrays and masks to enable the use of VMAP and parallel simulations. Food is randomly placed throughout the gridworld initially, and regrows stochastically to a per tile cap.

The following three serve as baseline regimes with a fixed movement cost. The first example shows a population collapse where the cost of movement is too high so the population simply collapses. The second shows a relatively stable population that oscillates roughly within the required band. Finally we have a boom-bust population. Movement cost is low so the population expands rapidly, consumes its resources, then starts dying out. The population is able to recover as the foragers can move around cheapily and consume the ever increasing food while the population is dying. These are real trajectories exported from the simulation, Green is food, dark dots are foragers, and the shaded strip on each chart is the target band.

The agent uses PPO, written in JAX in the style of PureJaxRL. Its observation is a three-dimensional vector of global aggregates (total food, population and mean energy) with no spatial features or history. The action space is discrete: 12 move costs from 0.0 to 2.75.

The band is a parameter of the environment, so I trained one policy for each of four bands, for 5M environment steps each.

Based on average return over 256 evaluation worlds, PPO outperformed the best fixed move cost in every target band. In this setup the wide band is essentially saturated: with a band about 60 foragers wide, the agent scores almost perfectly on average (496 of 500).

Each band trains for 5M steps, which takes about two minutes on a M4 Pro Macbook Pro. I used Adam with a learning rate of 2.5e-4 that decays to zero over the run.

It learns fast. Every band beats the best fixed cost within the first million steps (about 20 seconds). The wide, ceiling and floor bands level off by around 3M steps, while the tight band is still climbing when training stops.

These curves sample from the policy instead of taking its best action, so they sit a little below the scores above.

The agent learns a very basic threshold controller. As the policy in each band shows, it switches from a low cost to a high cost across the middle of the population band. Where that switch happens depends on the size and location of the band relative to the total population, food and grid settings.

It changes cost on most steps, around 290–350 times per 500-step episode, so it behaves like a thermostat rather than a smooth controller.

Each band has one training seed and one untuned set of hyperparameters (the PureJaxRL defaults), so nothing here says how reliably PPO finds these controllers. The tight and floor bands are the weakest.

The floor band's 27% extinction rate mostly comes from the start of each episode, when the population grows quickly and the agent then culls it too hard. Being able to observe the rate of change of the population may let the agent do better here, and that is what I will look at next.

Further experimentation into the floor control band, where the agent has to balance the risk of extinction against holding the population in band, is worth investigating.