Uncertainty-Aware Rank-One MIMO Q Network Framework for Accelerated Offline Reinforcement Learning
Authors: Thanh Nguyen, Tung Luu, Tri Ton, Sungwoong Kim, Chang D. Yoo
Organizations: School of Electrical Engineering, Korea Advanced Institute of Science and Technology, Daejeon 34141, Republic of Korea · Department of Artificial Intelligence, Korea University, Seoul 02841, Republic of Korea
Offline reinforcement learning (RL) has garnered significant interest due to its safe and easily scalable paradigm. However, training under this paradigm presents its own challenge: the extrapolation error stemming from out-of-distribution (OOD) data. Existing methodologies have endeavored to address this issue through means like penalizing OOD Q-values or imposing similarity constraints on the learned policy and the behavior policy. Nonetheless, these approaches are often beset by limitations such as being overly conservative in utilizing OOD data, imprecise OOD data characterization, and significant computational overhead. To address these challenges, this paper introduces an Uncertainty-Aware Rank-One Multi-Input Multi-Output (MIMO) Q Network framework. The framework aims to enhance Offline Reinforcement Learning by fully leveraging the potential of OOD data while still ensuring efficiency in the learning process. Specifically, the framework quantifies data uncertainty and harnesses it in the training losses, aiming to train a policy that maximizes the lower confidence bound of the corresponding Q-function. Furthermore, a Rank-One MIMO architecture is introduced to model the uncertainty-aware Q-function, \TP{offering the same ability for uncertainty quantification as an ensemble of networks but with a cost nearly equivalent to that of a single network}. Consequently, this framework strikes a harmonious balance between precision, speed, and memory efficiency, culminating in improved overall performance. Extensive experimentation on the D4RL benchmark demonstrates that the framework attains state-of-the-art performance while remaining computationally efficient. By incorporating the concept of uncertainty quantification, our framework offers a promising avenue to alleviate extrapolation errors and enhance the efficiency of offline RL.
Figures & tables
Fig. 1: Illustration of the proposed workflow: Our framework adopts a conventional actor-critic setup, where the online Q network and target Q network are alternated by the proposed Rank-One MIMO Q, while the policy is modeled by a stochastic Gaussian actor. As the framework prohibits direct interaction with the environment, transitions are sampled from a pre-collected offline dataset. These sampled transitions are then utilized to calculate the target, employing the minimum-valued head of the target network and policy entropy. Finally, the calculated target is backed up to the online network using mean square error during training.
Fig. 2: Demonstration of how to generate ensemble weights for two ensemble members from one rank-one layer. It’s important to note that the training weight stored in each rank-one layer consists of only the shared weight and two vectors for each member. The actual weights for each member are calculated on demand following equation 6 .
Task Name
BCQ [ 13 ]
IQL [ 24 ]
BEAR [ 25 ]
UWAV [ 57 ]
CQL [ 26 ]
MOPO [ 60 ]
TD3-BC [ 12 ]
EDAC-10 [ 3 ]
PBRL [ 6 ]
OURS
Random
HalfCheetah
2.3 ±0.0
14.4 ±2.3
2.3 ±0.0
2.3 ±0.0
17.5 ±1.5
35.9 ±2.9
11.0 ±1.1
13.4 ± 1.1
11.0 ±5.8
31.5±2.3
Hopper
10.5 ±0.2
11.1 ±0.1
3.9 ±2.3
2.7 ±0.3
7.9 ±0.4
16.7 ±12.2
8.5 ±0.6
16.9±10.1
26.8 ±9.3
31.3±0.6
Walker2d
4.2 ±1.5
5.8 ±0.3
12.8 ±10.2
2.0 ±0.4
5.1 ±1.3
4.2 ±5.7
1.6 ±1.7
6.7±8.8
8.1 ±4.4
21.3±0.2
Medium
HalfCheetah
40.1 ±0.5
41.7 ±6.2
43.0 ±0.2
42.2 ±0.4
47.0 ±0.5
73.1 ±2.4
48.3 ±0.3
64.1±1.1
57.9 ±1.5
68.6±2.6
Hopper
52.5 ±22.8
64.1±0.3
51.8 ±4.0
50.9 ±4.4
53.0 ±28.5
38.3 ±34.9
59.3 ±4.2
103.6±0.2
75.3 ±31.2
93.2±7.4
Walker2d
47.7±7.8
78.3±0.6
-0.2 ±0.1
75.4 ±3.0
73.3 ±17.7
41.2 ±30.8
83.7 ±2.1
87.6±11.0
89.6 ±0.7
92.2±0.3
TABLE I: The average normalized scores and standard deviations of all algorithms across four seeds in the D4RL benchmark [ 11 ] . The highest-performing scores are highlighted for each dataset. The state-of-the-art scores referenced in this study are from [ 6 , 59 ] .
Fig. 3: The forward time cost and memory of MIMO Q w.r.t the ensemble size. The result is relative to the single model cost.
Runtime (s/epoch)
GPU memory (GB)
CQL
32.4
1.4
PBRL
102.96
1.8
OUR
17.8
0.97
TABLE II: Computational costs on hopper-medium using Tesla V100
Fig. 4: An illustration of the behavior of the Rank-One Q Network on the synthetic regression dataset. yi denotes the prediction of each head.
K
2
5
10
15
20
Avg Return
0.19
92.9
112.8
23
0.4
Avg Q(s,π(s))
3e11
410
373.46
-5e6
-2e12
TABLE III: The effect of K on final performance and Q estimation on walker2d-medium-expert
Entropy
Likelihood
Avg Return
Avg Q(s,π(s))
✓
✓
112.9
373.4
✓
107.3
267.1
✓
111.0
258.9
109.6
453.2
TABLE IV: Component-wise analysis on the walker-2d medium-expert dataset.
Offline reinforcement learning (RL) aims to learn an effective policy from a static dataset, but its performance is fundamentally limited by dataset coverage. Action preference queries leverage expert feedback without additional environment interaction, enabling policy improvement during offline training. However, existing methods still face two key challenges: selecting informative preference queries and effectively exploiting the collected feedback. Current approaches typically rely only on the distance between policy actions and dataset actions for query selection, while enforcing fixed constraints that keep the policy close to queried preferences. Such strategies often lead to unstable policy updates and integrate poorly with value regularization. To address these limitations, we propose Conservative Query and Adaptive Regularization under Uncertainty Estimation, a lightweight framework that jointly improves preference querying and preference exploitation. Specifically, we employ a Morse network to estimate the uncertainty of policy actions with respect to the offline dataset. Based on this uncertainty, we introduce a conservative query strategy that selectively queries actions near the dataset to preserve Bellman-update stability, together with an uncertainty-aware adaptive regularization scheme that dynamically adjusts data-level constraints during policy optimization. We integrate our framework with CQL and evaluate it extensively on the D4RL benchmark. Experimental results demonstrate superior or competitive performance across a wide range of tasks.
Li-Rong Zhou, Qin-Wen Luo, Sheng-Jun Huang
Nanjing University of Aeronautics and Astronautics, Nanjing, China
Offline reinforcement learning requires careful conservatism to mitigate distribution shift, yet most existing methods apply a fixed penalty uniformly across all states regardless of local data coverage. We present UNIQ (Uncertainty-Informed Quantile), an offline RL method that introduces state-adaptive conservatism through conformally calibrated uncertainty estimation. Built on the Implicit Q-Learning (IQL) backbone, UNIQ trains a multi-expectile value ensemble, computes distribution-free uncertainty estimates using split conformal prediction, and maps the resulting signal to a state-dependent expectile that relaxes conservatism in well-covered regions while strengthening it in uncertain regions near the data frontier. On D4RL MuJoCo benchmarks, UNIQ consistently improves over IQL, with the largest gains observed on Walker2d and replay-heavy tasks. At the same time, UNIQ operates at near-IQL memory cost (approximately 250 MB peak VRAM), providing roughly a 10x reduction compared to EDAC. Rather than pursuing overall state-of-the-art performance, we position UNIQ as a practical mechanism contribution that improves the performance-efficiency trade-off in offline reinforcement learning.
Offline-to-Online Reinforcement Learning (O2O-RL) leverages an offline, pre-trained policy to minimize costly online interactions. Although data-efficient, O2O-RL is susceptible to shifts between offline and online distributions. Existing work aims to mitigate the harm of this shift by finetuning the policy on trajectory data sampled from a diffusion model. Inspired by this line of work, we propose DUAL: an efficient Diffusion Uncertainty-Aware framework for offline-to-online reinforcement Learning. DUAL utilizes the prior knowledge of the diffusion model to distill a fast-sampling diffusion actor policy and transition model in the offline phase. DUAL also employs a Laplace approximation and distance transition-state-shift detection, thereby using uncertainty quantification to improve exploration versus exploitation in the online phase. We formally show that our actor loss with the Laplace approximation provides a proxy for a principled estimate of epistemic uncertainty. Empirically, DUAL improves the online expected return over O2O-RL baselines across multiple settings and environments.
Ha Manh Bui, Metod Jazbec, Eric Nalisnick +1
Department of Computer Science, Johns Hopkins University, Baltimore, MD, U.S.A. · AMLab, University of Amsterdam, Amsterdam, Netherlands