FSCs as minimal agent models

Let us describe the agent’s decisions as actions – where a could be a movement in a given direction, an interaction with an object, or a measurement of some quantity – and observations – for instance, sensory cues of different types. In general, the action taken at time-step t by the agent may depend on the full history of past interactions with the environment, , through some probability distribution .

Rather than keeping a record of past experiences, a realistic agent encodes the history in its internal state . The agent’s behavior can then be described in terms of the probability of taking an action given its current internal state , and of the recurrent computation process , which updates the internal state by processing the information that has just been acquired. Both the policy and the internal state update g are conditional probability distributions, so that in general the process is stochastic. The description of the decision process is completed by a choice of the initial distribution of internal states . A model of the agent is thus defined as the tuple , as we summarize in Fig 1A.

An FSC is an agent characterized by a discrete set of internal states . For the tasks considered in this work, and are also discrete sets of cardinality and . We stress here that an FSC implicitly encodes both past observations and actions in through the recurrence of the internal computation g. In this sense, it generalizes the notion of belief about the state of the external environment that characterizes Bayesian agents [39]. Further, there is an exact mapping between FSCs and linearized gated recurrent units, a gating mechanism used in recurrent neural networks (RNNs) [40,41] (see Section D in S1 Appendix). This correspondence provides a bridge between the abstract representation of an FSC and the more biologically interpretable variables of an RNN, and generalizes input-output Hidden Markov Models [20–23] to allow internal states to retain information about past actions and use it to guide future decisions. Finally, we note that FSCs can always be made expressive enough to reproduce the behavior of an arbitrary agent over a time span T, albeit at the punishing price of an exponentially large number of internal states . In this work, however, we focus on FSCs with relatively small M, following the idea that internal states should represent how the agent compresses available information, rather than performing a complex act of mimicry.

Inference of FSCs from behavioral trajectories

We assume that we have behavioral trajectories, i.e., histories of interactions between agent and environment, each described by a sequence of actions and observations (see Fig 1B). Our goal is to infer a finite state controller that, at a given observation sequence, will reproduce the actions as closely as possible (Fig 1C). To this end, we seek FSC parameters that maximize the log-likelihood of the FSC actions given the observations, or equivalently solve the optimization problem

(1)

where

(2)

Above, the policy , the internal computation g, and the initial distribution of the internal states , which fully characterize the agent, are parametrized by softmax functions (see Materials and Methods). The optimization problem in Eq. (1) does not require a model of how observations are generated by the environment. This is a consequence of observations appearing only as conditional dependencies in Eq. (2). Thus, the environment may be arbitrarily complex, and observations may depend on the current action as well as on the entire history of past interactions.

The maximization of the likelihood in (1) is, in general, a non-convex optimization problem. Thus, in order to infer the parameters that define the FSC, we developed a global optimization method, MAPSO, whose core idea is that multiple FSCs are initialized with random parameters and then evolved according to an interacting swarm dynamics. MAPSO is based on Adaptive Particle Swarm Optimization [42], where particles are FSCs, but the topology of the interactions changes with the metric distance during the optimization. After a fixed number of iterations, the FSC that achieved the largest log-likelihood at any point during training is retained (see Materials and Methods and Section E in S1 Appendix for more details). We remark here that we do not add any regularization terms to the likelihood, such as sparsity constraints, although this is possible in principle. In this work, we infer FSCs with different numbers of internal states M and analyze them separately, and find that, in all cases considered here, there exists a value of M above which FSCs start to accurately reproduce relevant behavioral observables. We note that the number of parameters of an FSC is dominated by the internal computation, growing as M²AY (see Methods). Therefore, in finite datasets, this effectively limits the number of states that can be reliably inferred, especially in scenarios where A and Y are large.

Finally, we note that optimization can also be performed by a local search algorithm, e.g., gradient methods or an adapted Baum-Welch expectation maximization algorithm, augmented with random restarts. While we make no claim of computational supremacy, in our experience, MAPSO turns out to be more robust and efficient in locating global optima, at least for the cases discussed below. In Section F of the S1 Appendix, we show an explicit comparison with expectation maximization and stochastic gradient descent, showing that in some cases MAPSO outperforms them both.

Exact inference of an arithmetic agent

We first test our inference procedure on a purely synthetic agent performing an arithmetic task. The agent can take two actions, “A” or “B”. If the agent takes A, the environment returns a random digit between 0 and 4 with equal probability. Conversely, if it takes the B action, the observation is a digit between 6 and 9. The agent has a counter that is initially set to zero. At the first step, the agent takes a random action with equal probability, observes a digit, and adds it to the counter. At all subsequent times, the agent performs a parity check on the counter and takes action A if it is even and B if it is odd. Then, it receives a new digit and adds it to the counter (Fig 2A). We want to know if our inference algorithm can retrieve the decision-making process just from a dataset of sequences of actions and observations without any information about the inner workings of the environment and of the agent.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 2. Exact inference of synthetic agents.

(a) An agent can take two generic actions, “A” and “B”, and observe a digit. The agent adds the received digits to an internal counter that determines the following action according to a parity check (action A if it is even, B if it is odd). The first action is random. (b) Negative log-likelihood of an FSC with M = 3 internal states during training with N = 50 trajectories of T = 100 steps each. The inferred FSC perfectly reconstructs the agent’s behavior, reaching the expected negative log-likelihood (dashed gray line). (c) Inferred FSC. The state is the initial one, where the action is taken at random (π(A|) ≈ 0.48, reflecting the empirical distribution). The state is dedicated to the action A, whereas the one always takes B. The transitions between the states are always deterministic and solely depend on the parity of the observation, as expected. (d) An agent performing a spatial navigation task in a T-maze 7 tiles in length. The agent enters the maze with a → action and receives an “up” (U) or “down” (D) observation with equal probability. It then either keeps moving to the right (60% probability) or steps up (↑) or down (↓) with a slight bias that depends on the first observation (30% and 10% probability, respectively). Once it reaches the T-junction, the agent chooses ↑ or ↓ solely based on the first observation, ignoring all the ones it received in the corridor. (e) Negative log-likelihood of FSCs with different M during training with N = 1000 trajectories. With M = 5, the inferred FSC reaches the expected negative log-likelihood of the exact model (dashed gray line). (f) The state represents the initial condition, which transitions to one of two “corridor states”, either (first up observation) or (down observation). Once it reaches the end of the corridor, the transition → allows the FSC to correctly take ↑ , while → results in ↓. (g) A chemotactic agent in a one-dimensional concentration gradient. At each time, the agent may detect the presence of a chemical (observation •) or not detect it (), with a probability that depends on its concentration, which in turn depends on the position in space. If in the last two steps it received the same observation, the action is chosen at random between tumble (T), which reorients the agent without any displacement, and run (R), which moves the agent in its current direction. If the last observation was • (a detection) and the previous one was (no detection), then the agent runs. In the reversed case, it tumbles. (h) Negative log-likelihood of FSCs with different M during training with N = 250 trajectories of T = 100 steps each. With M = 4, the FSC reaches the expected likelihood. (i) The internal states of the FSC encode for the four combinations of past observations:  ∼ (y_t−2, y_t−1) = (×,×);  ∼ (×, •);  ∼ (•,×);  ∼ (•, •) and works as a shift-register memory.

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

It would be natural to identify the internal states of this agent with the counter’s value, but their number would be very large if sequences are long. Are there more compact descriptions of the agent’s decision-making process? In fact, our inference algorithm discovers an exact representation of the agent’s behavior with just M = 3 internal states, valid for any sequence length.

In Fig 2B, we plot the negative log-likelihood during training with MAPSO from N = 50 behavioral trajectories of the agent of length T = 100. We find that the inferred FSC reaches the negative log-likelihood expected from the agent, that is the conditional entropy of the empirical distribution of actions conditioned to observations, which in this case reduces to the entropy of the initial action distribution (see Section A in S1 Appendix). In Fig 2C, we draw the inferred FSC topology. The first state,  is the initial one where the action is taken at random. The other two internal states –  and  – are dedicated to taking the A and B action, respectively, and switch between them if an odd digit is observed. It is easy to check that this FSC exactly implements the agent’s decision-making rules without the need for a counter.

Inference of a spatial navigation task

We now consider the case of a synthetic agent performing a spatial navigation task (see, e.g., [43]) that requires the exact memorization of some observations, selectively ignoring other ones, and making associations between long-past and just-received observations (Fig 2D).

The agent is stepping on a tiled T-maze and can take three actions: if it takes action → , it makes a step to the right, whereas actions leave it at the same location (it bounces on the corridor walls). At the beginning of the task, the agent enters the maze from the left, always taking the → action, and receives an observation that can be either “up” or ”down” with equal probability. At subsequent steps, the agent either chooses action → most of the time or takes ↑ or ↓ with a slight bias towards the action that matches the observation received at the entrance step. As it travels along the corridor, the agent keeps receiving observations “up” or ”down” with equal probability until it reaches the T-junction and receives the observation “end”. At this time, it deterministically takes ↑ if the initial observation was “up” and ↓ if it was ”down”. All other observations made along the corridor do not matter – they are just noise.

In Fig 2E, we show the negative log-likelihood of the inferred FSCs with a different number M of internal states. With M = 5, the FSC perfectly reconstructs the agent’s behavior (Fig 2f), whereas with fewer internal states the controller fails to memorize and/or make the required associations. The node () represents the “blank slate” internal state of the agent as it enters the maze. Then, the FSC branches to one of two “corridor states” depending on the initial observation, and there it remains until the junction is reached.  remembers the first “up” observation, so that the corresponding policy has a slight bias in the up direction. After the “end” observation, the FSC switches from to , where the up action is taken deterministically. The other states,  and , represent instead the first “down” observation, allowing the FSC to perfectly match the agent’s behavior.

Inference of a chemotactic agent

Finally, we consider a synthetic chemotactic agent that moves in a one-dimensional concentration gradient. At each time, the agent may detect the presence of a chemical () or not () with a probability that depends on its position, and can either decide to run – keep on moving in its current direction, – or tumble – randomly reorient without moving, . The agent chooses the next action based on the last two observations according to the following policy. Assigning value y = 1 to the detection (•) and y = 0 to no detection (), the agent computes the difference between the last two observations . If the difference is positive, then the agent chooses to run; if it is negative it tumbles; and finally, if it is zero, it chooses at random with equal probability (see Fig 2G). At the beginning of the task, we assume that the agent has not detected a chemical in the last two time-steps. We remark that for this agent the histories of actions and observations are fully entangled: the probability of receiving an observation depends on the position, which in turn depends on the entire past history of actions, each of which depends on previous observations, and so on and so forth.

In Fig 2H, we show that the negative log-likelihood of the inferred FSCs for M = 4 reaches the value expected for the chosen agent. The corresponding FSC is shown in Fig 2G. Inspecting the structure of the transitions between internal states, one realizes that the nodes exactly encode the ordered pair of last observations and the FSC works as a shift-register memory: new observations enter to the right of the register while older observations are pushed to the left and eventually out of it.

We note in passing that if the internal computation matrices are low-rank, then the likelihood is invariant under some linear mixing of the internal states (which is not just a trivial permutation), leading to multiple FSCs with exactly the same behavior. This lack of identifiability indeed takes place for the chemotactic agent and is discussed in detail in Section D of the S1 Appendix. We stress that this is not a limitation of FSCs, but a general property of any behavioral model, as it signals a genuine symmetry of the task. Any equivalent FSC provides an equally valid description of the agent’s strategy.

Evidence accumulation in rats

We now consider an experimental paradigm where rats are trained to perform evidence accumulation during an auditory discrimination task (data from Ref. [44]). During each trial, rats initially insert their nose into a central point for a fixed amount of time (, signalled by a light-emitting diode). During this time, two trains of randomly timed auditory clicks are played simultaneously, both to the left and to the right of the rat, with different click rates and total durations. At the end of the fixation period, the rat decides to poke either to the left or to the right, and receives a water reward if it turns towards the side where more clicks were played (Fig 3A). The sum of the left and right click rates is fixed at , with subsequent trials varying in difficulty depending on the ratio of the click rates.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 3. Evidence accumulation in rats during perceptual decision making.

(a) Sketch of the experiment. During each trial, rats fixate their nose into a central port for . During this time, corresponding to the action “listen”, trains of randomly timed auditory clicks are played both to the left and to the right, with different frequencies and durations (between and ) which represent three observations: left, right, or silence (“” in the figure). After nose fixation ends, the rat receives an “end” observation, and pokes in one of the side ports (actions ← or →). If it turns to the side where more clicks have been played, it receives a water reward. Data of approximately 4500 trials (half for training, half for validation) from [44]. (b) Negative log-likelihood over the validation set of the inferred FSCs with different numbers of internal states. For , a two-layer structure, which reproduces the structure of the experiments, was imposed a priori to speed up convergence. Inset: negative log-likelihood over the training set during inference. (c) At , FSCs (grayscale dots) accurately reproduce rats’ performance (dashed sand line, averaged across trials and individuals) in choosing right as a function of the difference between the number of clicks. A further improvement can be seen for M = 5 onward. (d) Inferred FSC with M = 2 internal states. The FSC learns to take a ← or → action only after the “end” observation, but cannot perform any computation due to the lack of internal states, and thus must choose at random. (e) Inferred FSC with M = 4 internal states and a fixed two-layer structure. The two internal states on the top (,  ) are responsible for taking the “listen” action, acting as a computation layer where the agent accumulates evidence. The two at the bottom (,  ) decide whether to go left or right. The state remains unchanged if a click is heard on the left or no click is heard, whereas it switches to with a probability . If “end” is heard, the trial ends and switches with probability one to , where the agent chooses to go to the left. The situation is reversed for the states and , the latter of which goes to the right. (f) It can be shown that this FSC is computing an exponential smoothing of the difference between right and left observations with a smoothing factor (Eq. (3). Indeed, the fraction of times that the rats choose to go right is well-predicted by its value at the end of the trial. Since is captured by the internal state , the FSC prediction simply corresponds to the bisector of the plane.

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

From the behavioral trajectories, we can discretize time (we set ) and define four different observations and three actions. At each timestep, as long as the light is on, the agent may only take the action “listen” and receive a “left” () or “right” (y = 1) observation if it hears a click on the corresponding side, or a “silence” observation (y = 0) if it hears neither. When the light-emitting diode turns off, the agent receives an “end” observation and can take one of the two actions “go left” or “go right”.

Among the data in [44], we extracted the behavioral data of three rats, each across two different days. To average out noise in the rats’ behavior and accurately sample different task difficulties, we retained approximately 4500 trials, with 50% used for inferring the FSC parameters and 50% to validate their performance. In Fig 3B, we show the negative log-likelihood of the best-performing FSCs for different numbers of internal states M. Remarkably, we find a sharp decrease at M = 4, after which the performance quickly plateaus. Indeed, the FSCs with closely reproduce the observed rats’ performance (Fig 3C) with slight deviations visible only for M = 4. This signals that an FSC with a relatively small number of internal states is expressive enough to reproduce the behavioral trajectories.

A small FSC recapitulates the essential features of decision-making

With M = 2 and M = 3, the inferred FSCs are trivial, as their limited internal structure can only reproduce the fixation period during which rats listen until the “end” cue (see Fig 3d and Section D in S1 Appendix). For example, with M = 2, no computation can be done in the first internal state because “left” and “right” observations always circle back to this state. When “end” is observed, the FSC transitions to the other state and can only choose “go right” or “go left” at random, regardless of the past history. This model can only reproduce the average final decision across all data, and therefore cannot accumulate past evidence. Inferred FSCs with M = 3 nodes similarly fall short.

The structure of the interactions for this task hints at a layered organization of the internal states: a computation layer where the agent can only take the “listen” action, and a decision layer responsible for taking the “go” actions and accessible only upon receiving the ”end” signal. For computation to be possible, at least two nodes in the upper layer are required, whereas two are both necessary and sufficient for the final decision. We therefore focused on FSC architectures with computation layer nodes and 2 decision layer nodes, imposing this structure a priori for inference efficiency. A posteriori, we then check that these structures are indeed those that minimize the negative log-likelihood. In Section C of the S1 Appendix, we also show that the same FSCs can be recovered using a curriculum learning strategy, which allows the inference procedure to discover the two-layer structure directly from data.

Let us first focus on the simplest case, M = 4, shown in Fig 3E. We reiterate that, with the layered structure induced by the random interruption, only the two internal states of the top layer are used for internal computation. Despite the small number of states, this FSC is able to achieve a high accuracy in reconstructing behavior, with the state encoding for left clicks, and for right ones. The simplicity of the model allows to compute analytically the predicted probability of choosing an action at the end of the trial, , which obeys the recursion relation

(3)

with . Note that is an exponential smoothing of the difference between right and left observations, with smoothing factor . Here, is the probability of switching between the two upper memories (see Section D in S1 Appendix). In Fig 3F, we show that indeed the fraction of trials in which the rats choose to go right is well predicted by the value of at the end of the trial. Thus, the simplest nontrivial FSC with M = 4 interprets the decision-making process as the following algorithmic procedure: compute the exponentially recency-weighted difference of right and left clicks and, at the end of the trial, bias the choice proportionally to this quantity.

Additional internal states allow longer memory of past events

We now seek to understand how the FSC with M = 5 improves upon the previous one, allowing it to reproduce the rats’ behavior with greater accuracy and to saturate the likelihood (see Figs 3B and 3C).

In Fig 4A, we draw the inferred FSC with five internal states. Similarly to the previous case, we find two nodes in the computation layer ( and ) that encode the prevalence of left and right clicks, respectively. Once the trial ends, they deterministically transition to the corresponding state in the decision layer ( and ) where the agent goes to the left and to the right, respectively. However, the transitions between and are now mediated by an intermediate node . If the “end” observation is received while the FSC is in this state, it transitions to the decision layer with a slight left bias, reflecting an overall observed preference by the rats. Importantly, the state is also the one where the FSC is initialized.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 4. Inferred FSC with a two-layer structure and M = 5 internal states for evidence accumulation.

(a) Similarly to the case of M = 4, the leftmost () and rightmost () internal states in the computation layer are responsible for encoding the prevalence of left and right clicks, respectively, and go in the corresponding direction once the trial ends (respectively and ). Importantly, the switch between them is regulated by the internal state, which transitions to the left or the right when the corresponding click is heard with a probability . The FSC is initialized at this internal state. If the trial ends in the state, it transitions to either or approximately at random (0.56 and 0.44, respectively). (b) Probability of occupying one of the upper internal states (, , ) right before the trial ends, averaged over the observation sequences in the data, as a function of the difference between right and left clicks. The central state is mostly active in an uncertainty region where the difference in the number of clicks is close to zero. (c-d) For a fixed observation sequence y_0:t, the probability of occupying a given internal state may be different between the FSC with M = 5 (colored dots and solid lines) and with M = 4 (dashed lines). In particular, when a set of congruent observations is followed by incongruent ones (e.g., a series of “right” observations first and then a series of “left” ones), the central state acts as a buffer and maintains the occupancy of the right state at higher levels. This reflects in a slower and delayed decay of the probability of taking the action “go right” if the trial were to end at time t (bottom panel).

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

In Fig 4B, we plot the probability of occupying one of the computation states (,  ,  ) right before the trial ends as a function of the difference between the number of clicks heard on the right and on the left. When most observations are from the right, the FSC activates the state, and vice versa for the state. Crucially, the occupation of the middle state peaks when the trial ends in a region of uncertainty, where the difference between clicks is close to zero.

This intermediate node plays an important role. In Fig 4C, we consider a fixed observation sequence y_0:T from the data, characterized first by a series of “right” observations followed by noisy ones where the rat mostly hears “left”. We compute the probability of occupying a given internal state at time , , and compare it with the M = 4 case. We find that, after an initial transient, both FSCs reach a high level of occupancy of , since only right clicks were heard. However, when the incongruent observations start to appear, the state acts as a buffer and delays the transition to , which is instead quickly occupied when M = 4.

We can compute the predicted probability of choosing an action at the end of the trial for this FSC as well (see Section D in S1 Appendix). Remarkably, the calculations reveal that the rats’ decision is captured by a leaky competing accumulator model (LCA) [46] with lateral inhibition. Namely, the preference for choosing right is proportional to the difference of two accumulators – favoring state and , respectively – at the final time. These obey the equations

(4)

and

(5)

with the probability of exiting the central state , and the probability of exiting one of the extremal states ( and ). Thus, the passive accumulator – the one that contrasts with the current evidence – inhibits the active one. Due to this lateral inhibition, the probability of taking the right action if the trial were stopped at time t decays more slowly, allowing the FSC with M = 5 to reproduce the rats’ behavior more accurately and helping identify the side on which more clicks were played, as shown in Fig 4D. Notice that the fact that M = 5 provides a better model for the rats’ decision-making process is also compatible with previous results showing that individual rats exhibit diverse temporal weighting, represented by the and parameters in the inferred FSC.

Mice decision-making in a reversal task

We now consider a different experiment, where mice are trained to perform a reversal task and make decisions in a changing and uncertain environment (data from [45]). At each step, a mouse can choose from two side ports. One of the two ports has a high probability, e.g., 0.8, of delivering a water reward (the high port), whereas the other has a probability 0.2 of rewarding the mice (the low port). After the completion of each trial, the high and the low ports have a probability 0.02 of switching. This results in blocks of consecutive trials, of variable length, where the environment is predictable, even if uncertain. When the environment switches, however, the mouse has to swiftly adapt its strategy to the change (Fig 5A).

Download:

• PNG
  larger image
• TIFF
  original image

Fig 5. Probabilistic decision-making strategies in mice.

(a) Sketch of the experiment. Mice are trained to select one of two side ports, representing two possible actions. Each port has a different probability of giving a water reward, either 80% (the high port) or 20% (the low port), corresponding to the two possible observations. At each time, the high port may randomly switch with a 2% probability, so that mice have to adjust their choices accordingly. Data of 159 behavioral trajectories (126 for training, and 33 for validation) from [45]. (b) Negative log-likelihood over the validation set of the inferred FSCs with different numbers of internal states. Pentagons represent the FSCs with minimum loss across restarts of the inference procedure, and selected restarts are represented by dots. Inset: evolution of the negative log-likelihood over the training set during inference. (c-d) The FSCs reproduce both the probability of selecting the higher reward port () and the probability of switching port selection when the high port changes (), with an accuracy that increases with the number of internal states M and saturates from M = 4. (e) Inferred FSC with M = 2, M = 3, and M = 4. From an internal state, green arrows represent the probability of taking the action left, and red arrows the probability of choosing right. Transitions between the states depend on the observation, either when a reward is received (blue) or not (black). For ease of visualization, actions with low probability and the corresponding state transitions have a higher transparency. In all cases, we can identify two extremal states ( and ) primarily encoding the left or right action, respectively. For M > 2, the central states ( and ) buffer the transitions between the extremal ones, similarly to Fig 4. (f) Marginalizing over the actions emphasizes the linear structure of the FSC. For M > 2, the FSCs have a high probability of transitioning to the central states in the absence of a reward, whereas they drift towards the extremal states when a reward is obtained. (g-i) Probability of occupying the internal states as a function of time and for a specific sequence of high ports (right, left, right), averaged over 10⁴ trajectories. With M = 3 and M = 4, the central states quickly activate when the port changes, playing the role of switch detectors. This allows the FSC to mimic the mice’s strategy and switch to the high-reward action.

https://doi.org/10.1371/journal.pcbi.1014585.g005

From the behavioral trajectories, we can immediately identify two possible actions – “choose left” or “choose right” – and two observations – “water” and “no water”. Importantly, and differently from the previous case, the observations received by the agent are now explicitly dependent on the action it has taken. Furthermore, since any observation can follow from an action, we do not impose any a priori structure on the FSC. We consider the 159 trials reported in [45], with 80% of them used for inference and 20% for validation. In Fig 5B, we plot the negative log-likelihood of the inferred FSC for different numbers of internal states. Once more, we find that the likelihood plateaus already at relatively small values of M. However, at variance with the previous section, we find pronounced improvements when moving from M = 2 to M = 4. In Section B of the S1 Appendix, we explicitly compare the inferred FSCs with different types of hidden Markov models. We find that, while some HMMs can infer the immediate reward-to-action mapping, they cannot reproduce the mice’s commitment due to their inability to encode past actions in their internal states.

Simple FSCs successfully capture the switching behavior

To quantitatively assess the ability of the FSC in capturing the decision process of the mice, we compute the probability of selecting the high reward port and the probability of switching port selection when the high port changes , both as a function of the time from a switch [45]. We find that the inferred FSCs with M = 4 or above are able to reproduce the data very accurately. Remarkably, even for M = 2 and M = 3, the FSCs are still able to capture the timescales at which mice switch their behavior on average.

In Fig 5E, we draw the structure of the inferred FSCs with M = 2, 3, 4. In all cases, we find two extremal states ( and ) where the probability of choosing left and right, respectively, is significantly higher. This suggests that these states encode the agent’s confidence that the corresponding side is the one with a higher chance of delivering the reward. For M = 3, a central state emerges, from which the action is taken almost at random. Similarly to the previous example, this state buffers the transition between the extremal ones, delaying the commitment to a particular action. For M = 4, this effect is ascribed to two separate central states (, ) where the policy still allows taking both actions, but is skewed towards the left and right ones, respectively. By marginalizing over the actions, the transitions between the internal states at fixed observation emphasize that all FSCs have a linear structure (Fig 5F). That is, the inferred internal computation g drives the agent towards the extremal states when a reward is obtained, whereas it transitions to the central ones in the absence of a reward. This linear structure is maintained even for larger M, suggesting that it can consistently model the internal decision process of the mice. The diminishing return in terms of increased accuracy from M = 4 onward is further highlighted by the fact that at large M some inferred internal states become disconnected from the rest, signaling that they are not needed to reproduce the observed behavior (see Section C in S1 Appendix).

Emergence of internal switch detectors accurately reproduces the decision-making process

We now focus on the mechanisms that enable FSCs with larger M to better represent mice behavior. To this end, we generate behavioral trajectories from the FSCs for a fixed sequence of high ports (right, left, right) and track the probability of occupying a given internal state (Figs 5G, 5H and 5I. For M = 2, the FSC is able to correctly follow the sequence of environments, but the probability of occupying the state corresponding to the worst action remains high (e.g.,  when the high port is the right one). This results in a lower probability of taking the best action, as also shown in Fig 5C.

For M = 3, accuracy substantially increases thanks to the middle state . In Fig 5H we show that the probability p(m_t = ) peaks right after the high port changes. Hence,  acts as a switch detector, quickly activating when the number of rewards decreases more than expected in the previous state. The FSC with M = 4 is able to further improve its model of behavior thanks to the emergence of two dedicated switch detectors. As we show in Fig 5I, the activates when the high port changes from the right to the left, whereas is responsible for detecting the left-to-right switch. The policies π(a|) and π (a|) are still dominated by the right and left action, respectively, but are stochastic, allowing these states to probe the opposite action. The presence of these two switch detectors allows the accurate reconstruction of the mice behavior.

FSC parametrization

A finite state controller with M internal states is defined by three probability distributions – the policy, the internal computation, and the initial state distribution – which we parametrize using softmax functions. The policy models the probability of taking an action from a specific internal state , with and . Hence,

where are the parameters we seek to infer. The internal computation models the probability of transitioning to a state from a state m, conditioned on the action a taken by the agent and the observation it received, with . As before,

where . Finally, the initial internal state of the FSC is specified by

with . Therefore, the total number of parameters is M(1 + A + MAY).

Comparison with Hidden Markov Models

Finite state controllers may be viewed as a generalized type of state space models (SSMs) [62] with recurrent transitions. One of the crucial features of FSCs is that transitions between internal states explicitly depend on both actions and observations. This recurrence allows them to encode previous actions and observations in their internal states. In Fig 6 we explicitly compare FSCs (Fig 6A) with different types of hidden Markov models (HMMs).

Download:

• PNG
  larger image
• TIFF
  original image

Fig 6. Comparison between finite state controllers and other state space models.

(a) The FSC model of the agent, where transitions between internal states are fully recurrent - i.e., they depend on the previous observation, action, and state. (b) In a hidden Markov model (HMM), observations are ignored, and transitions between internal states are autonomous. The policy is known as the emission probability, and solely depends on the current internal state (as in FSCs). (c) In a generalized linear model HMM (GLM-HMM), the transitions between internal states remain autonomous, but the policy depends explicitly on the observation via a generalized linear model [22]. (d) In a general input-output HMM [20], both internal transitions and the policy depend on the observation, but not on the emissions, as in FSCs.

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

In a standard HMM (Fig 6B) transitions between internal states are autonomous and depend only on the previous state, while the policy – also known as the emission probability – depends on the current internal state, as in FSCs. Thus, observations do not influence transitions. A possible generalization is HMMs where emissions are parametrized by a generalized linear model (GLM-HMMs), so that the policy explicitly depends on observations [22] (Fig 6C). More generally, in input-output hidden Markov models (IO-HMMs), the dependence on observations enters both the emission probability and the internal computation [20] (Fig 6D). In Section B of the S1 Appendix, we show that these types of models – as well as extended versions of GLM-HMMs whose policy depends on the previous actions as well – fail to capture the mice behavior in the reversal task due to the fact that the dynamics of the internal states is disentangled from that of the emissions, which is instead a key dependency in FSCs.

FSC inference

The goal is to identify the FSC which is as “close” as possible to the true unknown agent, i.e., solve the minimization problem

where FSC is a shorthand for the set of parameters . Above, is the expectation over the distribution of trajectories of actions and observations generated by the interaction of the agent and the environment, D is the Kullback-Leibler divergence, and

where the policy, the internal computation, and the initial distribution are parametrized as described in the previous section. Noting that

and replacing the expectation value with the empirical distribution over a set of N behavioral trajectories for , one arrives at the maximum-likelihood problem (1)

Metric-adaptive particle swarm optimization

The main idea of MAPSO is sketched in Fig 7. First, FSCs are initialized with random parameters. We can think of each FSC as a particle in the space. The parameters of each particle are updated according to a swarming dynamics, where the velocity of each particle is controlled by three different terms. First, an inertia term decreases the velocity with the number of epochs. The second term is a so-called cognitive coefficient, where the velocity vector of the p-th particle evolves to point towards the best set of parameters ever seen by the particle (). Finally, the last term is a social coefficient, where the velocity tends to point towards the best set of parameters ever seen by one of the neighbors of the particles (), where the neighbors are specified by a network of interaction that changes as the algorithm progresses. At the beginning, during a “global phase” all particles interact with one another. Then, a “local metric phase” takes place, where at each step of the algorithm, each particle interacts only with a finite number of neighbors defined by the Euclidean distance between them. The number of neighbors decreases until the particles interact only with their nearest neighbors, before increasing back. The algorithm ends with a global phase to guarantee that all particles share the same set of best parameters.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 7. Sketch of the basic structure of metric-adaptive particle swarm optimization (MAPSO).

A given number of FSCs are initialized randomly, representing a swarm of particles in the high-dimensional parameter space of the policy, computation, and initial distribution. During an initial global phase, their parameters are updated according to an all-to-all swarming dynamics, where each FSC tends to move towards the best set of parameters it has seen so far, as well as towards the best set of parameters ever found by the swarm. After each step of the swarm evolution, this global best is mutated along a random direction and is kept if the mutation improves the likelihood, and the swarm dynamics is updated according to pre-defined adaptive rules [42]. The global phase is followed by a local metric phase, where the swarm dynamics only involve a finite number of neighbors defined by the Euclidean distance between the particles. During this phase, the FSCs evolve towards the local best set of parameters, i.e., the best set of parameters ever found by the neighbors. The number of neighbors decreases before increasing back, and the algorithm ends with a second global phase.

https://doi.org/10.1371/journal.pcbi.1014585.g007

All the parameters governing the relative strength of these terms are adaptively set according to the fuzzy rules introduced in [42]. Furthermore, at each step, the swarming dynamics is hybridized with a genetic algorithm, where each is mutated along a random direction and retained if it reduces the corresponding negative log-likelihood. MAPSO runs for a fixed number of epochs, where each epoch corresponds to a velocity update, a particle update, a parameter strength update, a computation of the metric distances between the particles and the resulting topology, and a mutation step. See Section F of the S1 Appendix for more details and an analysis of MAPSO robustness against hyperparameter choices.

MAPSO training schedule

To select the best inferred FSC, we consider a cross-validation scheme. We randomly split the data into a training set (50% of the available trajectories in the rats’ experiments, and 80% of the available trajectories in the mice’s experiments) and a validation set (the remaining trajectories). Due to the random split, we note that this implies that the training and validation sets will have trials that, in general, are interdispersed in time. In scenarios where long-timescale effects such as satiation may be important, more careful splits may be chosen. After running MAPSO for a fixed number of epochs over the training set, the best set of parameters ever found by a particle is returned. These parameters correspond to due to the final global phase where all particles interact with one another. Since the optimization problem is non-convex, we run the training procedure multiple times for a given number of internal states M, each with a random set of initial conditions for MAPSO. Then, we compute the negative log-likelihood over the validation set of each of the inferred FSCs, and retain the one with the smallest.

In this work, the initial parameters were extracted from a multivariate Gaussian distribution centered around zero and with diagonal covariance, but other choices are possible to bias the initial search in specific regions of the parameter space. We also note that MAPSO allows for other schedules, such as sequential restarts, where the particles’ initial position is not independent but is centered around the best parameters previously found with varying variance. See Section F of the S1 Appendix for an analysis of MAPSO’s results across restarts.