Episodic Memory for Learning Subjective-Timescale Models

TL;DR

This paper introduces a novel subjective-timescale model (STM) for model-based learning that dynamically adjusts its temporal horizon, inspired by human time perception, leading to improved exploration and decision-making in AI agents.

Contribution

The paper proposes a new STM framework that learns transition dynamics over subjective timescales, enabling flexible planning and better exploration, which is validated in an AI environment.

Findings

STM can vary its prediction horizon systematically.

STM encourages exploration by predicting future salient events.

STM outperforms baseline models in the Animal-AI environment.

Abstract

In model-based learning, an agent's model is commonly defined over transitions between consecutive states of an environment even though planning often requires reasoning over multi-step timescales, with intermediate states either unnecessary, or worse, accumulating prediction error. In contrast, intelligent behaviour in biological organisms is characterised by the ability to plan over varying temporal scales depending on the context. Inspired by the recent works on human time perception, we devise a novel approach to learning a transition dynamics model, based on the sequences of episodic memories that define the agent's subjective timescale - over which it learns world dynamics and over which future planning is performed. We implement this in the framework of active inference and demonstrate that the resulting subjective-timescale model (STM) can systematically vary the temporal extent of its predictions while preserving the same computational efficiency. Additionally, we show that STM predictions are more likely to introduce future salient events (for example new objects coming into view), incentivising exploration of new areas of the environment. As a result, STM produces more informative action-conditioned roll-outs that assist the agent in making better decisions. We validate significant improvement in our STM agent's performance in the Animal-AI environment against a baseline system, trained using the environment's objective-timescale dynamics.