aind-dynamic-foraging-models

AIND library of models for dynamic foraging tasks: generative agents fitted per session by maximum likelihood, logistic regression on choice history, hierarchical Bayesian fits across sessions and subjects, and GLM-HMMs with discrete latent states.

User documentation available on readthedocs.

Installation

pip install aind-dynamic-foraging-models            # logistic regression only
pip install aind-dynamic-foraging-models[rl]        # + generative agents
pip install aind-dynamic-foraging-models[bayes]     # + hierarchical Bayesian fitting
pip install aind-dynamic-foraging-models[glmhmm]    # + GLM-HMM

• [bayes] needs Python 3.12+: every arviz 1.x wheel requires it, and pip on 3.10 or 3.11 fails with "No matching distribution found for arviz". pyproject.toml explains the pin.
• [bayes] and [glmhmm] bound jax below 0.7, to match the disRNN wrapper's jax pin. CI tests both extras on Python 3.12, and the core package on 3.10 and 3.11.

To develop the code, clone the repo to your local machine, and run

pip install -e .[dev]

Reinforcement Learning (RL) models with Maximum Likelihood Estimation (MLE) fitting

Overview

RL agents that can perform any dynamic foraging task in aind-behavior-gym and can fit behavior using MLE.

Code structure

• To add more generative models, please subclass DynamicForagingAgentMLEBase.

Implemented foragers

• ForagerQLearning: Simple Q-learning agents that incrementally update Q-values.
  • Available agent_kwargs:

      number_of_learning_rate: Literal[1, 2] = 2,
      number_of_forget_rate: Literal[0, 1] = 1,
      choice_kernel: Literal["none", "one_step", "full"] = "none",
      action_selection: Literal["softmax", "epsilon-greedy"] = "softmax",
    

• ForagerLossCounting: Loss counting agents with probabilistic loss_count_threshold.
  • Available agent_kwargs:

      win_stay_lose_switch: Literal[False, True] = False,
      choice_kernel: Literal["none", "one_step", "full"] = "none",
    

• ForagerActorCritic: Actor-critic agents. A softmax actor holds a preference per side; a single state-value critic gives the reward baseline.
  • Available agent_kwargs:

      choice_kernel: Literal["none", "one_step", "full"] = "none",
    

• ForagerCompareThreshold: Compare-to-threshold (CTT) agents. One value is compared with a threshold through a sigmoid: above it the agent is more likely to stay, below it more likely to switch.
  • Available agent_kwargs:

      number_of_learning_rate: Literal[1, 2] = 1,
      choice_kernel: Literal["none", "one_step", "full"] = "none",
      reset_to_threshold: Literal[True, False] = True,
    

• Action selections (readthedoc)

Published-model reproductions

forager_published_bandits.py re-implements published two-armed-bandit models as MLE agents with fixed structure (no agent_kwargs). docs/published_bandit_reproductions.md records each model's equation source, the author code it was checked against, and every deliberate departure from it.

The Findling et al. Weber-imprecision Bayesian-inference model (findling_weber.py) is not an MLE agent: its likelihood averages over a stochastic belief state with a particle filter. Fit it with fit_findling_weber_map over the 1,000-candidate grid from findling_weber_parameter_grid.

Here is the list of foragers in the RL model playground. ForagerCollection().get_all_foragers() lists every forager and preset in the installed version.

Usage

• Jupyter notebook
• See also these unittest functions.

RL model playground

Play with the generative models here.

Logistic regression

See this demo notebook.

Choosing logistic regression models

Su 2022

$$ logit(p(c_r)) \sim RewardedChoice+UnrewardedChoice $$

Bari 2019

$$ logit(p(c_r)) \sim RewardedChoice+Choice $$

Hattori 2019

$$ logit(p(c_r)) \sim RewardedChoice+UnrewardedChoice+Choice $$

Miller 2021

$$ logit(p(c_r)) \sim Choice + Reward+ Choice*Reward $$

Encodings

• Ignored trials are removed

Some observations:

1. $RewardedChoice$ and $UnrewardedChoice$ are orthogonal
2. $Choice = RewardedChoice + UnrewardedChoice$
3. $Choice * Reward = RewardedChoice - UnrewardedChoice$

Comparison

Regularization and optimization

The choice of optimizer depends on the penality term, as listed here.

• lbfgs - [l2, None]
• liblinear - [l1, l2]
• newton-cg - [l2, None]
• newton-cholesky - [l2, None]
• sag - [l2, None]
• saga - [elasticnet, l1, l2, None]

Hierarchical Bayesian fitting

aind_dynamic_foraging_models.hierarchical_bayes fits the foragers as hierarchical Bayesian models (population → subject → session) with NumPyro, as a partially-pooled counterpart to the per-session MLE fitting in generative_model.

pip install aind-dynamic-foraging-models[bayes]

• Two forager likelihoods, chosen by model= (hierarchical_bayes.model.MODELS): "hattori2019" (the default) and "ctt" (compare-to-threshold).
• heldout.py scores a held-out subject: it adapts the subject to its context sessions with the population frozen, then scores the subject's test sessions.
• Start here: a worked end-to-end example on synthetic data with known ground truth -- simulate a cohort, fit it both ways, and check that the parameters come back: notebook/demo_hierarchical_bayes.ipynb
• Design decisions and their rationale: src/aind_dynamic_foraging_models/hierarchical_bayes/docs/design-decisions.md
• Validation against the reference Stan implementation, and performance measurements: src/aind_dynamic_foraging_models/hierarchical_bayes/benchmarks/RESULTS.md

GLM-HMM

aind_dynamic_foraging_models.glm_hmm fits an input-driven HMM with GLM emissions (Ashwood et al. 2022, Nat Neurosci 25:201) over left / right / ignore choices, on dynamax.

pip install aind-dynamic-foraging-models[glmhmm]

• Two emissions, both exactly normalised over the three outcomes: hurdle (per state, an engagement GLM on unsigned history and a side GLM on signed history) and reference (per state, one multinomial GLM with ignore as the zero-logit class, as in Faeze's ssm fits). How they differ, and which to use: glm_hmm/README.md.
• Regressors (design.py): RewC and UnrC lags (signed), Ignored and Rewarded lags (unsigned), and a bias; trial 0 of every session sees the bias alone.
• Fitting (model.py): EM over many variable-length sessions, bucketed by length with exact masking. The emission M-step is exact Newton with a Gaussian prior; pi and A use a Dirichlet MAP. Priors are arguments, so one model serves a near-flat global fit (GLOBAL_PRIOR) and per-subject MAP refits started from it (SUBJECT_PRIOR). Prior.mean centres the weight prior, for example on the global fit's weights for a subject refit; the default None is zero.
• fit_with_restarts(..., max_workers=N) runs EM restarts at once and shares their row compression; the GLM-HMM README gives the speed-up on a CPU.
• Scoring: predictive_log_probs is the forward filter's one-step-ahead log p(y_t | y_<t) per class. It is the held-out quantity, and it never uses later trials. state_posteriors gives filtered and smoothed posteriors and Viterbi paths.
• tfp-nightly is pinned to one build in the extra: dynamax leaves it unpinned, and only some nightlies work with jax 0.6.1 and numpy 2.4.
• Device: float64 on the CPU by default; AIND_GLMHMM_DEVICE=gpu moves the fits to a GPU. GPU linear solves need the CUDA solver libraries pinned, as model._compute explains. A GPU is faster for a large pooled fit, but not for one mouse's refit.
• Benchmark: glm_hmm/benchmarks/RESULTS.md times CPU against V100, A100 and L40S on study 14's 4.9M-trial cohort, with the scripts to rerun it. The GPU decision is tracked in AllenNeuralDynamics/aind-dynamic-foraging-bfm-wrapper#111.
• Validation: glm_hmm/recovery/ simulates from known parameters and refits them, and cross-checks the likelihood against ssm (RESULTS.md). The study version, including a factorial of designed strategies, is AllenNeuralDynamics/aind-dynamic-foraging-bfm-studies#50.

See also

• Foraging model simulation, model recovery, etc.: https://github.com/hanhou/Dynamic-Foraging

Contributing

Linters and testing

There are several libraries used to run linters, check documentation, and run tests.

• Please test your changes using the coverage library, which will run the tests and log a coverage report:

coverage run -m unittest discover && coverage report

• Use interrogate to check that modules, methods, etc. have been documented thoroughly:

interrogate .

• Use flake8 to check that code is up to standards (no unused imports, etc.):

flake8 .

• Use black to automatically format the code into PEP standards:

black .

• Use isort to automatically sort import statements:

isort .

Pull requests

For internal members, please create a branch. For external members, please fork the repository and open a pull request from the fork. We'll primarily use Angular style for commit messages. Roughly, they should follow the pattern:

<type>(<scope>): <short summary>

where scope (optional) describes the packages affected by the code changes and type (mandatory) is one of:

• build: Changes that affect build tools or external dependencies (example scopes: pyproject.toml, setup.py)
• ci: Changes to our CI configuration files and scripts (examples: .github/workflows/ci.yml)
• docs: Documentation only changes
• feat: A new feature
• fix: A bugfix
• perf: A code change that improves performance
• refactor: A code change that neither fixes a bug nor adds a feature
• test: Adding missing tests or correcting existing tests

Semantic Release

The table below, from semantic release, shows which commit message gets you which release type when semantic-release runs (using the default configuration):

Documentation

To generate the rst files source files for documentation, run

sphinx-apidoc -o docs/source/ src

Then to create the documentation HTML files, run

sphinx-build -b html docs/source/ docs/build/html

More info on sphinx installation can be found here.