Uncertainty-Aware Rank-One MIMO Q Network Framework for Accelerated Offline Reinforcement Learning

Abstract: 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.

• The paper introduces a rank-one MIMO critic that approximates an ensemble with nearly single-network compute, reducing memory and forward-pass costs while preserving uncertainty estimates.
• The framework combines minimum-head lower-confidence-bound targets, entropy regularization, and dataset-action likelihoods to limit extrapolation error without costly out-of-distribution action sampling.
• The method achieves an average D4RL score of 83.6 versus 74.37 for PBRL, runs 5.87 times faster than PBRL, and uses 0.97 GB of GPU memory, although performance is highly sensitive to ensemble size.

Overview

This paper addresses a persistent tension in offline reinforcement learning (RL): pessimism is necessary to suppress extrapolation error on out-of-distribution (OOD) actions, but the standard mechanism for calibrated pessimism—a Q-ensemble—is computationally expensive. The authors propose an Uncertainty-Aware Rank-One Multi-Input Multi-Output (MIMO) Q Network framework that retains ensemble-based uncertainty quantification while reducing its cost to nearly that of a single network (2602.19917). The framework builds on PBRL-style bootstrapped uncertainty but replaces the naive ensemble of independent critics with a single shared network augmented by rank-one adapters, and it modifies both policy evaluation and policy improvement losses so that OOD data is exploited selectively rather than uniformly penalized.

Motivation and related approaches

The core failure mode motivating the work is extrapolation error: when the Bellman backup evaluates greedy next-actions $(s', a')$ that rarely appear in the offline dataset, deep value fitting produces severe overestimation, which compounds through bootstrapping. Prior families of solutions each carry drawbacks. Policy-constraint methods (BCQ, BEAR, TD3-BC) tie the learned policy to an estimated behavior policy, which is suboptimal for non-expert datasets and fragile when behavior estimation is difficult. Conservative penalty methods such as CQL penalize all OOD actions uniformly, yielding overly conservative value functions. Uncertainty-aware methods—UWAC (dropout), EDAC (gradient-diversified ensembles), and PBRL (bootstrapped ensembles with explicit OOD sampling)—achieve state-of-the-art results but require separate Q-networks whose forward cost and memory scale linearly with ensemble size $K$, plus additional hyperparameters and, in PBRL's case, costly OOD action sampling.

The paper also situates itself within efficient ensembling literature from supervised learning. Multi-head architectures share a trunk but lack member diversity; full MIMO networks allow distinct paths per member but reportedly struggle beyond two subnetworks. BatchEnsemble-style rank-one factorization offers a middle ground, and this work transfers that idea to offline RL, where empirical evidence for such architectures had been limited.

Rank-One MIMO Q network

The architectural contribution models an ensemble of $K$ critics as one network of stacked rank-one layers. Each layer stores a shared weight matrix $W \in \mathbb{R}^{m \times n}$ common to all members, plus per-member vector pairs $v_k \in \mathbb{R}^m$ and $s_k \in \mathbb{R}^n$. Member-specific weights are materialized on demand as:

$W_k = W \circ (v_k s_k^\top)$

where $\circ$ denotes element-wise multiplication. Because the actual weights are computed via matrix vectorization, all $K$ members can be evaluated in a single batched forward pass:

$Y = \Phi(((X \circ V)\, W) \circ S)$

Memory scales as $K$0 rather than $K$1 for a naive ensemble of $K$2 layers, and the authors report that forward time and memory remain essentially flat in $K$3—equivalent to a single network. Member diversity, critical for meaningful uncertainty estimates, is induced by initializing the individual vectors as random sign vectors, avoiding the extra computational cost of explicit diversity losses used in EDAC.

A limitation worth noting: the diversity argument rests entirely on random sign initialization rather than learned diversification, and the paper does not provide a theoretical guarantee that the resulting members approximate independent bootstrapped critics; the claim of "the same capability" as a naive ensemble is supported empirically, not formally.

Uncertainty-aware training objectives

On top of the architecture, the framework introduces pessimistic losses built around the lower confidence bound (LCB). In policy evaluation, the target uses the minimum over the $K$4 MIMO heads:

$K$5

Invoking Royston's approximation for the expected minimum of Gaussian realizations, the authors show that $K$6 approximates the ensemble mean minus a standard-deviation penalty scaled by a coefficient depending only on $K$7. This yields two practical benefits: LCB estimation requires a single hyperparameter ($K$8), and backpropagation flows only through the minimum-valued head rather than all members, making training cost insensitive to ensemble size.

Two auxiliary components stabilize training without any OOD sampling scheme. First, an entropy bonus on policy-generated next-actions discourages over-reliance on high-value OOD actions, mitigating divergence. Second, policy improvement maximizes a combination of the minimum MIMO Q-value, policy entropy, and the log-likelihood of dataset actions—an in-distribution prior that proves essential on low-coverage expert datasets. A "lazy" policy update schedule (updating the actor every two critic updates) further reduces cost and improves evaluation stability.

Benchmark results

Experiments cover the D4RL Gym suite (HalfCheetah, Hopper, Walker2d across random, medium, medium-replay, medium-expert, and expert datasets), trained for 3 million steps and averaged over 4 seeds. The headline result is an average normalized score of 83.6, versus 74.37 for PBRL—the strongest prior method—a margin of +9.23, and roughly double the scores of BCQ (49.4) and BEAR (38.78). Gains are largest on noisy, low-coverage data: e.g., walker2d-random at 21.3 versus 8.1 for PBRL, and hopper-medium-replay at 102.9.

Method	Avg normalized score	Runtime (s/epoch)	GPU memory (GB)
CQL	67.35	32.4	1.4
PBRL	74.37	102.96	1.8
Proposed	83.6	17.8	0.97

The efficiency claims are substantiated: on hopper-medium with a Tesla V100, the method runs 5.87× faster than PBRL and 1.82× faster than CQL while using the least memory. The speed advantage stems from eliminating OOD sampling, avoiding diversity-loss computation, and restricting backward passes to the minimum head. One caveat: hyperparameters ($K$9 searched over 2–20; $K$0 searched up to $K$1 for expert datasets) were tuned per-dataset via random search, so the reported scores reflect a tuning budget comparable to baselines but not zero-shot transferability of settings.

Ablation findings

Three ablations clarify the framework's behavior. A synthetic regression task confirms that head disagreement grows monotonically outside the training support, validating the uncertainty signal. Sweeping $K$2 on walker2d-medium-expert shows the expected optimism–pessimism trade-off, though with striking sensitivity: average return peaks at 112.8 for $K$3, but collapses to 0.19 at $K$4 (with Q-values exploding to $K$5) and to 0.4 at $K$6 (Q-values collapsing to $K$7). This indicates that while $K$8 is the sole pessimism knob, performance is sharply non-monotonic in it, and the paper does not offer a principled procedure for selecting $K$9 beyond search. Component-wise analysis shows entropy and likelihood terms each contribute modestly (107.3 and 111.0 alone versus 112.9 combined on walker2d-medium-expert), and that removing the in-distribution likelihood term makes learning on expert datasets highly unstable—consistent with the authors' claim that low-coverage data is where this component matters most.

Limitations and open questions

The paper concedes several points implicitly. Evaluation is confined to continuous-control Gym tasks; generalization to higher-dimensional or discrete-action domains is untested. The equivalence to a true bootstrap ensemble is asserted architecturally rather than proven, and the extreme sensitivity to $W \in \mathbb{R}^{m \times n}$0 suggests fragility under suboptimal tuning. The lazy policy update interval and the $W \in \mathbb{R}^{m \times n}$1 ranges are heuristic choices. Open questions include whether rank-one adapters preserve calibration under distribution shift more severe than D4RL's, and whether the min-head LCB approximation remains accurate for small $W \in \mathbb{R}^{m \times n}$2, where the Gaussian assumption underlying Royston's formula is weakest.

Conclusion

This work demonstrates that ensemble-quality epistemic uncertainty for offline RL need not carry ensemble-level compute or memory costs. By combining a rank-one MIMO critic, min-head LCB targets, entropy regularization, and in-distribution likelihood maximization, the framework achieves the best reported D4RL averages among compared methods while running faster and using less memory than CQL and PBRL. Its main open issues—sensitivity to ensemble size and validation beyond standard benchmarks—define the natural next steps for this line of research.