A probabilistic successor representation for context-dependent prediction

The different strategies that animals use for predicting reward are often classified as model-based or model-free reinforcement learning (RL) algorithms. Model-based RL involves explicit simulation the future to make decisions while model-free strategies rely on learning associations between stimuli and predicted reward by trial and error. An alternative, intermediate strategy for RL is based on the "successor representation" (SR), an encoding of environmental states in terms of predicted future states. A recent theoretical proposal suggests that the hippocampus encodes the SR in order to facilitate prediction of future reward. However, this proposal does not take into account how learning should adapt under uncertainty and switches of context. Here, we introduce a theory of learning SRs using prediction errors which includes optimally balancing uncertainty in new observations versus existing knowledge. We then generalise that approach to a multi-context setting, allowing the model to learn and maintain multiple task-specific SRs and infer which one to use at any moment based on the accuracy of its predictions. Thus, the context used for predictions can be determined by both the contents of the states themselves and the distribution of transitions between them. This probabilistic SR model captures animal behaviour in tasks which require contextual memory and generalisation, and unifies previous SR theory with hippocampal-dependent contextual decision making.