# Uncertainty Estimation Using Riemannian Model Dynamics for Offline Reinforcement Learning

arXiv: https://arxiv.org/abs/2102.11327 · HTML: https://mlsanity.com/paper/2102.11327/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 47 of its 81 references that are in the map.
- First posted on arXiv 22 Feb 2021.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 47)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Is Pessimism Provably Efficient for Offline RL?](https://mlsanity.com/paper/2012.15085/) (2020)
- [Decoupling Representation Learning from Reinforcement Learning](https://mlsanity.com/paper/2009.08319/) (2020)
- [Randomized Prior Functions for Deep Reinforcement Learning](https://mlsanity.com/paper/1806.03335/) (2018)
- [Variational Lossy Autoencoder](https://mlsanity.com/paper/1611.02731/) (2016)
- [Revisiting Fundamentals of Experience Replay](https://mlsanity.com/paper/2007.06700/) (2020)
- [A Hierarchical Latent Variable Encoder-Decoder Model for Generating Dialogues](https://mlsanity.com/paper/1605.06069/) (2016)
- [Latent Space Oddity: on the Curvature of Deep Generative Models](https://mlsanity.com/paper/1710.11379/) (2017)
- [Optimizing the Latent Space of Generative Networks](https://mlsanity.com/paper/1707.05776/) (2017)
- [Safe Policy Improvement with Baseline Bootstrapping](https://mlsanity.com/paper/1712.06924/) (2017)
- [What are the Statistical Limits of Offline RL with Linear Function Approximation?](https://mlsanity.com/paper/2010.11895/) (2020)
- [The Importance of Pessimism in Fixed-Dataset Policy Optimization](https://mlsanity.com/paper/2009.06799/) (2020)
- [Unsupervised Learning of Disentangled and Interpretable Representations from Sequential Data](https://mlsanity.com/paper/1709.07902/) (2017)
- [BAIL: Best-Action Imitation Learning for Batch Deep Reinforcement Learning](https://mlsanity.com/paper/1910.12179/) (2019)

## Built on by (most cited first, 10 of 10)

- [When Demonstrations Meet Generative World Models: A Maximum Likelihood Framework for Offline Inverse Reinforcement Learning](https://mlsanity.com/paper/2302.07457/) (2023)
- [Embedding-Aligned Language Models](https://mlsanity.com/paper/2406.00024/) (2024)
- [Delphic Offline Reinforcement Learning under Nonidentifiable Hidden Confounding](https://mlsanity.com/paper/2306.01157/) (2023)
- [Riemann2: Learning Riemannian Submanifolds from Riemannian Data](https://mlsanity.com/paper/2503.05540/) (2025)
- [Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2505.13144/) (2025)
- [Representation-Driven Reinforcement Learning](https://mlsanity.com/paper/2305.19922/) (2023)
- [Emergent Riemannian geometry over learning discrete computations on continuous manifolds](https://mlsanity.com/paper/2512.00196/) (2025)
- [Geometry-aware similarity metrics for neural representations on Riemannian and statistical manifolds](https://mlsanity.com/paper/2603.28764/) (2026)
- [Short and Straight: Geodesics on Differentiable Manifolds](https://mlsanity.com/paper/2305.15228/) (2023)
- [World-Task Factorization for Robot Learning](https://mlsanity.com/paper/2606.02027/) (2026)

## Abstract (arXiv)

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.
