Generalised Linear Models in Deep Bayesian RL with Learnable Basis Functions

TL;DR

GLiBRL introduces a fully tractable Bayesian approach with learnable basis functions for deep Bayesian RL, improving task representation and performance on benchmarks.

Contribution

It presents a novel GLiBRL method that enables exact Bayesian inference and establishes a structural link between task representations and kernel methods.

Findings

Achieves up to 1.8× performance improvement on MuJoCo and MetaWorld benchmarks.

Provides a closed-form relationship between task representation distance and kernel similarity.

Enables seamless integration with on-policy and off-policy RL algorithms.

Abstract

Bayesian Reinforcement Learning (BRL), a subclass of Meta-Reinforcement Learning (Meta-RL), provides a principled framework for generalisation by explicitly incorporating Bayesian task parameters into transition and reward models. However, classical BRL methods assume known forms of transition and reward models. While recent deep BRL methods incorporate model learning to address this, applying neural networks directly to joint data and task parameters necessitates variational inference. This often yields indistinct task representations, compromising the resulting BRL policies. To overcome these limitations, we introduce Generalised Linear Models in Deep Bayesian RL with Learnable Basis Functions (GLiBRL). Our approach features fully tractable Bayesian inference over task parameters and model noise, alongside exact marginal likelihood evaluation for learning transition and reward models. The permutation-invariant nature of exact Bayesian inference in GLiBRL enables seamless integration with both on-policy and off-policy RL algorithms. We further show that GLiBRL admits a closed-form relationship between the $\mathcal{L}_2$ distance of its task representations and empirical kernel-based correspondence between task samples, which is to our knowledge the first such structural result for online deep BRL. GLiBRL is compared against representative and recent Meta-RL methods, and improves state-of-the-art performance on both MuJoCo and MetaWorld benchmarks by up to 1.8$\times$.