Two-armed bandit task.

We present our model in the context of a generic two-armed bandit task widely used decision-making task to study behavior [1, 15]. At each trial of time t, the agent is forced to choose between two possible actions, a (denoted L and R). Next, a stochastic reward, r, is delivered. The reward probability of each action is fixed during a block of B trials and selected randomly at the beginning of each block from a pre-existing setting (see Materials and methods for more details).

t-RNN architecture.

Our model as depicted in Fig 1F, is based on a recurrent neural network (RNN) consisting of fully connected (FC) layers and gated recurrent units (GRU; [25]). At each step, the inputs to the t-RNN are the two observed time-dependent variables, the action a_t and the reward r_t, and it produces two outputs: currents estimate of the agent RL parameters (Θ denotes a vector of RL parameters; subscripts t denotes time-varying RL parameters), and prediction of the future action .

t-RNN training.

The two outputs that t-RNN produces, namely the hidden RL parameters and future action, are processed through different loss terms. The RL parameters are predicted in two stages. First as a multiclass classification task, after employing quantization of each RL parameter to a discretized space (i.e., ; see Materials and methods), and then as a regression problem. It is important to note that quantization was used only to stabilize the network training, while the network’s final output remains continuous. This procedure is based on previous work [26, 27] which has shown that MSE loss is much harder to optimize and often converges to a fixed output that reflects the mean. To ensure that quantization did not affect our performance in any meaningful way, we ran multiple experiments in which we varied the number of the RL parameters quantization and found similar results (see S1 Table).

In the first stage, a_t (coded using one-hot transformation) and r_t is passed through a single layer of 32 GRU cells and the output is projected to m = |Θ| FC layers with ReLU activation to obtain the network intermediate prediction of the categorical class of each RL parameter. We then compute a cross-entropy loss (denoted as ) with regard to the true class for each RL parameter separately.

Next, to allow a continuous recovery of each RL parameter, in the second stage, the network categorical class prediction (i.e., the logits) is passed through m additional regression heads (each FC layer of size 1; each FC layer is responsible for a different RL parameter). We compute the mean-square error (denoted as ) with regard to the true continuous values of each RL parameter separately.

In addition to estimating the parameters, our network also predicts the future action, a_t+1. For this aim, we concatenate the network categorical class predictions (i.e., the logits) together with a_t and r_t and passed this vector through a second single layer of 32 GRU cells. We project the output to an FC layer with SoftMax activation to obtain a probability distribution over possible future actions and computed a cross-entropy loss (denoted as ) with regards to the true future action a_t+1.

In S6 Fig, we provide an additional experiment to examine the necessity of concatenation and its impact on the results. To investigate this, we employed ablation techniques and trained a version of the t-RNN model without concatenation.

We now define the combined loss function used to train the network, (1) where λ^CE and λ^MSE are hyperparameters, set early during the development process such that each loss term is approximately of equal magnitude (see Materials and methods for further details).

Theoretical RL model.

To generate a training set of t-RNN, we simulated Q-learning [29, 30] agents performing a two-armed bandit task. Specifically, for the Q-learning agents, at each experience of time t, the agent updates the Q-values (initialized to zero at the beginning of each block) of the action selected a_t using the obtained reward r_t by, (2) where 0 ≤ α ≤ 1 is a learning-rate free-parameter, which determines the extent to which newly acquired information overrides old information, and δ_TD is the temporal-difference (TD) error that drives learning.

These Q-values are in turn transformed to a stochastic choice-rule policy in accordance with the Boltzmann distribution, (3) where β ≥ 0 is the inverse-temperature free-parameter, which determines the randomness of action selection.

Training t-RNN was done with a synthetic training set of 2,000 simulated Q-learner agents operating with Eqs 2 and 3. Each agent performed a two-armed bandit task for 10 blocks of 100 trials each, and the agents’ RL parameters were sampled from a uniform distribution. Importantly, to stabilize the training procedure, we applied two techniques. First, to allow for continuous recovery of each RL parameter, we performed a quantization preprocessing step, where we discretized each RL parameter into evenly spaced bins, which acted as an intermediate prediction target during training. Second, to facilitate t-RNN’s flexibility to recover time-varying RL parameters, we simulated agents with non-stationary parameters. For each agent, there was a small probability per trial that the agent’s parameters would be re-sampled from the same uniform distribution (see Materials and methods for more details).

For comparison purposes, we included two additional baseline methods for estimating hidden RL parameters. The first baseline method was a stationary Q-learning model (denoted Q-stationary), in which stationary RL parameters were estimated for each agent individually, using a maximum-likelihood estimation. The second baseline method was Bayesian particle filtering (denoted Bayesian) adopted from [28]. Here, underlying RL parameters are assumed to drift slowly according to a Gaussian random walk and are estimated for each agent on a trial-by-trial basis using Bayesian particle filtering (see Materials and methods for more details).

Results.

To validate our method, after training t-RNN, we simulated 30 additional test agents (Materials and methods) on which the network was not trained and recorded the predictions made by t-RNN as well as the two baseline methods. First, we calculated the binary cross-entropy (BCE) between the true actions and predicted actions of each of the three models (stationary Q-learning, Bayesian, and t-RNN). We found that t-RNN had the lowest error compared to the two other models (p < .01 when comparing Q-stationary and t-RNN using Wilcoxon signed rank test, and p < .05 when comparing Bayesian and t-RNN; see BCE_action in Fig 2A;). Second, we calculated the mean-squared error (MSE) between the true and estimated RL parameters and found that t-RNN outperformed the two alternatives methods in terms of MSE_β (p < .01 when comparing Q-stationary and Bayesian to t-RNN using Wilcoxon signed rank test; see Fig 2A). Note that although t-RNN exhibited the lowest MSEα compared to both baseline methods, it was not significant. This finding confirmed that our network can accurately recover hidden RL parameters of test agents operating with various underlying RL parameter dynamics (see S2 Table for raw results).

Download:

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

Fig 2. Validating the ability of t-RNN to predict both agents’ actions and to infer RL parameters of simulated data.

(A) Action prediction (measured with binary cross entropy; BCE) and RL parameter estimation error (measured with mean-square error; MSE) averaged across 30 artificial test agents (black lines indicate s.e.m). We found that our method exceeded alternative methods (Q-stationary [Maximum-likelihood], green; Bayesian [particle filtering], yellow) in both action prediction (left) and parameter estimation (middle and right). (B) Recovery of theoretical RL α learning-rate (left) and β inverse-temperature (right) parameters of three example simulated agents with different RL parameter trajectories. These trajectories included abruptly changing (top), random-walk drift (middle), and stationary (bottom) RL parameters. Overall, we found that our network (blue dotted line) was able to effectively recover the true dynamic RL parameters (red solid line). This demonstrates the versatile ability of t-RNN to accurately capture a wide range of time-varying RL parameter dynamics. Significant effect at * p < .05, ** p < .01; n.s., a nonsignificant effect (Wilcoxon signed rank test).

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

To further examine and illustrate the t-RNN capabilities, we plotted (see Fig 2B) the true and recovered RL parameters of three types of test agents: abruptly changing RL parameters (top), gradually changing RL parameters (middle), and agents with stationary RL parameters (bottom). Visual inspection of Fig 2B indicates that t-RNN generated predictions that closely matched the true RL parameter trajectories. Overall, these results indicate t-RNN flexible ability to recover underlying RL parameters while keeping high predictive accuracy. Additionally, we examine the recovery of agents simulated in a volatile environment, where we observed that t-RNN maintains superior performance compared to both baseline methods (see S3 Table). We also examine the ability of t-RNN to estimate stable parameters compared to the stationary Q-learning method as a function of the number of observations (see S1 Text). Our findings demonstrate that t-RNN performs well for stable parameters and is comparable to the stationary method when more than one parameter is included in the theoretical model. Moreover, for a low number of observations (below 50 trials), t-RNN outperforms the stationary estimation.

Theoretical RL model.

To model subjects’ behavior (and to train t-RNN), we consider the same RL Q-learning model as mentioned above (see Eqs 2 and 3) with an additional action preservation free-parameter added to the SoftMax choice-rule policy, (4) where −0.5 ≤ κ ≤ 0.5 is an action preservation free-parameter determining the agent’s tendency to stick with (or switch) a previous action regardless of the reward [31].

Results.

Following training of the t-RNN with a synthetic training set of Q-learner agents using Eqs 2 and 4, we fixed the network’s weights and recorded the predictions made by the t-RNN on the behavioral dataset of humans who performed a two-armed bandit task [9]. Additionally, we fitted a stationary Q-learning with preservation model (denoted QP-stationary) to each subject, which estimated stationary RL parameters via maximum-likelihood estimation, as well as a Bayesian model fitted with particle filtering. Furthermore, we fitted a data-driven RNN (d-RNN) to subject behavior, as described in [9], which solely produced action predictions (see Materials and methods for more details).

As expected, we found that t-RNN was able to maintain a higher action prediction of human behavior compared to the QP-stationary model but not compared to d-RNN (see Fig 1B; p < .05 when comparing QP-stationary and t-RNN using Wilcoxon signed rank test, and p < .0005 when comparing t-RNN and d-RNN). To further explore this result, we divided the subjects according to their diagnostic groups and examined the action predictions within each group. We found that t-RNN superiority over QP-stationary and Bayesian models was mainly due to the improved action prediction of bipolar and depressed individuals (p < .01 comparing QP-stationary and Bayesian to t-RNN in the bipolar group; p < .01 comparing QP-stationary to t-RNN and p < .001 comparing Bayesian to t-RNN in the depressed group; Wilcoxon signed rank test), while there was no significant difference in its prediction accuracy for healthy controls (p = 0.36 comparing QP-stationary to t-RNN; p = 0.07 comparing Bayesian to t-RNN; see Fig 3A and S4 Table for the raw results).

Download:

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

Fig 3. Application of t-RNN to behavioral data of psychiatric and non-psychiatric individuals [9].

(A) Action prediction (measured with BCE; black lines indicate s.e.m) divided by diagnostic labels. Bipolar and depressed subjects at the group level are explained significantly better by t-RNN (blue) compared to the QP-stationary model (green) that assumes that the RL parameters are fixed, as well as to the Bayesian model (yellow). (B) Boxplot of the time-varying κ preservation parameter estimates by t-RNN across the three diagnostic groups (middle black solid lines denote the median; light dots indicate a single trial estimate of the κ preservation parameter). Result suggests higher volatility in the clinical groups mainly the bipolar group, compared with the healthy group. (C) Distribution of the Pearson correlation between stay probability (calculated using moving average of 10 trials) and time-varying κ preservation parameter produced by t-RNN for each subject individually. Result shows a strong relation between t-RNN time-varying RL κ parameters estimation and moving average of the stay probability. (D) Sequence of selected actions of two example subjects, with bipolar on the left and depressed on the right. The top panel shows the action prediction, and the bottom panel shows the RL κ parameter estimation produced by t-RNN (blue) and QP-stationary model (green) for a 50-trial segment (see S4 Fig for all trials). Both subjects seemed to switch between selecting the same action for several trials to repeatedly alternating between both actions every trial (red dots). We found that the t-RNN model was able to detect this visually apparent shift in behavior, estimating a change in the κ action perseveration parameter. In contrast, the QP-stationary model was not able to capture this transition. (E) Trial-by-trial κ preservation parameter estimates by t-RNN averaged over the first 100 trials of each block and for each diagnostic group separately (shaded area signifies the s.e.m). In all three groups, subjects show a steady increase in their tendency to perseverate their actions as the block progresses. Significant effect at * p < .05, ** p < .01, *** p < .001; n.s., a nonsignificant effect (Wilcoxon signed rank test).

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

We next turn to explore the dynamic estimation of the theoretical parameters. First, we assessed the volatility in κ preservation estimates estimated trial-by-trial using t-RNN. Our results revealed that the bipolar group exhibited lower values and greater volatility in these estimates compared to the healthy group (M = 0.15, SD = 0.30 for the bipolar group, M = 0.35, SD = 0.14 for the healthy group; see Fig 3B). This finding might correspond with the clinical nature of the disorder where volatility is a main feature [23]. However, for the aim of the current study, this is mainly a proof-of-concept to the t-RNN’s ability to inform us regarding volatility in underlying parameters, while also providing better action predictions compared with a stationary QP-stationary model. In S5 Table, we show additional distribution of all RL parameters produced by t-RNN.

Second, to gain insight into the association between RL parameter estimation and observed behavior, we selected two clinical participants with bipolar disorder and depression whose behavior displayed clear transitions between high (sticking with the same action) and low (switching the action selection every trial) levels of perseveration. We plotted the sequence of selected actions alongside the action predictions and inferred κ estimates generated by t-RNN, as well as the predictions of the QP-stationary model for a 50-trial segment (see Fig 3D and S4 Fig for all trials). We found that the observed change in action perseveration was closely coupled with a change in the κ parameter, and the t-RNN model accurately detected the visually apparent shift in behavior. In contrast, the stationary model, which assumes fixed RL parameters, was unable to capture this transition, resulting in lower action prediction accuracy. These results highlight the utility of t-RNN in capturing dynamic changes in RL parameters and their association with observed behavior.

Third, to further support that t-RNN parameters estimation indeed corresponds well with changes in observed behavior across all individuals, we turn to a well-established and easy-to-interpret model-agnostic measurement and calculated a moving average of the stay probabilities (probability of repeating the previous action; window size of 10 trials). If indeed the t-RNN κ perseverate dynamic parameter is doing a good job of capturing fluctuations in perseveration, we should expect a high correlation with the moving average model-agnostic estimate. We calculated a Pearson correlation between the moving average and t-RNN κ perseverate trial-by-trial estimation for each individual (see Fig 3C). Overall, we found that the average Pearson across individuals was very high (Median = 0.64; see S1 Fig for further detailed examples). We also included a similar analysis with the β inverse-temperature parameter (see S2 and S3 Figs).

Finally, we plotted the trial-by-trial κ perseveration average for each group, showing that overall, across groups individuals increased their tendency to perseverate with the previous action selection which might reflect an overall reduction motivation as the block progresses (see Fig 3E). Overall, we were able to show that t-RNN is able to produce trial-by-trial estimates of RL parameters, while also keeping higher predictive accuracy compared with QP-stationary model. In the current dataset, we provided a proof-of-concept on how t-RNN might be able to inform of higher volatility in RL parameters in clinical psychiatric groups.

Theoretical RL model.

To model subjects’ behavior, we consider the same model hybrid exploration model presented in Gershman, 2018 [22]. The learning model included a Kalman filtering algorithm [15, 32] where agents recursively compute the posterior mean and variance for each action by, (5) (6) where Q_t+1(a) and is the posterior mean and variance estimates over the value of action a, respectively and K_t is the Kalman gain given by, (7) where is the error variance. At the beginning of each block the initial values estimate for both actions were set to the prior mean, namely, Q₀ (a) = 0, and (see Materials and methods for more details).

Action values estimate were transformed to action policy with a hybrid of random and directed exploration [24] choice-rule, (8) where Φ is the cumulative distribution function (CDF) of the standard Gaussian distribution (i.e., probit), 0 ≤ β ≤ 4 is a free-parameter that controls the contributions of random exploration (similar in spirit to the inverse-temperature parameter of the SoftMax choice-rule of Eq 3) and 0 ≤ γ ≤ 1 is a free-parameter that controls the contributions of directed exploration.

This model can also be interpreted as a combination of Thompson sampling [33] and UCB policies [34], where the first term in the brackets promotes a random exploration that is proportional to the difference in the value estimates, Q_t (L) − Q_t (R) together with the total uncertainty much like Thompson sampling. The second term in the brackets, σ_t (L) − σ_t (R), is the so-called relative uncertainty, synonyms with UCB policy which promotes a directed exploration toward the less certain option (i.e., information bonus).

Results.

After the training of the t-RNN with a synthetic training set of Q-learner agents using Eqs 5, 6, 7 and 8, we fixed the network’s weights and recorded the predictions made by the t-RNN on the behavioral dataset of humans who performed a two-armed bandit task [22]. Additionally, we fitted a stationary hybrid exploration model, a Bayesian model, and a d-RNN to subject behavior (Materials and methods). We found significantly higher action prediction for t-RNN compared to both the stationary hybrid exploration model (p < .05; Wilcoxon signed rank test) and the Bayesian model (p < .01), but no significant difference in prediction accuracy when comparing t-RNN to d-RNN (p = 0.12; see Fig 4A and S6 Table for raw results).

Download:

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

Fig 4. Application of t-RNN to investigate changes in random and directed exploration [22].

(A) Action prediction (measured with binary cross entropy; BCE; black lines indicate s.e.m). We found a significant improvement in the action prediction of the t-RNN model (blue) compared to the hybrid exploration model (green; * p < .05) and the Bayesian (yellow; ** p < .01). (B) Example of a single instantiation of ungreedy action. The top panel depicts the sequence of chosen actions, obtain rewards, and t-RNN action predictions as red dots, red numbers, and blue solid lines, respectively. The subject consecutively chooses the higher value action and suddenly switches and chooses the lower value action (denoted with a downward arrow). The bottom panel depicts t-RNN β random exploration parameter estimation, where it estimates a step decrease in the parameter after the ungreedy action selection. (C) Dynamics of the random exploration parameter averaged across all instantiations of ungreedy action. t-RNN estimates a sharp and sudden decrease in the random exploration parameter after an ungreedy action is taken. (D) Correct choice rate of individuals’ actions divided by blocks where the true value difference is high (≥ 19; easy blocks; purple) and blocks where the true value difference is low (≤ 1; hard blocks; grey). Subjects are more accurate when it is easy to distinguish the correct action compared to when it is hard. (E) Dynamics of the β random exploration parameter estimation produced by t-RNN average overall easy (purple) and hard (grey) blocks. Random exploration increases in hard blocks (grey) where it is difficult to tell apart which action is best, whereas in the easy blocks (purple) random exploration decreases as subjects become more deterministic in their action selection. (F) Dynamics of the γ directed exploration parameter average overall easy (purple) and hard (grey) blocks. Directed exploration is high at the beginning of the block where uncertainty is equal in both conditions, however only in the easy blocks (purple) where it is easy to differentiate between the better/worst arm, directed exploration decreases as the block evolves. In panel C, E and F solid line/dots signifies the mean estimation and the shaded area correspond to s.e.m.

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

To investigate the trial-by-trial dynamics of the t-RNN’s theoretical parameter estimation, we conducted a manual search of the data to identify an instance in which an individual made an ungreedy action. Specifically, we located a trial where an individual chose option L despite having been rewarded for option R for the previous four consecutive trials. We observed a sharp decrease in the t-RNN’s β random exploration parameter estimation, indicating its sensitivity to changes in observed behavior (see Fig 4B). To determine whether this pattern generalizes across individuals and trials, we aggregated all trials (a total of 32 trials) in which individuals selected the better rewarding arm for three or more consecutive trials, were then rewarded, and subsequently switched to the alternative, poorer arm. Our analysis revealed an overall decrease in the β random exploration parameter produced by t-RNN in these cases (mean ± s.e.m. before: 2.1 ± 0.079, after: 1.40 ± 0.080; see Fig 4C). In S5 Fig, we provide a similar example for the γ directed exploration parameter.

Next, we investigated the overall dynamics of the two exploration parameters (random and directed) within blocks. Specifically, we examined blocks where the true expected value difference between the two arms was above or equal to 19 (referred to as “easy” blocks; total of 91 blocks) and blocks where the true expected value difference between the two arms was below or equal to 1 (referred to as “hard” blocks; total of 110 blocks). It is worth noting that participants were not informed whether a block was easy or hard and were required to start with a similar strategy in the first trial where the difficulty of the block was unknown. We hypothesized that participants would change their exploration strategy across time, and predicted that in easy blocks, individuals would reduce random exploration since the best action is easy to detect, whereas, in hard blocks, individuals would show a higher tendency for random exploration (indicating either noisy behavior due to block difficulty or a desire to test the arms in a more heterogeneous fashion to gain more information, where the information gain is relatively small for the participant).

We found that participants had a higher correct choice rate in easy blocks compared to hard blocks (mean ± s.e.m. easy: 0.90 ± 0.007, hard: 0.56 ± 0.003; see Fig 4D). Consistent with our hypothesis, t-RNN estimated an increase in random exploration (i.e., a decrease in the β parameter) across hard blocks and a decrease across easy blocks (see Fig 4E). We also observed the same pattern for the γ directed exploration, where overall directed exploration was high at the beginning of the block, possibly reflecting the individual’s lack of knowledge about the environment. However, there was a selective decrease in directed exploration only for easy blocks, where the best action was easy to locate (see Fig 4F).

Our results indicate that individuals modify their exploration strategies based on both the block difficulty and phase, as evidenced by changes in both random and directed exploration parameters. Importantly, t-RNN was successful in capturing these nuanced changes in exploration strategies, thus providing a second proof of the utility of our method.

Task structure.

Agents performed a two-armed bandit task for 10 blocks of 100 trials each. At each trial the agent received a stochastic binary reward of r = 1 for rewarded trials, or r = 0 otherwise. The reward probability of each action is fixed during a block of 100 trials and is selected randomly at the beginning of each block from three possible Bernoulli distribution: .

t-RNN training.

Training t-RNN was done with a synthetic training set of 2,000 simulated Q-learners agents operating with Eqs 2 and 3. Each agent performed a two-armed bandit task for 10 blocks of 100 trials each. Each agents RL parameters were sampled from a uniform distribution of α ∼ U(0, 1) and β ∼ U(0, 10). In the quantization preprocessing step, we discretized the α parameter into five evenly spaced bins (of size 0.2), and the β parameter into five evenly spaced bins (of size 2).

Importantly, to facilitate t-RNN flexibility to recover time-varying RL parameters, we simulated agents with non-stationary parameters. For each agent, there was a small probability of (p = 0.005) per trial that the agent’s parameters will be re-sampled from the same uniform distribution. To keep the set from changing too rapidly, the overall number of trials in which the parameters were re-sampled was limited to four, and after each trial of parameters re-sampling the probability of another change was fixed to zero for 100 trials.

The combined loss function (Eq 1) included the following λ hyperparameters values: , , , , so that each loss term is approximately of equal magnitude.

Stationary Q-learning.

Here, RL parameters were estimated for each agent individually, using a maximum-likelihood estimation approach (Python SciPy [49] minimized function with L-BFGS-B optimization, and five different starting search locations to avoid local maxima). This approach allowed for an estimation of a single and fixed set of RL parameters for each agent [1].

Bayesian particle filtering.

Here, underlying RL parameters are assumed to drift slowly according to a Gaussian random walk, (9) and are estimated for each agent on a trial-by-trial basis using Bayesian particle filtering. We implemented the same model as described by Samejima et al., 2003 [28]. To recover the simulated data, we set 1,000 particles to approximate the distribution of two hidden RL parameters {α, β}. Initial distribution of the particles were set to be Gaussians, with means of {0, 0} and variances of {3, 1}, respectively. We set the hyper-prior for the dynamics of α and β as σ_α = 0.1 and σ_β = 0.05, respectively (we also tested using hyper-prior as low as σ_α/σ_β = 0.005/0.0005 and us high as σ_α/σ_β = 0.05/0.005 and σ_α/σ_β = 0.1/0.5, these setting resulted in poorer results). We also clipped the β to range between (0,10) to match the estimation range of other models.

Test agent.

In order to validate and compare the performance of t-RNN against baseline methods, we simulated a test set of 30 artificial Q-learning agents (Eqs 2 and 3). Each agent was simulated on the same task structure, consisting of 10 blocks of 100 trials, and his RL parameters were sampled from a uniform distribution of α ∼ U(0, 1) and β ∼ U(0, 10). The test set included three groups of 10 agents each. The first group had stationary RL parameters, while the second group had RL parameters that abruptly changed with a small probability for each trial (p = 0.005). The third group had RL parameters that drifted slowly according to a Gaussian random walk (Eq 9) with sigma values of σ_α = 0.1 and σ_β = 0.05.

Behavioral dataset and task structure [9].

The behavioral dataset includes N = 101 individuals that performed the bandit task for 12 blocks, with a free-choice rate of 40 second per block. At each trial subjects received a stochastic binary reward of r = 1 for rewarded trials, or r = 0 otherwise. The reward probability of each action was fixed during a block and selected at the beginning of each block from six possible settings:. Each possible reward setting was repeated twice, for a total of 12 blocks. The data included task trials of 33 individuals diagnosed with bipolar disorder, 34 with depression and 34 individuals who were matched healthy controls (see Dezfouli et al., 2019 [9] for further details).

t-RNN training.

Training t-RNN was done with a synthetic training set of 2,000 simulated Q-learners agents operating with Eqs 2 and 4. Each agent performed a two-armed bandit task for 12 blocks of 100 trials each. Each agents RL parameters were sampled from a uniform distribution of α ∼ U(0, 0.2), β ∼ U(0, 10) and κ ∼ U(−0.5, 0.5). In the quantization preprocessing step, we discretized the α parameter into three evenly spaced bins (of size 0.066), the β parameter into five evenly spaced bins (of size 2) and the κ parameter into five evenly spaced bins (of size 0.2). Again, to facilitate t-RNN flexibility to recover time-varying RL parameters, we simulated agents with non-stationary parameters, with the same setting as described above.

The combined loss function (Eq 1) included the following λ hyperparameters values: , , , , , , so that each loss term is approximately of equal magnitude.

d-RNN model.

Data driven RNN was implemented following the same RNN model as described by Dezfouli et al., 2019 [9] (term here d-RNN). d-RNN is composed of a single GRU layer of 10 hidden units. The inputs to the d-RNN are the current action a_t and the reward r_t received after taking that action. The outputs are probabilities over the action of the next step a_t+1. We train the d-RNN with and Adam optimizer, learning rate of 0.001 and batch size of 1000. We assess performance with the same leave-one-out cross-validation schema depicted in [9]. At each CV round one of the subjects was withheld and d-RNN was train using data from the remaining subjects. Importantly, this procedure is done separately for each diagnostic group. For example, in order to predict the actions of a withheld bipolar subject, we use data from only the remaining bipolar subjects to train the d-RNN model.

Stationary Q-learning with preservation (QP-stationary).

Here, stationary RL parameters are estimated for each subject individually, with maximum-likelihood estimation (Python SciPy [49] minimized function with L-BFGS-B optimization, and five different starting search locations to avoid local maxima).

Bayesian particle filtering.

For each individual, we set 1000 particles to approximate the distribution of three hidden RL parameters {α, β, κ}, Gaussian with means of {−2, 1, 0} and variances of {1.5, 1, 1.5}, respectively. The hyper-prior for the dynamics of σ_α = 0.05, σ_β = 0.005, and σ_κ = 0.05. We clipped the beta to range between (0,10) to match the estimation range of other models.

Behavioral dataset and task structure [22].

The behavioral dataset includes N = 44 individuals that performed the bandit task for 20 blocks of 10 trials each. At each trial subjects received a stochastic reward drawn from a Gaussian distribution with means of μ_L,μ_R (for actions L and R respectively) and a variance (i.e., error variance) of . The mean reward probability of each action was fixed during a block and was drawn at the beginning of each block from a Gaussian distribution, with mean 0 and variance 100 (see Gershman, 2018 [22] for further details).

t-RNN training.

Training t-RNN was done with a synthetic training set of 10,000 simulated Q-learners agents operating with Eqs 5, 6, 7 and 8. Each agent performed a two-armed bandit task for 20 blocks of 10 trials each. Each agents RL parameters were sampled from a uniform distribution of β ∼ U(0, 4) and γ ∼ U(0, 1). In the quantization preprocessing step, we discretized the β parameter into five evenly spaced bins (of size 0.8), the γ parameter into five evenly spaced bins (of size 0.2). Again, to facilitate t-RNN flexibility to recover time-varying RL parameters, we simulated agents with non-stationary parameters, with the same setting as described above. For each agent, there was a small probability of (p = 0.02) per trial that the agent’s parameters will be re-sampled from the same uniform distribution. The overall number of trials in which the parameters were re-sampled was limited to four, and after each trial of parameters re-sampling the probability of another change was fixed to zero for 10 trials.

The combined loss function (Eq 1) included the following λ hyperparameters values: , , , , so that each loss term is approximately of equal magnitude.

d-RNN model.

Similar to the one described above, with the two differences being are that the batch size was 200 and the LOOCV schema was performed over all 44 individuals.

Stationary hybrid exploration model.

We adopted a similar model as described in [22], in which we estimated two stationary (i.e., fixed) parameters for each individual separately. Namely, a β random exploration parameter and a γ directed exploration parameter using maximum-likelihood estimation.

Bayesian particle filtering.

For each individual, we set 1,000 particles to approximate the distribution of two hidden RL parameters {β, γ}. Initial distribution of the particles were set to be Gaussians, with means of {0, −2} and variances of {1, 1}, respectively. The hyper-prior for the dynamics of β was σ_β = 0.1, and for γ was σ_γ = 0.1. We clipped the β to range between (0,4) to match the estimation range of other models.