A transparent model for learning in volatile environments

Sound principles of statistical inference dictate that uncertainty shapes learning. In this work, we revisit the question of learning in volatile environments, in which both the first and second-order statistics of environments dynamically evolve over time. We propose a new model, the volatile Kalman filter (VKF), which is based on a tractable state-space model of uncertainty and extends the Kalman filter algorithm to volatile environments. Algorithmically, the proposed model is simpler and more transparent than existing models, and encompasses the Kalman filter as a special case. Specifically, in addition to the error-correcting rule of Kalman filter for learning observations, the VKF learns volatility according to a second error-correcting rule. These dual updates echo and contextualize classical psychological models of learning, in particular hybrid accounts of Pearce-Hall and Rescorla-Wagner. At the computational level, compared with existing models, the VKF is more accurate, particularly in estimating volatility, as it is based on more faithful approximations to the exact inference. Accordingly, when fit to empirical data, the VKF is better behaved than alternatives and better captures human choice data in a probabilistic learning task. The proposed model provides a transparent and coherent account of learning in stable or volatile environments and has implications for decision neuroscience research.


Introduction
Our decisions are guided by our ability to associate environmental cues with the outcomes of our chosen actions. Accordingly, a central theoretical and empirical question in behavioral psychology and neuroscience has long been how humans and other animals learn associations between cues and outcomes. According to both psychological and normative models of learning [1][2][3][4], when animals observe pairings between cues and outcomes, they update their belief about the value of the cues in proportion to prediction errors, the difference between the expected and observed outcomes.
Importantly, the degree of this updating depends on a stepsize or learning rate parameter. Although some accounts take this as a free parameter, analyses based on statistical inference, such as the Kalman filter [5], instead demonstrate that the learning rate should in principle depend on the learner's uncertainty.
The dynamics of uncertainty -and hence, of learning rates -then depend on the assumed or learned dynamics of the environment. For instance, the Kalman filter is derived assuming that the true associations fluctuate randomly, but at a known, constant speed. In this case the asymptotic uncertainty, and learning rate, are determined by how quickly the associations fluctuate and how noisily they are observed. However, in volatile environments, in which the speed by which true associations change might itself be changing, uncertainty (and learning rates) should fluctuate up and down according to how quickly the environment is changing [6,7]. This normative analysis parallels classical psychological theories, such as the Pearce-Hall model [3], which posit that surprising outcomes increase the learning rate while expected ones decrease it. Those models measure surprise by the absolute value of the discrepancy between the actual outcome and the expected value, i.e. the unsigned prediction error [3].
Behavioral studies have also shown that human learning in volatile environments is consistent with the predictions of this class of models [6,8]: learning rates fluctuate with the volatility of the environment. There is also evidence for a neural substrate for such dynamic learning rates: for instance, neuroimaging studies have shown that the dorsal anterior cingulate cortex covaries with the optimal learning rate [6] and recent work suggests a mechanistic model of adaptive learning rate [9]. Theoretical and empirical work also suggests that neuromodulatory systems, particularly acetylcholine, norepinephrine and serotonin, might be involved in encoding uncertainty signals necessary for computing learning rate in stable and volatile environments, respectively [7,8,[10][11][12]. Finally, various approximate learning algorithms have been fruitful for studying individual differences in learning [13][14][15][16].
The theoretical studies establish the question of why the learning rate should be dynamically adjusted, and the empirical studies provide evidence that it does so. However, a complete understanding of a learning system requires understanding of how these theories could be realized at the process or algorithmic level [17]. This is as yet much less clear: as discussed below, statistically grounded theories of learning under volatility tend to be somewhat impractical and opaque. Furthermore, while their general analogy with the more rough and ready psychological theories (like Pearce-Hall) seems clear, there is not a direct mapping comparable to the way for example the Kalman filter encompasses and rationalizes the classical Rescorla-Wagner theory.
A fully Bayesian account of learning turns on two different aspects. The first is a generative model, that is a set of explicit probabilistic assumptions about how the environment evolves and generates outcomes. In situations more complicated than the smooth world of the Kalman filter, however, exact inference is generally intractable and additional approximations are required to achieve a practical, algorithmic-or process-level inference model. An influential model, proposed by Behrens and colleagues [6], comprises a hierarchical generative model for learning, in which a variable governing the observed associations fluctuates with a speed determined by a second higher-level, fluctuating variable. Although this model has been conceptually influential, it lacks a tractable and biologically-plausible inference model at the algorithmic level: instead, its predictions were simulated by brute-force integration. Another model, called the hierarchical Gaussian filter (HGF) [8,18], proposed a related generative model also consisting of multiple levels of hierarchy, in which the speed of change at each level is determined by the preceding level. An approximate inference rule was also proposed, based on the so-called mean-field approximation [19], that the posterior at each time point is fully independent of all other time points. Although this model provides a biologically-plausible algorithm for learning, the assumption that the probability distribution at each time point is fully independent of all the other time points poses important theoretical and practical challenges, particularly in estimating volatility and learning rate. Furthermore, the resulting learning algorithm is much more convoluted than classical models, such as the Pearce-Hall model or the Kalman filter.
Here we revisit the question of approximate inference in a hierarchical model of learning in volatile environments. Rather than assuming a mean-field approximation, we show the importance of preserving the key statistics reflecting volatility of the environment, in particular temporal autocorrelations. Informed by this reasoning, we propose a novel model for learning in volatile environments that is conceptually aligned with that of Behrens et al. [6] and the HGF but more clearly related to the Kalman filter and to psychological models such as Pearce-Hall. In particular, the proposed model is based on the hierarchical structure of Behrens et al. [6] and principles of probabilistic graphical models [20]. The resulting algorithm, called VKF, is a generalization of Kalman filter algorithm to volatile environments and resembles models that hybridize the error-driven learning from the Rescorla-Wagner model and the Kalman filter with Pearce-Hall's dynamic learning rate (as proposed by different authors, for example by Li et al. and Le Pelley [21,22]). Notably, in volatile environments, the learning rate fluctuates with larger and smaller than expected prediction errors, as suggested by models such as Pearce-Hall.
In the next section, we review the Kalman filter algorithm and present the generative model underlying the VKF and the resulting learning algorithm. The full formal treatment is given in the S1 Appendix. Next, we show that the proposed model outperforms existing models in predicting empirical data.

Theory
Kalman filter for tracking in environments with stable dynamics The Kalman filter is the cornerstone of statistical tracking theories, with widespread applications in many technological and scientific domains including psychology and neuroscience [4,7,23,24]. The Kalman filter corresponds to optimal statistical inference for a particular class of linear state space environments with Gaussian dynamics. In those environments, the hidden state of the environment is gradually changing across time (according to Gaussian diffusion) and the learner receives an outcome on each time depending on the current value of the state (plus Gaussian noise). In these circumstances, the posterior distribution over the hidden state is itself Gaussian, thus tracking it amounts to maintaining two summary statistics: a mean and a variance.
Formally, consider the simplest case of prediction: that of tracking a noisy, fluctuating reward (e.g. that associated with a particular cue or action), whose magnitude " is observed on each trial . Assume a state space model in which, on trial t, the hidden state of the environment, " (the true mean reward), is equal to its previous state "&' plus some process noise which depends on the noise parameters and the posterior variance "&' on the previous trial. Note that " < 1 in all trials and it is larger for larger values of . On every trial, the posterior variance also gets updated: Note that although it is not required for prediction (that is for updating " ), it is also possible to compute the auto-correlation or posterior covariance (given observations ' ,…, " ) between consecutive states, "&'," = cov[ "&' , " ], which is given by This equation indicates that when the Kalman gain is relatively small, the autocorrelation is large, which means that information transmitted by observing a new outcome is expected to be quite small. We will see in the next section that this autocorrelation plays an important role in inference in volatile environments.

VKF: A novel algorithm for tracking in volatile environments
We next consider a volatile environment in which the dynamics of the environment might themselves change. In the language of the state space model presented above, the process noise dynamically changes. Thus, the variance of the process noise ( " in Equation 1) is a stochastic variable changing with time ( Figure 1A). To build a generative model, we need to make some assumptions about the dynamics of this variable. Our approach here is essentially the same as that taken by Smith and Miller [25] and by Gamerman et al. [26] (see also West et al. [27]). Consider a problem in which the process variance dynamically changes. In the previous section, we saw that the state " diffused according to additive noise. Because variances like are constrained to be positive, it makes sense to instead assume their diffusion noise is multiplicative to preserve this invariant. Therefore, we assume that the current value of precision (inverse variance), " , is given by its previous value multiplied by some independent noise. Formally, the current state of " is given by where " is an independent random variable on trial , which is distributed according to a rescaled beta distribution (as detailed in S1 Appendix), such that the mean of " is 1 (that is, conditional expectation of " is equal to "&' ) but has spread controlled by a free parameter 0 < < 1. The value of noise, " , is always positive and is smaller than 1 − &' . Therefore, while "on average", " is equal to "&' , " can be smaller than "&' , or larger by a factor up to 1 − &' . Thus, the higher the parameter , the faster the diffusion.
To build up the full model, first consider a simplified case in which the latent state "~" , " &' is directly observed at each step and, further, has some known mean " . This subproblem provides the foundation for the inference in the full model, using estimates of these quantities.
It can be formally shown (see S1 Appendix) that in this case, estimating the process variance, i.e. the posterior over " at each step given all previous observations, is tractable. Specifically, it can be shown that the posterior distribution over " takes the form of a Gamma distribution, whose inverse mean, " , is updated according to the observed sample variance: Note that this equation amounts to an error-correcting rule for updating " , in which the error is given by " − " G − "&' : the difference between the observed and expected squared prediction error.
Now we are ready to build the full generative model in which both " and " dynamically evolve over time ( Figure 1). In this model, as before, the observation, "~( " , ω), follows a Gaussian distribution with the mean given by the hidden variable " , which itself is given by " = "&' , " &' : a diffusion whose precision (i.e. inverse variance) is given by another dynamic random variable, " . The value of " in turn depends on its previous value multiplied by some noise that finally depends on parameter according to Equation 2. Therefore, this generative model consists of two chains of random variables, which are hierarchically connected to each other. Unlike each of those two chains separately, however, when they are conjoined in this hierarchical model, inference is not tractable and therefore we need some approximation. We use structural variational inference for approximate inference in this model [28,29]. This technique assumes a factorized approximate posterior distribution and minimizes the mismatch between this approximate posterior and the true posterior using principle of variational inference.
The resulting learning algorithm is very similar to the Kalman filtering algorithm and encompass Kalman filter as a special case. Importantly, the new algorithm also tracks volatility on every trial, denoted by " , which is defined as the inverse of expected value of " . Therefore, we call the new algorithm the "volatile Kalman filter". In this algorithm, the update rule for the posterior mean " and variance " over where the expectation should be taken under the approximate posterior over "&' and " . Therefore, the volatility update rule takes a form of error correcting, in which the error is given by with the noise parameter, , as the step size. Thus, the higher the noise parameter, the higher the speed of volatility update is. Therefore, we call the "volatility update rate". Also, note that the expectation in this equation depends on the autocorrelation.
It is then possible to write " − "&' G in terms of variance and covariance of "&' and " to obtain Equation 8 below and complete the VKF learning algorithm: where is the constant variance of observation noise. In addition to the volatility update parameter, , which is constrained in the unit range, this algorithm depends on another parameter, K > 0, which is the initial value of volatility. Notably, the Kalman filter algorithm is a special case of the VKF in which = 0 and the process variance is equal to K on all trials. In the next section, we test this model with synthetic and empirical datasets. influence. On each trial, , outcome, " , is generated based on a Gaussian distribution, its mean is given by the hidden random variable " , and its variance is given by the constant parameter . This variable is itself generated according to another Gaussian distribution, which its mean is given by "&' , and its variance is given by another hidden random variable, " &' . This variable is itself generated based on its value on the previous trial, "&' , multiplied by some positive noise distributed according to a scaled Beta distribution governed by the parameter . The inverse mean of this variable on the first trial is assumed to be given by another constant parameter, K . See Equations 1-2 for further explanation.

Binary VKF
We have also developed a binary VKF for situations in which observations are binary. The generative model of the binary VKF is the same as that of VKF with the only difference that binary outcomes are generated according to Bernoulli distribution with the parameter given by where ( " ) is the sigmoid function, which maps the normally distributed variable " to the unit range.
Further approximation is required to make inference here, because observations are not normally distributed and Equation 1 does not hold. For binary VKF, we assumed a constant variance and employed moment matching (which is sometimes called assumed density filtering) to obtain the mean. The resulting algorithm is then very similar to the original VKF with the only difference that the update rule for the mean (replacing Equation 5) is slightly different ).

Equation 9
The update rules for the Kalman gain, variance, covariance and the volatility are the same as those for the VKF (Equations 4,6-8). Note that the outcome " is binary here and therefore " − ( "&' ) is always between -1 and +1.

Simulation analyses
We first applied the VKF to a typical volatility tracking task (similar to [6,8]). In this task, observations are drawn from a normal distribution whose mean is given by the hidden state of the environment. The hidden state is constant +1 or -1. Critically, the hidden state was reversed occasionally. The frequency of such reversals itself changes over time; therefore, the task consists of blocks of stable and volatile conditions. We investigate performance of VKF in this task because its variants have been used for studying volatility learning in humans [6,8]. In particular, it has been shown that humans learning rate is higher in the volatile condition. Figure 2 shows the performance of VKF in this task. As this figure shows, the VKF tracks the hidden state very well and its learning rate, " , is higher in the volatile condition.
Furthermore, the volatility signal increases when there is a dramatic change in the environment. Note that as for other models previously fit to tasks of this sort, the switching dynamics at both levels of hidden state here are not the same as the random walk generative dynamics assumed by our model. These substitutions demonstrate that the principles relating volatility to learning are general, and so the resulting models are robust to this sort of variation in the true generative dynamics. In another simulation analysis, we studied the performance of the binary VKF in a similar task, but now with binary observations, which were used in previous studies [6,8]. Here, observations are drawn from a Bernoulli distribution whose mean is given by the hidden state of the environment. The hidden state is a constant probability 0.8 or 0.2, except that it is reversed occasionally. Figure 3 shows the performance of the binary VKF model in this task. As this figure shows, predictions of the binary VKF match with the hidden state. Furthermore, the volatility signal increases when there is a dramatic change in the environment, and the learning rate, " , is higher in the volatile condition. We then compared VKF and HGF (with three levels), as the latter is the most commonly used algorithm for learning in volatile environments. These two models differ in two ways. First, although both models employ the same conditional dependency assumptions in their generative models, the distribution over the volatility variable is different. The HGF assumes normal distribution at both levels, and therefore exponential transformations are needed to ensure variances are positive, which requires further approximation for inference. Our model of dynamic variance (with multiplicative rather than additive Gaussian diffusion) is tractable. Another important difference between the two models is in the assumptions about the approximate posterior. The approximate inference model of HGF assumes that the posterior at each time point is fully independent of all the other time points. We argue that this assumption poses critical challenges, particularly because the autocorrelation between hidden states " and "&' is neglected. Note that when the process noise variance, , is constant, the Kalman filter indicates that the posterior covariance is nonzero and has an inverse relation with the learning rate. In the VKF model, unlike the HGF, this covariance plays a key role as it influences volatility update because To demonstrate these issues, we generated a sequence of observations using the generative model of the HGF (see Methods). We then used the generated sequence as the input to the HGF inference algorithm, using true generative parameters. We also performed inference using VKF on these generated time-series, by fitting its free parameter, volatility learning rate, , to data. Note that we have set up this experiment to favor the HGF, in that any performance differences due to the models assuming different generative dynamics will work against the VKF. We repeated this process 500 times. Note that we have set up this experiment to favor the HGF, in that any performance differences due to the models assuming different generative dynamics will work against the VKF. Across all simulations, in 51 simulations, the HGF encountered numerical problems: its inferred trajectory conflicted with the assumptions of the inference model, resulting in negative posterior variance at the third level. Only for the remaining 449 simulations were we able to calculate the median error of predicted outcomes (defined as the absolute normalized difference between predicted and observed outcome) across all time points. We found that in 325 of simulations (i.e. 72%), the error of the VKF was smaller than the error of the HGF, despite the simulations being conducted under the HGF's own generative assumptions. That is, the VKF's more accurate approximate resulted in overall better performance even given the mismatch between the assumed and true generative models.

Testing VKF using Empirical data
We tested VKF using an empirical dataset from a probabilistic learning task (previously published in [30]).
In the experiment ( Figure 4A), 44 participants were presented with facial cues and were asked to make either a go-or a no-go-response (i.e., press a button, or withhold a press, respectively) for each of four facial cues in order to obtain monetary reward or avoid monetary punishment. The cues were different combinations of the emotional content of the face image and its background color; together, these indicated the chance, at the end of each trial, that a go or no-go response would earn the better of two outcomes (monetary reward vs nothing, or nothing vs loss) would occur. The response-outcome contingencies for the cues were probabilistic and manipulated independently, and reversed after a number of trials, varying between 5 and 15 trials, so that the experiment consisted of a number of blocks with varying trial length ( Figure 4B). Within each block, the probability of a win was fixed. The task was performed in the scanner, but here we only focus on behavioral data. 120 trials were presented for each trial-type (480 trials in total). Participants learned the task effectively: performance, quantified as the number of correct decisions given the true underlying probability, was significantly higher than chance across the group (t(43)=14.68, p<0.001). The focus of the original studies using this dataset was on socio-emotional modulation of learning and choice [30,31]. Here, however, we use this dataset because the task is a volatile environment in which the contingencies change frequently. We considered a model space including both VKF and HGF, and also a Rescorla-Wagner (RW) model that does not take into account uncertainty or vary its learning rate. We used the softmax rule as a choice model to generate probability of choice data according to action values derived for each model as well as value-independent biases in making a go response. Note that the choice model controlled value-independent biases in making or avoiding a go response due to the emotional or reinforcing content of the cues (see Methods).
This model space was then fit to choice data using hierarchical Bayesian inference (HBI) procedure [32], a Bayesian algorithm to perform concurrent model fitting and model comparison with the advantage that fits to individual subjects are constrained according to the group-level statistics (i.e. empirical priors) taking into account the possibility that different subjects might express different models. The HBI also performs random effects model comparison by quantifying model evidence across the group (goodness of fit penalized by the complexity of the model [33]). Across participants, VKF was the superior model in 37 out of 44 participants. Furthermore, the protected exceedance probability in favor of VKF (i.e. the probability that a model is more commonly expressed than any other model in the model space [34,35] taking into account the null possibility that differences in model evidence might be due to chance [35]) was indistinguishable from 1 ( Figure 5). These results provide evidence that volatility affects choice and VKF is more parsimonious than HGF in tracking volatility. Note that since the outcome is binary here, we used the binary version of HGF, as described in [18] (considering constraints described by Mathys et al. [36], see Methods), and the binary version of VKF.

Discussion
In this work, we have introduced a novel model for learning in volatile environments. The proposed model has theoretical advantages over existing models of learning in volatile environments, mainly because it better captures the temporal dependencies and it is based on a novel generative model of uncertainty.
Using empirical choice data in a probabilistic learning task, we showed that this model captures human behavior better than the state-of-the-art HGF.
The Kalman filter is the cornerstone of tracking theories, with widespread applications in many technological and scientific domains including psychology and neuroscience [4,7,23,24]. For example, in movement neuroscience, the Kalman filter has been used as a model of how the brain tracks sensory consequences caused by a motor command. In learning theories, the Kalman filter provides a normative foundation for selective attention among multiple conditioned stimuli predicting a target stimulus, such as food. Nevertheless, the Kalman filter is limited to environments in which the structure of process noise is constant and known. The VKF fills this gap by extending the Kalman filter to inferring the process variance in volatile environments in which that variance is itself dynamically changing. In particular, VKF contains two free parameters, a volatility update rate (i.e. ) indicating the extent of noise in process variance dynamics and the initial value of volatility ( K ). The Kalman filter is indeed a special case of VKF, in which the volatility update rate is zero and the constant process noise is equal to K .
A complete understanding of a learning system requires understanding of how computational theories should be realized at the algorithmic level [17]. Previous works have shown that at this level, the normative perspective might shed light on or even encompass related psychological models, as for example temporal difference models of reinforcement learning encompass the classical Rescorla-Wagner theory. The proposed model builds a normative foundation for the intuitive hybrid models combining critical features of Rescorla-Wagner and Pearce-Hall theories for conditioning [21,22,31]. Specifically, the learning process of VKF depends on two components. The first component is the classical prediction error signal, which is defined as the difference between the observed and expected outcome, similar to Rescorla-Wagner error signal. The second component is learning rate modulated by volatility estimate, which is itself a function of surprise (i.e. the squared prediction error) according to Equation 3. Therefore, although the detailed algebraic computation of the surprise term slightly differs, the structure of the model is consistent with the classical Pearce-Hall model, and like it the rate of learning in VKF depends on surprise. This construction clarifies the relationship between the Pearce-Hall associability, its update, and volatility inference in hierarchical learning models such as the HGF and that of Behrens et al. [6].
The generative model of VKF is based on a novel state-space model of variance dynamics of Gaussian-distributed observations with known mean, in which the inference (for the volatility level considered in isolation) is tractable. This particular generative model leads to exact inference without resorting to any approximation. To build a fully generative model for volatile environments, we then combined this state-space model with the state-space model of Kalman filter. Therefore, the full generative model of VKF contains two temporal chains (Figure 1), one for generating the mean and the other one for generating the variance. Although the inference is tractable within each chain, the combination of both chains makes the exact inference intractable. Therefore, we used structured variational techniques for making approximate inference. This can be seen as a minimal mean-field assumption in which, unlike the HGF, all temporal dependencies are persevered.
The state-of-the-art algorithm for learning in volatile environments is HGF [8,18]. There are two critical differences between the VKF and the HGF. First, although the generative conditional dependencies of our model and HGF are the same, there is a critical difference between the generative process underlying uncertainty in the two models. The HGF assumes a log-normal distribution to maintain positivity constraint on variance, and uses variational approximations to update both its posterior mean and variance on every trial. In contrast, the generative model of process variance in the VKF makes it possible to make exact inference by tracking the expected value of process variance (i.e. volatility). This feature of the VKF not only makes it an extension of Kalman filter, but also makes it algorithmically simpler than HGF. Secondly, the HGF is based on more severe approximations to the exact inference and therefore is less accurate. Specifically, while the VKF is based on a structured variational approximation that keeps all temporal dependencies, the HGF is based on a full mean-field approximation and therefore neglects covariance between consecutive states. Note that the mean-filed assumption across timepoints can be less accurate in online algorithms like HGF that do not iterate in updating posteriors. We compared the performance of VKF and HGF in predicting human choice data in a probabilistic learning task, similar to those tasks that have been modeled with HGF in the recent past [8]. Bayesian model comparison showed that the VKF is more parsimonious than the HGF and predicts choice data better in vast majority of participants.
Bayesian models have recently been used for online inference of sudden changes in the environment [37,38]. Although those situations can be modeled with generative processes with discrete change-point detection, Behrens et al. [6] showed that models with volatility estimate might be as good as or even better than models with specific discrete change-points. Our simulations also showed that VKF can be successfully applied to those situations. In such situations, the volatility signal plays the role of a continuous change-point estimator, which substantially increases after a major change. This is because those sudden changes in the environments cause a large "unexpected uncertainty" signal [7,11], which substantially increases the volatility.
In this article, we introduced a unifying model for learning under uncertainty. The VKF has theoretical and practical advantages over existing models and explains human choice data better than alternatives.

Ethics statement
Data from human subjects are from a study [30] that is approved by the local ethical committee ("Comissie Mensgebonden Onderzoek" Arnhem-Nijmegen, Netherlands).

Implementation of models for analysis of human choice data
We have fitted a 3-level (binary) HGF [18] to choice data, which contains three free parameters: the variance of the third level, > 0, > 0 determining the extent that the third level affects the second level uncertainty and , which indicates the tonic level variance at the second level. Formally, if V " is the third level random variable on trial , the (generative) variance of the second level is given by exp( V " + ). As suggested by Mathys et al. [36], we assumed an upper bound at 1 for both ν and . We also considered a constant parameter for at -4, as Iglesias et al. [8]. However, since the original HGF with a free outperformed this model (using maximum-a-posteriori estimation and random effects model comparison [34,35]), we included the original HGF in the model comparison with binary VKF. Similar to the HGF toolbox, we assumed that the initial mean at the second and third levels are 0 and 1, respectively, and the initial variance of the second and third levels are 0.1 and 1, respectively.
For implementing the binary VKF, we assumed an upper bound of 10 for the initial volatility parameter, K . The initial value of " and " were assumed to be 0 and equal to , respectively.
Each of the learning models was then combined with a choice model to generate probabilistic predictions of choice data. Expected values were used to calculate the probability of actions, ' (go response) and G (no-go response), according to a sigmoid (softmax) function: " G = 1 − " ( ' ) where " was equal to " for the VKF, and it was equal to 2 ' " − 1 for the HGF (as implemented in the HGF toolbox). Moreover, is the decision noise parameter encoding the extent to which learned contingencies affect choice (constrained to be positive) and ( " ) is the bias towards ' due to the stimulus presented independent from learned values. The bias is defined based on three free parameters, representing bias due to the emotional content (happy or angry), b , bias due to the anticipated outcome valence (reward or punishment) cued by the stimulus, c , and bias due to the interaction of emotional content and outcome, d . No constraint was assumed for the three bias parameters. For example, a positive value of b represents tendencies towards a go response for happy stimuli and for avoiding a go response for angry stimuli (regardless of the expected values). Similarly, a positive value of c represents a tendency towards a go-response for rewarding stimuli regardless of the expected value of the go response. Critically, we also considered the possibility of an interaction effect in bias encoded by d .
Therefore, the bias, ( " ), for the happy and rewarding stimulus is b + c + d , the bias for the angry and punishing stimulus is − b − c + d , the bias for the happy and punishing stimulus is b − c − d and the bias for the angry and rewarding stimulus is − b + c − d .

Simulation analysis using HGF generative model
We generated data according to the HGF generative model (with Gaussian observations) according to these parameters: = 0.5, = 1 and = −3. The observational noise was assumed to be 1. The initial mean at the second and third levels were 0 and 1, respectively. These parameters were also used for inference using the HGF algorithm. For inference using the VKF algorithm, we used the generative initial volatility (i.e. exp V K + = 0.135) and generative observational noise (i.e. 1). We fitted the VKF volatility learning rate parameter to data for each time-series using a maximum-a-posteriori procedure, in which the prior mean and variance were assumed to be 0 and 15.23, respectively.

Model fitting and comparison
We used a hierarchical and Bayesian inference method, HBI [32], to fit models to choice data. The HBI performs concurrent model fitting and model comparison with the advantage that fits to individual subjects are constrained according to hierarchical priors based on the group-level statistics (i.e. empirical priors). Furthermore, the HBI takes a random effects approach to both parameter estimation and model comparison and calculates both the group-level statistics and model evidence of a model based on responsibility parameters (i.e. the posterior probability that the model explains each subject's choice data). The HBI quantifies the protected exceedance probability and model frequency of each model, as well as the group-level mean parameters and corresponding hierarchical errors ( Figure 5). This method fits parameters in the infinite real-space and transformed them to obtain actual parameters fed to the models. Appropriate transform functions were used for this purpose: the sigmoid function to transform parameters bounded in the unit range or with an upper bound and the exponential function to transform parameters bounded to positive value. The initial parameters of all models were obtained using maximum-a-posteriori procedure, in which the prior mean and variance for all parameters were assumed to be 0 and 6.25, respectively. This initial variance is chosen to ensure that the parameters could vary in a wide range with no substantial effect of prior. Specifically, for any parameter bounded in the unit range, the log-effect of this prior is less than one chance-level choice (i.e log0.5) for any value of that parameter between 0.05 and 0.95. The HBI algorithm is available online at https://github.com/payampiray/cbm.
Supporting information S1 Appendix. Formal treatment of the VKF model.
In this appendix, we give a formal treatment of the VKF. We first review state space models and principles of making inference in those models. Next, we explain two specific state space models with different types of noise. The first one Next, the projected estimate is adjusted by actual measurement, giving rise to the posterior given all observations: where c t = p(u t |u 1 , ..., u t−1 ) is constant with respect to s t . Note that since c t is independent of s t , it is only needed to be calculated if one is interested in the likelihood of data, which is easy to see to be given by Furthermore, the posterior over two consecutive states s t−1 and s t is given by where C is a normalization constant independent of s t−1 and s t .

A.2 Kalman filter
Consider a state space model in which, on trial n, hidden state of the environment, x t , is equal to its state on trial n − 1 plus some Gaussian noise.
where e t is a zero-mean Gaussian random variable with variance v. Therefore, the generative probabilistic model of x t is The initial state is also assumed to be Gaussian Observations are given by a Gaussian distribution: where ω is the observation variance. Using equations (A.2-A.3), we obtain: and k t is the learning rate (also called Kalman gain) and is given by: Furthermore, we can make use of (A.5) to obtain the covariance between consecutive states, x t−1 and x t :

A.3 A tractable filter for variance
Let's now consider another state space model in which hidden state on trial t, z t , depends on its value on previous trial, z t−1 multiplied by some (independent) noise, t . Then, a Gaussian random variable, x t , is drawn with a known mean, µ t , and a precision given by z t , Therefore, conditional dependencies indicate that equations (A.2-A.3) hold for the posterior over z t .
Our goal in this section is to build a rich model of dynamic (inverse) variance with tractable and simple update equations. First, we inspect the functional form of h t (z t ) with respect to z t : Therefore, h t (z t ) takes a form of gamma distribution in which information about new observations is always coming through the rate and the shape is constant.

Moreover, equation (A.3) indicates that ifq(z t ) is a gamma distribution, then
q(z t ) takes the form of gamma distribution. Thus, for the moment, we assume that q(z t ) is a gamma distribution with a constant shape a and a dynamic rate In the previous section, we saw that the transition noise was additive. For variance, it makes sense to assume a multiplicative noise. Thus, we suppose that the value of z t is equal to z t−1 multiplied by some noise Assuming that the noise is bounded between zero and R ≥ 1, we can write this equation as where t = Rε t , in which ε t is constrained in the unit range and has a beta distribution. Assuming that conditional expectation of z t is z t−1 , we obtain A general beta distribution with this property is given by: where η = R −1 and ν is a constant positive parameter. We show below that there is a specific value for ν, which makes the inference tractable.
Using transform theorem for random variables, the conditional distribution of z t is given by Applying equation (A.2) toq(z t ), we obtain If we take a to be equal to ν, the integral is tractable resulting in a gamma distribution forq(z t ). Note that evaluating those terms that are independent of z t is not even necessary becauseq(z t ) is a probability distribution. Therefore, we haveq Combiningq(z t ) with h(z t ), we obtain q(z t ): If we now take the shape of q(z t ) to be the same as the shape of q(z t−1 ), the 5 inference is similarly tractable for all next trials. Consequently, ν is given by Therefore, the posterior distribution is a gamma distribution with a fixed shape and an evolving rate We can also write b t in terms If we define λ = 1 − η, we can write this equation in the form of an errorcorrecting update rule: It is also important to note that E[z t ] underq(z t ) is given by v t−1 according to equation (A. 19). Furthermore, as equation (A.2) indicatesq(z 1 ) = p(z 1 ), we take the probability of initial state to be given by where b 0 > 0 is a free parameter, which we can equivalently write it according to the initial volatility value v 0 . Note that as a consequence of making the inference tractable, the form of multiplicative transition noise is restricted according to equation (A.21). However, this is not a major restriction, as the mean of noise is still free to be chosen. Note that it is also possible to assume other forms for ν, although those assumptions might make the inference slightly more complicated.

A.4 Volatile Kalman filter
In this section, we provide a solution to the general prediction problem in volatile environments. Specifically, consider a problem in which on trial t a latent ran-6 dom variable, x t , is given by its previous value, x t−1 , plus some Gaussian noise with a precision given by another dynamic random variable, z t : The dynamic of z t is given by equations (A.17-A.18). The observation on trial t, o t , is then drawn based on the value of x t according to a normal distribution: Therefore, the graphical model corresponding to this problem consists of two coupled chains of latent variables, x t and z t . Exact inference for this problem is not possible due to coupling between the two chains. However, it is possible to derive a variational approximation by inducing factorization between the two chains, while preserving all dependencies within each chain. Here, we assume p(x 1 , ..., x t , z 1 , ..., z t |o 1 , ..., o t ) q(x 1 , ..., x t )q(z 1 , ..., z t ) (A.28) Using variational calculations, we obtain where all terms independent of x 1 , ..., x t are absorbed into C and Since the conditional dependencies for Next, we use variational techniques to obtain q(z t ): where all terms independent of z 1 , ..., z t are absorbed into C and we have t−1 + w t−1 + m 2 t + w t − 2(m t−1 m t + w t−1,t ) =(m t − m t−1 ) 2 + w t−1 + w t − 2w t−1,t (A. 35) These results give rise to an efficient and simple algorithm for learning in volatile environments. On every trial t, we first compute E[z t ] −1 underq(z t ), which is given by v t−1 . Next, we update m t and w t and v t , according to equations (A.11-A.14) by replacing v with E[z t ] −1 = v t−1 . Finally, making use of equations (A.35), (A.12) and (A.14), we can simplify (A.34) and obtain equation (8). Note that here, unlike typical algorithms obtained based on variational Bayes, we only update statistics once to have an efficient algorithm for learning.

A.5 Binary VKF
In this appendix, we give a formal treatment of the binary VKF, in which observations are drawn using the sigmoid function: where σ(.) is the sigmoid function: σ(x t ) = 1 1 + exp(−x t ) (A.37) Note that for binary observations, (A.3) indicates that q(x t ) is not normally distributed even if we assume thatq(x t ) is normal. Therefore, some approximate strategies are required for binary observations.
A well-known method approximates q(x t ) by matching its moments to those of p(o t |x t )q(x t ). For a Gaussian distribution, this means that first and second order moments should be matched. However, since matching variance might lead to negative values for the variance, we assume that the variance is constant (similar to Kalman filter) and match only the mean of distributions. Formally, we take the "message" from o t to x t to be an unnormalized Gaussian, g t (x t ) = c t N (x t |m t , ω), in which the variance ω is a constant and the mean will be obtained by matching the mean of q(x t ) = N (x t |m t , w t ) with that of p(o t |x t )q(x t ), in whichq(x t ) = p(x t |o 1 , ..., o t−1 ) = N (x t |m t−1 , w t−1 + v t−1 ).
We used a simple approximation for the arising integral, which is inspired by Mackay (1992): It is also possible to approximate m t using a slightly more complicated equation given by MacKay (1992). Using simulation analyses, however, we found that (A.38) results in better approximations. The variance of q(x t ) is then given by which results in equations (6)(7)(8).