Fast and robust Bayesian inference for modular combinations of dynamic learning and decision models

In cognitive neuroscience, there has been growing interest in adopting sequential sampling models (SSM) as the generative choice function for reinforcement learning (RLSSM) to jointly account for decision dynamics within and across trials. However, such approaches have been limited by computational tractability due to lack of closed-form likelihoods for the decision process or expensive trial-by-trial evaluation of complex reinforcement learning (RL) processes. We enable hierarchical Bayesian estimation for a broad class of RLSSM models, using Likelihood Approximation Networks (LANs) in conjunction with differentiable RL likelihoods to leverage fast gradient-based inference methods including Hamiltonian Monte Carlo or Variational Inference (VI). To showcase the scalability and faster convergence with our approach, we consider the Reinforcement Learning - Working Memory (RLWM) task and model with multiple interacting generative learning processes. We show that our method enables accurate recovery of the posterior parameter distributions in arbitrarily complex RLSSM paradigms, and moreover, that in comparison, fitting data with the equivalent choice-only model yields a biased estimator of the true generative process. Moreover, leveraging the SSM with efficient inference allows us to uncover a heretofore undescribed cognitive process within the RLWM task, whereby participants proactively adjust the decision threshold as a function of WM load.