aind_dynamic_foraging_models.generative_model package

Subpackages

Submodules

aind_dynamic_foraging_models.generative_model.act_functions module

aind_dynamic_foraging_models.generative_model.base module

Base class for DynamicForagingAgent with MLE fitting

class aind_dynamic_foraging_models.generative_model.base.DynamicForagingAgentMLEBase(agent_kwargs: dict = {}, params: dict = {}, **kwargs)[source]

Bases: DynamicForagingAgentBase

Base class of “DynamicForagingAgentBase” + “MLE fitting”

act(observation)[source]

Chooses an action based on the current observation. I just copy and paste this from DynamicForagingAgentBase here for clarity.

Parameters:

observation – The current observation from the environment.

Returns:

The action chosen by the agent.

Return type:

action

fit(fit_choice_history, fit_reward_history, fit_bounds_override: dict | None = {}, clamp_params: dict | None = None, k_fold_cross_validation: int | None = None, DE_kwargs: dict | None = {'workers': 1})[source]

Fit the model to the data using differential evolution.

It handles fit_bounds_override and clamp_params as follows: 1. It will first clamp the parameters specified in clamp_params 2. For other parameters, if it is specified in fit_bounds_override, the specified

bound will be used; otherwise, the bound in the model’s ParamFitBounds will be used.

For example, if params_to_fit and clamp_params are all empty, all parameters will be fitted with default bounds in the model’s ParamFitBounds.

Supports both single-session and multi-session fitting.

Multi-session format is a list/tuple of per-session 1D arrays. In multi-session fitting, latent states (Q-values, choice kernels, etc.) are reset at the start of each session.

Parameters:

  • fit_choice_history (np.ndarray or List[np.ndarray]) – Choice history for fitting. Can be either: - A single 1D array for single-session fitting (backward compatible) - A list of 1D arrays for multi-session fitting, one array per session

  • fit_reward_history (np.ndarray or List[np.ndarray]) – Reward history for fitting. Must match the format of fit_choice_history.

  • fit_bounds_override (dict, optional) – Override the default bounds for fitting parameters in ParamFitBounds, by default {}

  • clamp_params (dict, optional) – Specify parameters to fix to certain values, by default {}

  • k_fold_cross_validation (Optional[int], optional) –

    Whether to do cross-validation, by default None (no cross-validation). If k_fold_cross_validation > 1, it will do k-fold cross-validation and return the prediction accuracy of the test set for model comparison. Cross-validation behavior: - Single-session input: trial-level k-fold CV (backward compatible) - Multi-session input: session-level k-fold CV (entire sessions held out)

    Requires n_sessions >= k_fold_cross_validation.

  • DE_kwargs (dict, optional) –

    kwargs for differential_evolution, by default {‘workers’: 1} For example:

    workersint

    Number of workers for differential evolution, by default 1. In CO, fitting a typical session of 1000 trials takes:

    1 worker: ~100 s 4 workers: ~35 s 8 workers: ~22 s 16 workers: ~20 s

    (see https://github.com/AllenNeuralDynamics/aind-dynamic-foraging-models/blob/22075b85360c0a5db475a90bcb025deaa4318f05/notebook/demo_rl_mle_fitting_new_test_time.ipynb) # noqa E501 That is to say, the parallel speedup in DE is sublinear. Therefore, given a constant number of total CPUs, it is more efficient to parallelize on the level of session, instead of on DE’s workers.

Returns:

  • fitting_result (OptimizeResult) – The fitting result object containing optimized parameters and statistics.

  • fitting_result_cross_validation (dict or None) – Cross-validation results if k_fold_cross_validation is specified, else None.

get_agent_alias()[source]

Get the agent alias for the model

Should be overridden by the subclass.

get_choice_history()[source]

Return the history of actions in format that is compatible with other library such as aind_dynamic_foraging_basic_analysis

get_fitting_result_dict()[source]

Return the fitting result in a json-compatible dict for uploading to docDB etc.

get_latent_variables()[source]

Return the latent variables of the agent

This is agent-specific and should be implemented by the subclass.

get_p_reward()[source]

Return the reward probabilities for each arm in each trial which is compatible with other library such as aind_dynamic_foraging_basic_analysis

get_params()[source]

Get the model parameters in a dictionary format

get_params_str(if_latex=True, if_value=True, decimal=3)[source]

Get string of the model parameters

Parameters:

  • if_latex (bool, optional) – if True, return the latex format of the parameters, by default True

  • if_value (bool, optional) – if True, return the value of the parameters, by default True

  • decimal (int, optional)

get_reward_history()[source]

Return the history of reward in format that is compatible with other library such as aind_dynamic_foraging_basic_analysis

learn(observation, action, reward, next_observation, done)[source]

Updates the agent’s knowledge or policy based on the last action and its outcome. I just copy and paste this from DynamicForagingAgentBase here for clarity.

This is the core method that should be implemented by all non-trivial agents. It could be Q-learning, policy gradients, neural networks, etc.

Parameters:

  • observation – The observation before the action was taken.

  • action – The action taken by the agent.

  • reward – The reward received after taking the action.

  • next_observation – The next observation after the action.

  • done – Whether the episode has ended.

perform(task: DynamicForagingTaskBase)[source]

Generative simulation of a task, or “open-loop” simulation

Override the base class method to include choice_prob caching etc.

In each trial loop (note when time ticks):

agent.act() task.step() agent.learn()

latent variable [t] –> choice [t] –> reward [t] —-> update latent variable [t+1]

perform_closed_loop(fit_choice_history, fit_reward_history)[source]

Simulates the agent over a fixed choice and reward history using its params. Also called “teacher forcing” or “closed-loop” simulation.

Unlike .perform() (“generative” simulation), this is called “predictive” simulation, which does not need a task and is used for model fitting.

perform_closed_loop_multi_session(fit_choice_history_sessions, fit_reward_history_sessions)[source]

Simulates the agent over multiple sessions, resetting latent states at each session start.

Unlike .perform() (“generative” simulation), this is called “predictive” simulation, which does not need a task and is used for model fitting across multiple sessions.

Parameters:

  • fit_choice_history_sessions (list of np.ndarray) – List of choice history arrays, one per session

  • fit_reward_history_sessions (list of np.ndarray) – List of reward history arrays, one per session

Returns:

choice_prob_sessions – List of choice probability arrays, one per session

Return type:

list of np.ndarray

plot_fitted_session(if_plot_latent=True)[source]

Plot session after .fit()

  1. choice and reward history will be the history used for fitting

  2. laten variables q_estimate and choice_prob will be plotted

  3. p_reward will be missing (since it is not used for fitting)

Parameters:

if_plot_latent (bool, optional) – Whether to plot latent variables, by default True

plot_latent_variables(ax, if_fitted=False)[source]

Add agent-specific latent variables to the plot

if_fitted: whether the latent variables are from the fitted model (styling purpose)

plot_session(if_plot_latent=True)[source]

Plot session after .perform(task)

Parameters:

if_plot_latent (bool, optional) – Whether to plot latent variables, by default True

set_params(**params)[source]

Update the model parameters and validate

aind_dynamic_foraging_models.generative_model.base.negLL(choice_prob, fit_choice_history, fit_reward_history, fit_trial_set=None)[source]

Compute total negLL of the trials in fit_trial_set given the data.

aind_dynamic_foraging_models.generative_model.forager_loss_counting module

aind_dynamic_foraging_models.generative_model.forager_q_learning module

aind_dynamic_foraging_models.generative_model.foragers module

aind_dynamic_foraging_models.generative_model.learn_functions module

Functions for update latent variables in generative models.

aind_dynamic_foraging_models.generative_model.learn_functions.learn_RWlike(choice, reward, q_value_tminus1, forget_rates, learn_rates)[source]

Learning function for Rescorla-Wagner-like model.

Parameters:

  • choice (int) – this choice

  • reward (float) – this reward

  • q_value_tminus1 (np.ndarray) – array of old q values

  • forget_rates (list) – forget rates for [unchosen, chosen] sides

  • learn_rates (_type_) – learning rates for [rewarded, unrewarded] sides

Returns:

array of new q values

Return type:

np.ndarray

aind_dynamic_foraging_models.generative_model.learn_functions.learn_actor_critic(choice, reward, actor_preference_tminus1, value_tminus1, choice_prob_t, learn_rates)[source]

One-step actor-critic update for a (state-less) two-armed foraging task.

This implements the classic Sutton & Barto actor-critic with a softmax actor:

  1. Critic (state value V) is updated by the TD error. Because the dynamic foraging task is effectively state-less (a single state / bandit), there is no bootstrap term, so the TD error reduces to the reward-prediction error:

    delta_t = reward_t - V_{t-1}
V_t     = V_{t-1} + alpha_critic * delta_t

    Here V acts as a running reward baseline shared across actions.

  2. Actor (per-action preference H) is updated along the softmax policy gradient, using the same TD error as the critic’s teaching signal:

    H_t(a) = H_{t-1}(a) + alpha_actor * delta_t * (1{a == choice} - pi(a))

    where pi is the policy (choice probability) that was actually used to select the action on this trial.

Parameters:

  • choice (int) – This trial’s choice (action index).

  • reward (float) – This trial’s reward.

  • actor_preference_tminus1 (np.ndarray) – Array of old actor preferences H, one entry per action.

  • value_tminus1 (float) – Old critic state value V (scalar).

  • choice_prob_t (np.ndarray) – The policy pi (choice probability per action) that was used to make this trial’s choice. Used for the softmax policy-gradient term.

  • learn_rates (list) – Learning rates as [alpha_actor, alpha_critic].

Returns:

(new actor preferences H, new critic value V)

Return type:

tuple(np.ndarray, float)

Notes

The actor’s overall scale (alpha_actor) and the policy’s softmax inverse temperature (beta) are not jointly identifiable, because both scale the preferences inside the softmax. When fitting, it is therefore recommended to clamp softmax_inverse_temperature (e.g. to 1.0) and let learn_rate_actor absorb the scale. For that reason, beta is intentionally not multiplied into the actor gradient here.

aind_dynamic_foraging_models.generative_model.learn_functions.learn_choice_kernel(choice, choice_kernel_tminus1, choice_kernel_step_size)[source]

Learning function for choice kernel.

Parameters:

  • choice (int) – this choice

  • choice_kernel_tminus1 (np.ndarray) – array of old choice kernel values

  • choice_kernel_step_size (float) – step size for choice kernel

Returns:

array of new choice kernel values

Return type:

np.ndarray

aind_dynamic_foraging_models.generative_model.learn_functions.learn_loss_counting(choice, reward, just_switched, loss_count_tminus1) → int[source]

Update loss counting

Returns the new loss count

Module contents