Deep Recurrent Q-Learning Captures the Behavioral Dynamics Observed in Deterministic and Stochastic Task Switching

Abstract 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 the switch 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 the current state of the task. Significance Statement The proposed DRQL model learns to perform the PST, even when reward probabilities vary across different blocks of trials. Behaviorally, the trained model requires different amounts of time to commit to a task switch, depending on the uncertainty of the reward information, with quicker switches for more certain outcomes. The learned belief state and other computed measures may give insight into the neural mechanisms underlying task switching in the primate brain. In contrast with other approaches, the DRQL model presents a more biologically tractable solution to CF, as it is easily adapted to other forms of the PST task, including altering the number of possible actions and the rules for providing rewards. Competing Interest Statement The authors have declared no competing interest.