Benchmark comparison.

First, we sought to validate our ANN method and compare its performance to existing methods by testing it on standard likelihood-tractable cognitive models of different levels of complexity in the same task: 2-parameter (2P − RL) and 4-parameter (4P − RL) reinforcement learning models commonly used to model behavior on reversal tasks [7, 14, 53, 54], as well as Bayesian Inference model (BI) and Bayesian Inference with Stickiness (S − BI) as an alternative model family that has been found to outperform RL in some cases [55–57]. We estimated model parameters using multiple traditional methods for computing (maximum likelihood and maximum a-posteriori estimation; MLE and MAP) and approximating (Approximate Bayesian Computation; ABC) likelihood. We used the results of these tools as a benchmark for evaluating the neural network approach. Next, we estimated parameters of these models using two variants of RNNs: with gated recurrent units (GRUs) or Long-Short-Term-Memory units (LSTM).

We used the same held-out data set to evaluate all methods (the test set the ANN has not observed yet, see simulation details). For each of the methods, we extracted the best-fit parameters, and then quantitatively estimated the method’s performance as the mean squared error (MSE) between estimated and true parameters across all agents. Methods with lower MSE indicated better relative performance. All of the parameters were scaled for the purpose of loss computation, to ensure comparable contribution to loss across different parameters. To quantify overall loss for a cognitive model, we averaged across all individual parameter MSE scores; to calculate fitting method’s MSE score for a class of cognitive models (e.g. likelihood tractable models) we averaged across respective method’s MSE scores for those models (see Methods for details about method evaluation).

First, we examined the performance of standard model-fitting tools (MLE, MAP and ABC). The standard tools yielded a pattern of results that are expected based on noisy, realistic-size data sets (with several-hundred trials per agent). Specifically, we found that MAP outperformed MLE (Fig 2A, average MSEs: MLE =.67, MAP =.35), since the parameter prior applied in MAP regularizes the fitting process. ABC was also worse compared to MAP (Fig 2A, average MSE: ABC =.53). While the fitting process is also regularized in ABC, worse performance in some models can be attributed to signal loss that arises from approximation to the likelihood. Next, we focused on the ANN performance; our results showed that for each of the models, ANN performed better than or just as well as the traditional methods (Fig 2A, average MSEs for different RNN variants: GRU =.32, LSTM =.35). Better network performance was more evident for parameter estimation in more complex models (e.g., models with a higher number of parameters such as 4P-RL and S-BI; average MSE across these 2 models: MLE =.95, MAP =.43, ABC =.71, GRU =.38, LSTM =.44).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2.

A) Parameter recovery loss from the held-out test set for the tractable-likelihood models (2P-RL, 4P-RL, BI, S-BI) using each of the tested methods. Loss is quantified as the mean squared error (MSE) based on the discrepancy between true and estimated parameters. Bars represent loss average for each parameter across all agents, with error bars representing standard error across agents. B) Parameter recovery from the 4P-RL model using MAP and GRU. ρ values represent the Spearman ρ correlation between true and estimated parameters. Red line represents a unity line (x = y) and black line represents a least squares regression line. All correlations were significant at p <.001.

https://doi.org/10.1371/journal.pcbi.1012119.g002

Next, we visualized parameter recovery. We found that for each of the cognitive models the parameter recovery was largely successful (Spearman ρ correlations between true parameter values and estimated values: β ρ_MAP, ρ_GRU = [.90, .91], α⁺ρ_MAP, ρ_GRU = [.53, .52], α⁻ρ_MAP, ρ_GRU = [.88, .89], κ: ρ_MAP, ρ_GRU = [.78, .79], Fig 2B; all correlations were significant at p <.001). For conciseness, we only show recovery of the more complex model parameters from the RL model family (and only MAP method as it performed better compared to ABC and MLE, as well as only GRU since it performed better than LSTM), as we would expect a more complex model to emphasize superiority of a fitting method more clearly compared to simpler models. Recovery plots of the remaining models (and respective fitting methods) can be found in S2–S5 Figs. Our results suggest that 1) ANN performed as well as traditional methods in parameter estimation based on the MSE loss; 2) more complex models may limit the accuracy of parameter estimation in traditional methods that neural networks appear to be more robust against. We note that for the 4P − RL model, parameter recovery was noisy for all methods, with some parameters being less recoverable than others (e.g. α⁺, Fig 2B). This is an expected property of cognitive models applied to realistic-sized experimental data as found in most human experiments (i.e., a few hundred trials per participant). To check whether the limited recovery can be attributed to parameter identifiability rather than pitfalls of any specific method, we looked at the correlation between parameter estimates obtained using the standard model fitting method (MAP) and the ANN (GRU) (S10 Fig)—with parameters that are not well recovered (e.g., α⁺ in 4P-RL model) being of particular interest. High correlation between estimated parameters obtained via 2 methods imply systematic errors in parameter identification that apply to both methods—thus suggesting that the weaker correlation between true and fit parameters for some parameters is more likely due to limitations in the model applied to the data set than method specifications such as poor optimization performance. We further discuss the implications in the discussion section—highlighting that computational models should be carefully crafted and specified regardless of the tools used for model fitting.

Testing in cognitive models with intractable likelihood.

Next, we tested our method in two examples of computational models with intractable likelihood. As a comparison method, we implemented Approximate Bayesian Computation (ABC), alongside our ANN approach to estimate parameters. The two example likelihood-intractable models we used had in common the presence of a latent state which conditioned sequential updates: RL with latent attentive state (RL − LAS), and a form of non-temporal hierarchical reinforcement learning (HRL, [27]). Since we cannot fit these models using MAP or MLE we used only ABC as a benchmark. Because we found LSTM RNN to be more challenging to train and achieve similar results when compared to GRU, we focused on GRU for the remainder of comparisons. We found that average MSE was much lower for the neural network compared to ABC for both RL-LAS (Fig 3A, average MSEs: ABC =.62, GRU =.21) and HRL (Fig 3A, average MSEs: ABC =.28, GRU =.19). Spearman correlations were noisier for ABC compared to GRU in both models (Fig 3B, RL-LAS: β ρ_ABC, ρ_GRU = [.72, .91], α ρ_ABC, ρ_GRU = [.83, .95], T ρ_ABC, ρ_GRU = [.5, .81]; HRL: β ρ_ABC, ρ_GRU = [.86, .89], α ρ_ABC, ρ_GRU = [.85, .9]; all correlations were significant at p <.001). Furthermore, some parameters were less recoverable than others (e.g., the T parameter in RL-LAS model, which indexed how long participants remained in an inattentive state); this might be in part due to less straightforward effect of T on behavior (S6 Fig). Note that, in order to obtain our ABC results we had to perform an extensive exploration procedure to select summary statistics—ensuring reasonable ABC results. Indeed, the choice of summary statistics is not trivial and represents an important difficulty of applying basic rejection ABC [33, 38], that we can entirely bypass using our new neural network approach. We acknowledge that recent methods that rely on ANNs replaced standard ABC methods by automating (or semi-automating) the construction of summary statistics [38, 40–44, 51]. However, we aimed to explore an alternative approach, independent of explicit optimization of summary statistics, and focused on the ABC instantiation that has been most frequently implemented in the field of cognitive science as a benchmark [33–35].

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3.

A) Parameter recovery loss from the held-out test set for the intractable-likelihood models (RL-LAS, HRL) using ABC and GRU network. Loss is quantified as the mean squared error (MSE) based on the discrepancy between true and estimated parameters. Bars represent MSE average for each parameter across all agents, with error bars representing standard error across agents ((S17 Fig) shows variability across seeds). B) Parameter recovery from the RL-LAS and HRL models using ABC (green) and GRU network (yellow). ρ values represent the Spearman ρ correlation between true and estimated parameters. Red line represents a unity line (x = y) and black line represents a least squares regression line.All correlations were significant at p <.001.

https://doi.org/10.1371/journal.pcbi.1012119.g003

Uncertainty of parameter estimates.

Thus far, we have outlined a method that provides point estimates of parameters based on input data sequences, as is typically the use for much lightweight cognitive modeling (e.g., maximum likelihood estimation or MAP). However, it is sometimes also valuable to compute the uncertainty associated with these estimates [21]. It is possible to extend our approach to add this capability. While there are various alternative ways to do so (e.g. Bayesian neural networks), the approach we have opted for is incorporating evidential learning into our method [58]. Evidential learning differs from Bayesian networks in that it places priors over likelihood function, rather than network weights. The network leverages this property to learn both statistical (aleatoric) and systematic (epistemic) uncertainty during the process of estimating a continuous target based on the input data sequences. This marks a shift from optimizing a network to minimize errors based on average prediction, without considering uncertainty.

We applied our method with integrated evidential learning to tractable and intractable versions of the RL models (2P-RL and RL-LAS, Fig 4). We found that incorporating this modification did not compromise the point estimate parameter recovery (e.g. compared to our baseline method focused only on maximizing the accuracy of the predictions). Additionally, it enabled the estimation of the uncertainty around the point estimate, as demonstrated by [58]. This extension appears to be more computationally expensive (with longer training periods) than our original method, but not to a prohibitive extent.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 4.

Using evidential learning to evaluate uncertainty of parameter estimates for A) the 2-parameter RL model (tractable likelihood) and B) the RL model with latent attention states (intractable likelihood). Vertical lines around point estimates illustrate model uncertainty. We are showing only 100 data points for cleaner visualization, Spearman ρ values are computed based on the total number of agents in the held-out test data (3k).

https://doi.org/10.1371/journal.pcbi.1012119.g004

Robustness tests: Influence of different input trial sequence lengths.

ANNs are sometimes known to fail catastrophically when data is different from the training distribution in minor ways [59–62]. Thus, we investigated the robustness of our method to differences in data format we might expect in empirical data, such as different numbers of trials across participants. Specifically, we conducted robustness experiments by varying the number of trials in each individual simulation contributing to training or test sets, fixing the number of agents in the training set.

To evaluate the quality of parameter recovery, we used the coefficient of determination score (R²), which normalizes different parameter ranges. We found that the ANNs trained with a higher trial number reach high R² scores in long test trials. However, their performance suffers significantly with a smaller number of test trials. The results also show a similar trend in model identification tasks except, that training with higher trial number doesn’t guarantee better performance. For instance, the classification accuracy between HRL task models of the ANN trained with 300 trials reaches 87% while the ANN trained with 500 trials is 84%.

Data-augmentation practices in machine learning increase robustness of models during training [63] by introducing different types of variability in the training data set (e.g., adding noise, different data sizes). Specifically, slicing time-series data into sub-series is a data-augmentation practice that increases accuracy [64]. Thus, we trained our ANN with the fixed number of simulations with different trial numbers. As predicted, we found that the ANNs trained with a mixture of trial sequence lengths across simulations (purple line) consistently yielded better performance across different numbers of test trials for both parameter recovery and model identification (Fig 6A and 6B).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 6. Robustness checks using different training (different line colors) and testing (x-axis) trial sequence lengths.

A) Parameter estimation in both RL-LAS and HRL models shows that training with a mixture of trial sequence lengths (purple line) yields more robust out-of-sample parameter value prediction compared to fixed trial sequence lengths. B) Best model identification results, performed on different combinations of model candidates, were also yielded by mixed trial sequence length training. The number of agents/simulations used for training was kept constant across all the tests (N agents = 30k).

https://doi.org/10.1371/journal.pcbi.1012119.g006

Robustness tests: Prior parameter assumptions.

We also tested the effects of incorrect prior assumptions about the parameter range on method performance. Specifically we 1) trained the network using data simulated from a narrow range of parameters (theoretically informed) and 2) trained the network based on broader range of parameter values. Next, we tested both networks in making out-of-sample predictions for test data sets that were simulated from narrow and broad parameter ranges, respectively. The network trained using a narrow parameter range made large errors at estimating parameters for data simulated outside of the range it was trained on; training the network on a broader range overall resulted in smaller error, with some loss of precision for the parameter values in the range of most interest (e.g., the narrow range of parameters the alternative network is trained on). We observed similar results with MAP, where we specified narrow/broad prior (where narrow prior would place high density on a specific parameter range). Notably, training the network using a broader range of parameters while oversampling from a range of interest yielded more accurate parameter estimation compared to MAP with broad priors (Approach described in S9 Fig).

Robustness tests: Model misspecification.

In addition to testing the effects of incorrect priors, we also tested the effect of model misspecification on standard method and ANN performance (focusing on MAP and GRU network, as they performed the best in parameter recovery tests on benchmark models). We fit the Bayesian inference model (without stickiness) to the data simulated from the Bayesian inference model with stickiness using MAP. For the ANN, we trained the neural network to estimate parameters of the Bayesian inference model, and tested it on the separate test set data simulated from the Bayesian inference model with stickiness. For each method, we looked at the correlation between the ground truth Bayesian inference with stickiness model parameters, and the method’s parameter estimates (S13 Fig). Our results suggest that the parameters shared between the 2 models are reasonably recoverable using both MAP and ANN (e.g., the recovery is noisier but comparable to that of parameters in Bayesian models without model misspecification (S4 and S5 Figs); furthermore, the correlation between ground truth and estimated values is similar for the two methods.

To make the model misspecification more extreme, we additionally simulated data from a Bayesian inference model, and estimated RL model parameters from the simulated data. We did this using standard methods (MAP) and ANN, and repeated the same process in reverse (simulating data from an RL model, and fitting Bayesian inference model parameters). We found that both MAP and ANN exhibited similar patterns. That is, in the case of simulating the Bayesian inference model and fitting RL model parameters, the estimated β captured the variance from the true β and p_switch, while the estimated α parameter captured the variance driven by the Bayesian updating parameters p_reward and p_switch (S14 Fig). In the case of simulating RL model and fitting Bayesian inference model parameters, the p_switch parameter captured the noise in the simulated data coming from the β parameter, and the variance from the α parameter was attributed to the p_reward parameter (S15 Fig). We also correlated parameter estimates generated by the two methods. High correlation implies that MAP and GRU generate similar parameter estimates, suggesting that they are impacted by model misspecification in a similar way (S11 Fig).

Probabilistic reversal learning task.

We have simulated data from different models (see the Models section) on a simple probabilistic reversal learning task (PRL; [73, 74]). In the task, an agent chooses between two actions on each trial, and receives a binary outcome (r = 1 [reward] or r = 0 [no reward]). One of the two actions is correct for a number of trials; a correct action is defined as the action that gets rewarded with higher probability (e.g., p(r = 1|action = correct) = 0.80), with 1 − p probability of getting no reward if selected. After a certain number of trials, the correct action reverses; thus the action that was previously rewarded with low probability becomes the more frequently rewarded one (S1 Fig). This simple task (and its variants) has been extensively used to provide mechanistic insights into learning from reinforcement, inferring probabilistic structure of the environment, and people’s ability (or failure) to update the representation of a correct choice.

Hierarchical reinforcement learning task.

We developed a novel task environment that can be solved using a simple but plausible model with intractable likelihood. In this task, an agent observes N arrows (in different colors), each pointing at either left or right direction. The agent needs to learn which arrow is the correct one, by selecting an action that corresponds to either left or right side (consistent with the direction the arrow is pointing at) in order to get rewarded. Selecting the side the correct arrow is pointing at rewards the agent with high probability (p =.9); choosing an action by following direction of other arrows leads to no reward (r = 0) with same high probability. The correct arrow changes unpredictably in the task, which means that the agent must keep track of which arrow most reliably leads to the reward, and update accordingly when the change occurs. We refer to this task structure as hierarchical because the choice policy (left/right) depends on the higher-level rule (color) agents choose to follow.

PRL task models.

We implemented multiple models of the PRL task to test the artificial neural network (ANN) approach to parameter estimation. First, we cover the benchmark models; these are the models that we can fit using traditional methods (MLE, MAP), as well as the ANN, to ensure that we can justify using the ANN if it performs at least just as well as (or better than) the traditional methods.

Reinforcement learning models family.

Two-parameter reinforcement learning model. We simulated artificial data on the PRL task using a simple 2-parameter reinforcement learning model (2P-RL). The model assumes that the agent tracks the value of each action contingent on the reward history, and uses these values to inform the action selection on each trial.

The model uses simple delta rule to update action values on each trial upon outcome observation, by first computing the reward prediction error (RPE, δ) as the discrepancy between the expected and the observed outcome, and then adjusting the value of the chosen action using the RPE scaled by the learning rate (α) [75]: (1)

We also allowed for counterfactual updating, where the value of the non-chosen action also gets updated on each trial [7, 76]: (2)

The action values are transformed into action probabilities using the softmax function, thus defining a policy where actions with higher value are chosen with higher probabilities. The β parameter controls how deterministic the choices are with higher values of β corresponding to more deterministic choices: (3)

The 2p-RL model contained following free parameters: learning rate (α) and softmax beta (β).

Four-parameter reinforcement learning model. The four parameter RL (4P-RL) model follows the same updating and policy structure as the 2-parameter RL, with 2 main differences. The 4P-RL model differentiates between positive and negative feedback ([77]), by using different learning rates—α⁺ and α⁻ for updating action values after positive and negative outcomes respectively:

Furthermore, 4P-RL model also includes the stickiness parameter κ which captures the tendency to repeat choice from the previous trial: (4)

Like in the 2P-RL we also included counterfactual updating of values for non-selected actions. The 4P-RL model included following free parameters: positive learning rate (α⁺), negative learning rate (α⁻), softmax beta (β) and stickiness (κ).

Bayesian models family.

Bayesian inference model. Bayesian inference model (BI) assumes that an agent infers the latent state in the environment, updates the latent state based on new observations, and uses the inference process to make rewarding choices. For instance, in the PRL task, the agent infers a latent state corresponding to the correct action (C_t: a_right = cor or C_t: a_left = cor) at time t. The agent tracks and updates their belief over which one of the two actions is currently the correct one based on 1) their estimate of the switch frequency (p_switch) and 2) how noisy the reward is (p_reward) from the history of observations up to the previous trial H_t−1. On each trial, the belief is updated according to the Bayes’ rule—based on the prior belief (agent’s model of the task) and likelihood of observed evidence (the outcome given the choice): where i and j are in [left/right], p(C_t = i|H_1:t−1) is the prior probability, and p(r_t|C_t = i, a_t) is the likelihood of outcome given the action. The likelihood is defined in accordance to whether the choice matches the latent state: where p_reward is the parameter controlling the probability of receiving the reward given the choice of correct action. Posterior belief for the correct action is updated to a prior belief for the upcoming trial in accordance with the p_switch parameter, which determines the probability that the correct action might have reversed on the current trial:

Like in the RL models, the action selection in Bayesian models also followed the softmax policy; however, instead of being informed by the Q values the action probabilities were determined by the belief W given the choice and reward history H and the choice parameter β: (9)

The BI model included following parameters: inferred probability of reward given the action determined by the current belief (p_reward), likelihood of the correct action reversing (p_switch) and softmax beta (β).

Bayesian inference model with stickiness. We also added a variation of the Bayesian inference model that accounts for sticky choice behavior (e.g., repeating actions) by introducing a stickiness parameter κ that augments the belief associated with the action chosen on the previous trial:

Intractable likelihood.

As a proof of concept, we implemented a simple model that assumes a latent state of agent’s attention (engaged/disengaged). This model can’t be fit using methods that rely on computing likelihood. While models can have intractable likelihood for a variety of reasons, we focused on leveraging latent variables (e.g., attention state), that are not readily observable in the data. Thus, in the data that is being modeled, only the choices are observed—but not the state the agent was in while executing the choices. The learned choice value which affects the choice likelihood depends on the trial history, including which state the agent was in. Thus, if there are 2 such states, there are 2^N possible sequences that may result in different choice value estimates after N trials. To estimate choice values and likelihood on any given trial one must integrate over the uncertainty of an exponentially increasing latent variable—thus making the likelihood intractable.

RL and latent engagement state.

We simulated a version of a 2p-RL model for a probabilistic reversal learning (PRL) task that also assumes that an agent might occupy two of the latent attention states—engaged or disengaged— during the task (RL-LAS). The model assumes that in the engaged state an agent behaves in accordance with the task structure (e.g., tracks and updates values of actions, and uses action values to inform action selection). In the disengaged state, an agent behaves in a noisy way, in that 1) it does not update the Q value of actions, and 2) chooses between the two actions randomly (e.g., uniform policy) instead of based on their value (e.g., through softmax). Note that assumption 1) is different from a previous version of the model our group considered [78, 79], and is the core assumption that renders the likelihood intractable. The agent can shift between different engagement states at any point throughout the task, and the transition between the states is controlled by a parameter τ. Specifically, for each agent we initialized a random value T between 10 and 30 (which roughly maps onto approximately how many trials an agent spends in a latent attention state), and then used a non-linear transformation to compute τ: 1-(1/T). The value of τ, thus quantifies the probability of transitioning between the two states. The agent was initialized to be in an attentive state at the onset of trials.

The likelihood of this model can be computed: where ls = latent state, l ∈ { 0 = disengaged state, 1 = engaged state }, corresponds to the history of actions and rewards up to the trial t. However, it is in practice intractable, because of the sum over latent states in the equation, which cannot be factored out.

Cognitive models of the HRL task.

Bayesian models of the HRL task. Bayesian models of the HRL task assume an inference process over the latent variable of which arrow is currently the valid arrow, and thus which side (R/L) (given the current trial’s set of arrows) is most likely to result in positive outcome. The inference relies on the generative model of the task determined by parameters p_switch and p_reward, history of trial observations O_t, set of arrows and stochastic choice based on this inference. Initial prior belief over arrows is initialized uniformly prior = 1/nA, where nA corresponds to the number of arrows.

To determine the agent policy over arrows at trial t, we first implemented a softmax function with decision parameter β and prior belief of which arrow is the correct one; once the arrow is selected, the agent implements an ϵ-greedy policy conditioned on the selected arrow A_t to choose a R/L side:

Likelihood p(r_t = 1|A_t, side(A_t)) and posterior are then updated into the prior belief for the next trial using the p_switch model of the task parameter:

This belief is subsequently used to inform arrow choices on the next trial. This model differs from the Bayesian Inference model for the probabilistic task in that 1) p_reward and p_switch parameters are not free/inferred and 2) the choice of the side is stochastic, allowing for a potential lapse in selecting the side that is not consistent with the selected arrow. This model, thus had has following free parameters: decision parameter β and noise parameter ϵ. Like in the Bayesian inference model for the PRL task, we also tested the model variant with stickiness κ parameter that biases beliefs associated with the arrow/side chosen on the previous trial. Both models have tractable likelihoods.

Hierarchical reinforcement learning.

We also simulated a simple hierarchical reinforcement learning (HRL) model to simulate the performance on a HRL task (S1 Fig). This model assumes that an agent tracks the value of each of the arrows, and chooses between the arrows noisily: (5)

We have also explored the model with an assumption that an agent has a tendency to repeat the choice from the previous trial, captured by the stickiness parameter κ: (6)

Once the agent chooses the arrow,it greedily chooses the direction based on which side (left/right) the arrow is pointing at (observable). Note that we only know the side the agent selected (left/right), because the arrow the agent chooses is non-observable. The agent then observes an outcome, and updates the value of the selected arrow based on the observed outcome:

In the case of this model, the likelihood is intractable because of the need to integrate over uncertainty of what rule (which arrow) the agent followed on all of the past trials; because the integration exponentially increases with each time point, the likelihood is not tractable beyond the first several trials: where a_t corresponds to the action dictating which side the agent selected (left/right), corresponds to the task history encoding rewards, selected actions/sides, arrow directions, and c correspond to identity/color of the correct arrow.

Maximum likelihood and maximum a posteriori estimation.

Maximum likelihood estimation (MLE) represents a cornerstone of modeling that leverages probability theory and estimation of likelihood (P(D|M, θ)) of the data given the model parameters and assumptions [18]. The parameter estimates are determined as the values that maximize the likelihood of the data:

Thus, to estimate best fitting parameters via MLE, the likelihood of the data is computed and maximized with respect to parameter values via an optimization algorithm (often a blackbox one, such as fmincon in MATLAB or optimize.minimize from scipy toolbox in python). Maximum a posteriori estimation (MAP) relies on much the same principle, with an addition of a prior p(θ) to maximize the posterior:

As a prior for the MAP approach, we used an empirical prior derived from the true simulating distribution of parameters. We note that this gives an advantage to the MAP method above what would be available for empirical data, allowing MAP to provide a ceiling performance on the test set.

Because MAP and MLE rely on likelihood computation, their use is essentially limited to models with tractable likelihood. We used MAP and MLE to estimate parameters of tractable-likelihood models as one of the benchmarks against which we compared our ANN approach. Specifically, we fit the models to the test-set data used to compute the MSE of the ANN, and compared fit using the same metric across methods (see main text).

Likelihood approximation methods.

Because models with tractable likelihood comprise only a small subset of all possible (and likely more plausible) models, researchers have handled the issue of likelihood intractability by implementing various likelihood approximation methods. While there are different likelihood approximation tools, such as particle filtering [31] and assumed density estimation [32], we focus on Approximate Bayesian Computation (ABC; [33, 36, 37, 70]), as it is more widely accessible and does not require more extensive mathematical expertise. ABC leverages large scale model simulation to approximate likelihood. Specifically, a large synthetic data set is simulated from a model, with parameters randomly sampled from a specific range for each agent. Summary statistics that describe the data (e.g., average accuracy or variance in accuracy) are used to construct the empirical likelihood that can be used in combination with classic methods.

We implemented a basic form of ABC—the rejection algorithm [33]. This algorithm first samples a set of model parameters θ, simulates the data from the model M using these parameters and computes the predetermined summary statistic of the simulated data which we refer to as the sample. The summary statistics of the real data S(D) and the sample are then compared—if the distance between the two sets of summary statistics ρ is greater than the predetermined criterion ϵ, the sample is rejected:

The distance metric, like the rejection criterion, is determined by the researcher. The samples that are accepted are the samples with distance to the real data smaller than the criterion, resulting in the conclusion that parameters used to generate the sample data set can plausibly be the ones that capture the target data. Thus, the result of the ABC for each data set is a distribution of plausible parameter values which can be used to obtain point estimates via the mean, median, etc.

ABC is a valuable tool, but standard ABC has serious limitations [33]. For instance, the choice of summary statistics is not a trivial problem, and different summary statistics can yield significantly different results. Similarly, in the case of rejection algorithm ABC, researchers must choose the rejection criterion which can also affect the parameter estimates. A possible way to address this is using cross validation to determine which rejection criterion is the best, but this also requires specification of the set of possible criteria values for the cross validation algorithm to choose from. Furthermore, one of ABC assumptions is the independence of data points, which is violated in many sequential decision making models (e.g. reinforcement learning).

To compare our approach to ABC, we used network training set data as a large-scale simulation data set, and then estimated parameters of the held-out test set also used to evaluate the ANN.

To apply ABC in our case, we needed to select summary statistics that adequately describe performance on the task. We used the following summary statistics to quantify the agent for the models simulated on the PRL task:

• Learning curves: We computed agents’ probability of selecting the correct action, aligned to the number of trials with reference to the reversal point. Specifically, for each agent we computed an average proportion of trials where a correct action was selected N trials before and N trials after the correct action reversal point, for all reversal points throughout the task. This summary statistic should capture learning dynamics, as the agent learns to select the correct action, and then experiences dip in accuracy once the correct actions switch, subsequently learning to adjust based on feedback after several trials.
• 3-back feedback integration: The 3-back analysis quantifies learning as well; however, instead of aligning the performance to reversal points, it allowed us to examine agents’ tendency to repeat action selection from the previous trial contingent on reward history—specifically the outcome they observed on the most recent 3 trials. Higher probability of repeating the same action following more positive feedback indicates learning/sensitivity to reward as reported in [11]
• Ab-analysis: The Ab-analysis allowed us to quantify probability of selecting an action at trial t, contingent both on previous reward and action selection history (trials t − 2 and t − 1, [11, 80]).

For the models simulated on a hierarchical task we used the learning curves as summary statistics (same as for the PRL), where reversal points were defined as the switch of the correct rule/arrow to follow. In addition, we quantified the agent’s propensity to stick with the previously correct rule/arrow, where the agent should be increasingly less likely to select the side consistent with the arrow that was correct before the switch, as the number of trials since the switch increases. Similarly, we used a version of the 3-back analysis where the probability of staying contingent on the reward history referred to the probability of potentially selecting the same cue across the trial window, based on observed choices of the agent. All summary statistics are visualized in S8 Fig.

Model comparison.

To perform benchmark model comparison, we used the Akaike Information Criterion (AIC) metric [5], commonly used to evaluate relative model fitness, with the aim of identifying the best model candidate that might have generated the data. The AIC score combines model log likelihood and number of parameters to quantify model fitness, while also penalizing for model complexity in order to prevent overfitting: where K corresponds to the number of parameters. The model with the lowest AIC scores corresponds to the best fitting model for the given data. A related metric that is commonly used is the Bayesian Information Criterion (BIC, [22]), which considers the number of observations (N) as well, and similarly uses the lowest score to signal the best fitting model:

We used AIC score as it outperformed BIC model comparison, and thus provided us with ceiling benchmark to evaluate the ANN.

To perform proper model comparison, it is essential to not only evaluate the model fitness (overall AIC/BIC score), but also to test how reliably the true models (that generated the data) can be identified and successfully distinguished from others. To do so, we constructed a confusion matrix based on the AIC score (Fig 5A). We used the test set data simulated from each model, and then fit all candidate models to each of the data sets while also computing the AIC score for each fit. If the models are identifiable, we should observe that AIC scores for true models (e.g., the models the data was simulated from) should be the lowest for that model when it’s fit to the data compared to other model candidates.

Parameter recovery.

To implement ANNs for parameter estimation we have used the relatively simple neural network structure inspired by the previous work [52]. In all experiments, we used 1 recurrent GRU layer followed by 3 fully connected dense layers with 2000 dimensional input embeddings (S1 Table). To train the network, we simulated a training data set using known parameters. For each model, we used 30000 training samples, 3000 validation samples, and 3000 test samples that are generated from simulations separately. For probabilistic RL, the input sequence consisted of rewards and actions. For hierarchical RL, the sides (left/right) of three arrow stimuli are added to the rewards and actions sequences. The network output dimension was proportional to the number of model parameters. We used a tanh activation in the GRU layer, ReLu activations in 2 dense layers, and a linear activation at the final output. Additional training details are given below:

• We used He normal initialization to initialize GRU parameters [81].
• We used the Adam optimizer with mean square error (MSE) loss and a fixed learning rate of 0.003. Early stopping (e.g. network training was terminated if validation loss failed to decrease after 10 epochs) was applied with a maximum of 200 epochs.
• We selected network hyperparameters with Bayesian optimization algorithms [82] applied on a validation set. Details of the selected values are shown in S1 Table.

All of the training/validation was run using TensorFlow [83]. The training was performed on Nvidia V100 GPUs with 25 GB memory.

Network evaluation.

The network predicted the values of parameter on the test set that is unseen in the training and validation. We also conducted robustness tests by varying trial numbers (input size).

To evaluate the output of both ANN and traditional tools we used the following metrics (ensuring our results are robust to the choice of performance quantification):

• The mean squared error (MSE): To evaluate parameter estimation accuracy, we calculated a mean squared error between true and estimated model parameter across all agents. Prior to calculating MSE all parameters were normalized, to ensure comparable contribution to MSE across all parameters. Overall loss for a cognitive model (across all parameters) was an average of individual parameter MSE scores. Overall loss for a class of models (e.g., likelihood-tractable models) was an average across all model MSE scores.
• Spearman correlation (ρ): We used Spearman correlation as an additional metric for examining how estimated parameter values relate to true parameter values, with higher Spearman ρ values indicating higher accuracy. We paired Spearman correlations with scatter plots, to visualize patterns in parameter recoverability (e.g., whether there are specific parameter ranges where parameters are more/less recoverable).
• R-Squared (R² or the coefficient of determination): R-Squared represents the proportion of variance in true parameters that can be explained by a linear regression between true and predicted values. It thus indicates the goodness of fit of an ANN model. We calculated an R-Squared score for each individual parameters across all agents and used it as an additional evaluation for how well the data fit the regression model.

Uncertainty estimation.

To compute uncertainty of parameter estimates we have incorporated evidential learning into our method [58]. In the application of evidential learning to continuous regression [58] observed targets follow a Gaussian distribution, characterized by its mean and variance. Conjugate Gaussian prior, normal inverse-gamma distribution, is created by placing a prior on both the mean (Gaussian distribution) and the variance (inverse-gamma distribution). By sampling from this distribution, specific instance of the likelihood function is obtained (based on both mean and the variance). This approach not only aims for accurate target predictions but also takes into account the uncertainty (quantified by the variance term). For more insights and details into evidential learning, refer to the work by [58]. Their research also introduces a regularization term, which is useful for penalizing incorrect evidence and data that falls outside the expected distribution.

For the purpose of visualization (Fig 4), we have created upper and lower bounds of targets by adding/subtracting variance from the predicted target values. We then re-scaled these values by applying the inverse scaler (e.g., from the scaler applied to normalize parameters for network training). This provides a scale-appropriate and more interpretable visualization of parameter recovery and uncertainty for each parameter.

Alternative models.

We have also tested the network with long short-term memory (LSTM) units since LSTM units are more complex and expressive than GRU units; nevertheless they achieved the similar performance as GRU units but are more computationally expensive, and thus we mostly focused on the GRU version of the model. Since LSTM worked, but not better than GRU, the LSTM results are reported in S2–S5 Figs.

Model identification.

The network structure and training process were similar to that of the network used for parameter recovery, with an exception of the output layer that utilized categorical cross-entropy loss and a softmax activation function. The network validation approach was the same as the one we used for parameter recovery (e.g., based on the held-out test set). We also observed a better performance when training with various trial numbers.

Robustness test: Influence of different input trial numbers.

For all robustness experiments, we followed the same training procedures as described previously while varying the training data. The details of training data generation are given below:

Parameter recovery.

We simulated 30,000 training samples with 2000 trials per simulation in the probabilistic reversal learning task. For shorter fixed trial sequence lengths per training samples (e.g., 500), we used the same training set truncated to the first 500 trials. To generate the training data with different trial numbers across training samples, we reused the same training set, with sequences of trials truncated to a given number. There were 6000 training samples of 50, 100, 500, 1000, 1500, and 2000 trials, each.

Model identification.

The process of data generation for model identification robustness checks was similar to parameter recovery. However, we only simulated 500 trials for each model because we found no significant increase in accuracy with higher trial numbers.