Cognitive Flexibility (CF) is the ability to switch between tasks, even under conditions when the need to switch is not explicitly cued. While the prefrontal cortex and its interaction with subcortical regions are considered central to CF, a key question remains: what are the underlying computational mechanisms that implement theswitch from one task to another? In particular, does the switch rely on 1)~learning processes in which synaptic changes directly alter action execution choice, or 2)~neural state processes that estimate a belief state from which actions can be chosen? Bartolo and Averbeck (2020) argue for the neural state change hypothesis, proposing a Bayesian belief state estimation model, and ruling out Reinforcement Learning as an approach to modeling CF tasks because of its reliance on synaptic changes to implement the switch. We propose instead a Reinforcement Learning-based Deep Recurrent Q-Learning (DRQL) model that simultaneously learns to update a belief state representation based on prior action outcomes, and an action preference representation based on this belief state. This model is presented with a repeated-trial, force-choice probability switching task (PST) in which actions are rewarded stochastically, and the reward probabilites switch between blocks of trials. Although the model is not explicitly cued to the task type, probability of reward, or time of switch, following training, the model performs the PST in the absence of synaptic changes. We show that the trained model produces behavior consistent with non-human primates performing a similar task, and that it develops a belief state representation that captures key information about thecurrent state of the task. 

Skeleton of repository. Code coming soon.