Reinforcement Learning with Fast and Forgetful Memory

TL;DR

This paper introduces Fast and Forgetful Memory, a new memory model for reinforcement learning that is faster, more efficient, and outperforms RNNs in various benchmarks by leveraging structural priors inspired by psychology.

Contribution

The paper proposes a novel memory model tailored for RL that replaces RNNs, offering faster training and better rewards without hyperparameter tuning.

Findings

Achieves greater reward than RNNs on multiple benchmarks.

Training speeds are two orders of magnitude faster than RNNs.

Uses structural priors to improve memory efficiency in RL.

Abstract

Nearly all real world tasks are inherently partially observable, necessitating the use of memory in Reinforcement Learning (RL). Most model-free approaches summarize the trajectory into a latent Markov state using memory models borrowed from Supervised Learning (SL), even though RL tends to exhibit different training and efficiency characteristics. Addressing this discrepancy, we introduce Fast and Forgetful Memory, an algorithm-agnostic memory model designed specifically for RL. Our approach constrains the model search space via strong structural priors inspired by computational psychology. It is a drop-in replacement for recurrent neural networks (RNNs) in recurrent RL algorithms, achieving greater reward than RNNs across various recurrent benchmarks and algorithms without changing any hyperparameters. Moreover, Fast and Forgetful Memory exhibits training speeds two orders of magnitude faster than RNNs, attributed to its logarithmic time and linear space complexity. Our implementation is available at https://github.com/proroklab/ffm.