Fast and Memory-Efficient Uncertainty-Aware Framework for Offline Reinforcement Learning with Rank One MIMO Q network

Abstract

Offline reinforcement learning (RL) has recently gained attention for its ability to train policies from existing datasets, eliminating the need for additional environment exploration. However, applying conventional RL algorithms directly to offline RL faces a challenge due to extrapolation errors from out-of-distribution (OOD) data. Existing methods often suffer from issues such as over-conservative value function, imprecise OOD data handling, and high computational costs. To address these limitations, we introduce FMU, a framework that stands for Fast and Memory-efficient Uncertainty quantification. FMU utilizes uncertainty quantification to effectively leverage reliable OOD data while reducing conservatism for in-distribution data. Our framework precisely learns a policy by maximizing the lower confidence bound of the corresponding Q-function and employs a Rank One Multi-Input Multi-Output (MIMO) architecture to effectively make the Q-function uncertainty aware.Extensive experiments on the D4RL benchmark show that FMU achieves state-of-the-art performance while remaining computationally efficient. By incorporating uncertainty quantification,FMU offers a promising solution to mitigate extrapolation errors and enhance offline RL efficiency.