A Local Temporal Difference Code for Distributional Reinforcement Learning

Recent theoretical and experimental results suggest that the dopamine system implements distributional temporal difference backups, allowing learning of the entire distributions of the long-run values of states rather than just their expected values. However, the distributional codes explored so far rely on a complex imputation step which crucially relies on spatial non-locality: in order to compute reward prediction errors, units must know not only their own state but also the states of the other units. It is far from clear how these steps could be implemented in realistic neural circuits. Here, we propose a local temporal difference code for distributional reinforcement learning that is representationally powerful and computationally straightforward. The code decomposes value distributions and prediction errors across three completely separated dimensions: reward magnitude (related to distributional quantiles), time horizon (related to eligibility traces) and temporal discounting (related to the Laplace transform of future immediate rewards). Besides lending itself to a local learning rule, the decomposition can be exploited by model-based computations, for instance allowing immediate adjustments to changing horizons or discount factors. Finally, we show that our code can be computed linearly from an ensemble of successor representations with multiple temporal discounts which, according to a recent proposal, might be implemented in the hippocampus.