Echo of Bayes: Learned Memory Functions Can Recover Belief States

Juan Liévano-Karim, Peter Koepernik, George Konidaris, Cameron Allen

Track: Extended Abstract Track

Keywords: Partially observable Markov decision process, recurrent neural networks, interpretability, Bayesian inference, state representations

TL;DR: We show that RNN memory learned in model-free RL approximately encodes Bayesian belief state updates, which we recover from hidden states with a memory probe in two POMDP environments.

Abstract: A classical approach to solving partially observable Markov decision processes (POMDPs) is maintaining a belief state---a distribution over hidden states updated iteratively using Bayes' rule. In fact, an optimal policy for a POMDP can always be expressed in terms of the belief state. However, explicitly computing exact posteriors is intractable beyond small-scale problems, and requires knowledge of the hidden state space and transition dynamics, which are usually unavailable to the agent. Recent work has proposed model-free algorithms that help agents learn memory functions that are useful for solving partially observable tasks, but these methods provide no way of interpreting what exactly is being remembered. We hypothesize that these learned memory functions are implementing approximate Bayesian inference. To investigate this, we study two environments where ground-truth state information is available to the experimenter but not to the agent. By probing the hidden states of the trained recurrent networks, we find that in both environments we can reconstruct belief state distributions that closely match the ground-truth.

Submission Number: 142