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.