Deterministic Uncertainty Propagation for Improved Model-Based Offline Reinforcement Learning

TL;DR

This paper introduces MOMBO, a deterministic method for propagating uncertainty in model-based offline reinforcement learning, which improves convergence speed and provides tighter suboptimality guarantees compared to Monte Carlo sampling methods.

Contribution

The paper proposes MOMBO, a novel deterministic approach using moment matching for uncertainty propagation, reducing sampling variance and enhancing convergence in offline RL.

Findings

MOMBO converges faster than Monte Carlo-based methods.

Tighter suboptimality guarantees are achieved with MOMBO.

MOMBO outperforms existing approaches on benchmark tasks.

Abstract

Current approaches to model-based offline reinforcement learning often incorporate uncertainty-based reward penalization to address the distributional shift problem. These approaches, commonly known as pessimistic value iteration, use Monte Carlo sampling to estimate the Bellman target to perform temporal difference-based policy evaluation. We find out that the randomness caused by this sampling step significantly delays convergence. We present a theoretical result demonstrating the strong dependency of suboptimality on the number of Monte Carlo samples taken per Bellman target calculation. Our main contribution is a deterministic approximation to the Bellman target that uses progressive moment matching, a method developed originally for deterministic variational inference. The resulting algorithm, which we call Moment Matching Offline Model-Based Policy Optimization (MOMBO), propagates the uncertainty of the next state through a nonlinear Q-network in a deterministic fashion by approximating the distributions of hidden layer activations by a normal distribution. We show that it is possible to provide tighter guarantees for the suboptimality of MOMBO than the existing Monte Carlo sampling approaches. We also observe MOMBO to converge faster than these approaches in a large set of benchmark tasks.