A PyTorch implementation of a recurrent actor-critic agent, trained with PPO, on a simulated version of a naturalistic foraging task: three patches ("boxes") that become rewarded at stochastic intervals, a noisy cue on each patch, and a cost for every push. The agent has to decide where to push, when to push, and when to do nothing, from a partial and noisy view of the world. It is the model-side companion to our primate foraging experiments.
Left, middle: the trained agent spaces its pushes to each patch according to that patch's reward interval (dashed: identity); with an unreliable cue (middle) the scaling flattens and the agent waits too long at the fast patch, the same undermatching seen in animals. Right: a ridge readout of true reward availability from the LSTM state; trial-by-trial decoding error correlates with the agent's behavioural error (grey: shuffle control), so the recurrent state carries a belief about each patch.
| Element | Definition |
|---|---|
| Patches | 3, with mean reward intervals drawn as a random permutation of (3, 6, 9) steps per episode |
| Reward timing | After each push on patch i, its next reward becomes available after a Gamma(shape 10, mean mu_i) interval |
| Actions | Push patch 1, 2, 3, or do nothing (action 4) |
| Reward | +1 if the patch was armed, a cost of 0.01 for every push; doing nothing is free |
| Observation | Per-patch elapsed time over interval, clipped to [0, 1], corrupted by multiplicative Gaussian noise with std obs_std drawn per episode from [0, 1); obs_std itself is given as input |
| Agent input | observation (4) + last action + last reward + normalised time (7 dims) |
| Episode | 1000 steps, 500 episodes run in parallel as a batch |
- Two LSTMs (256 units each): one for the actor, one for the critic; linear heads on top.
- PPO with clipped objective (eps 0.2), GAE (gamma 0.9, lambda 1), entropy bonus 0.01, 80 epochs of full-sequence updates per batch, Adam at 3e-4.
- Hidden states are carried across the update so truncated-sequence training sees the right context.
- Orthogonal / Xavier initialisation of recurrent weights, small-gain orthogonal policy head.
| Path | Purpose |
|---|---|
env.py |
ForagingEnv: the vectorised task (reward timing, noisy cues, push cost); fully documented |
config.py |
ConfigCore: all hyperparameters; saves and loads itself next to the checkpoints |
Model.py |
ActorCritic recurrent network and weight initialisation |
Agent.py |
Agent (action selection, GAE, PPO update, save/load) and rollout Buffer |
train.py |
Command-line trainer (python train.py --seed 0); writes checkpoints to runs/seed<k>/ |
train.ipynb |
The same training loop, interactive |
figures.ipynb |
Loads a trained agent, simulates it at several noise levels, and produces the figures above plus decoding analyses |
tests/test_env.py |
Unit tests for the environment (pytest) |
requirements.txt |
torch, numpy, scipy, scikit-learn, matplotlib, jupyter, pytest |
pip install -r requirements.txt
pytest # environment tests
python train.py --seed 0 --episodes 201 # ~200 episodes converge; a few minutes on a GPU
python train.py --seed 0 --episodes 5 --batch-size 16 --steps 100 --device cpu # 30-second smoke test
jupyter lab figures.ipynb # point agent.load(...) at a checkpoint name from runs/seed<k>/train.py picks cuda:0 when available, otherwise CPU. Trained checkpoints are not stored in the
repository.
Built in 2023 as the modelling half of a project on foraging under sensory uncertainty (see foraging-allocation-dynamics and the 2026 bioRxiv preprint linked there). The task abstracts a freely moving macaque arena with three push-button patches and cues of controlled reliability.
MIT (see LICENSE).