Correlation Priors for Reinforcement Learning

TL;DR

This paper introduces a Bayesian framework using Pólya-Gamma augmentation to model correlation structures in discrete environments for reinforcement learning, improving predictive performance with less data.

Contribution

It provides a novel principled method for capturing correlations in discrete decision-making environments, which was previously lacking.

Findings

Outperforms correlation-agnostic models in predictive tasks

Effective in imitation learning, subgoal extraction, and system identification

Requires significantly less training data for comparable accuracy

Abstract

Many decision-making problems naturally exhibit pronounced structures inherited from the characteristics of the underlying environment. In a Markov decision process model, for example, two distinct states can have inherently related semantics or encode resembling physical state configurations. This often implies locally correlated transition dynamics among the states. In order to complete a certain task in such environments, the operating agent usually needs to execute a series of temporally and spatially correlated actions. Though there exists a variety of approaches to capture these correlations in continuous state-action domains, a principled solution for discrete environments is missing. In this work, we present a Bayesian learning framework based on P\'olya-Gamma augmentation that enables an analogous reasoning in such cases. We demonstrate the framework on a number of common decision-making related problems, such as imitation learning, subgoal extraction, system identification and Bayesian reinforcement learning. By explicitly modeling the underlying correlation structures of these problems, the proposed approach yields superior predictive performance compared to correlation-agnostic models, even when trained on data sets that are an order of magnitude smaller in size.