Uncertainty Estimation Using Riemannian Model Dynamics for Offline Reinforcement Learning

TL;DR

This paper introduces a novel uncertainty estimation method for offline reinforcement learning that leverages Riemannian geometry of generative models to improve model accuracy and out-of-distribution detection.

Contribution

It combines parametric and nonparametric uncertainty estimation via a Riemannian metric on latent space, enhancing offline RL performance.

Findings

Significant improvement on continuous control benchmarks.

Enhanced out-of-distribution sample detection.

Better uncertainty quantification in model-based RL.

Abstract

Model-based offline reinforcement learning approaches generally rely on bounds of model error. Estimating these bounds is usually achieved through uncertainty estimation methods. In this work, we combine parametric and nonparametric methods for uncertainty estimation through a novel latent space based metric. In particular, we build upon recent advances in Riemannian geometry of generative models to construct a pullback metric of an encoder-decoder based forward model. Our proposed metric measures both the quality of out-of-distribution samples as well as the discrepancy of examples in the data. We leverage our method for uncertainty estimation in a pessimistic model-based framework, showing a significant improvement upon contemporary model-based offline approaches on continuous control and autonomous driving benchmarks.