Q-Distribution guided Q-learning for offline reinforcement learning: Uncertainty penalized Q-value via consistency model

TL;DR

Offline reinforcement learning suffers from distribution shift that biases Q-value estimates toward Out-of-Distribution actions. The paper introduces Q-Distribution Guided Q-Learning (QDQ), which learns a Q-value distribution via a consistency model and applies uncertainty-informed pessimism to OOD actions, while employing an uncertainty-aware objective to avoid excessive conservatism. It constructs a trajectory-level truncated Q dataset and uses a one-step consistency model to sample multiple Q-values per state-action pair, enabling robust uncertainty estimation. Theoretical guarantees show convergence of the truncated Q-distribution, favorable contraction properties of the QDQ Bellman operator, and proximity of the learned Q-value to the optimum; empirically, QDQ attains competitive results on D4RL benchmarks, particularly in wide-distribution regimes, with practical guidance for hyperparameters. Overall, QDQ provides a principled, efficient framework for safe yet effective offline Q-learning by tying uncertainty directly to OOD risk via a distributional Q-value model.

Abstract

``Distribution shift'' is the main obstacle to the success of offline reinforcement learning. A learning policy may take actions beyond the behavior policy's knowledge, referred to as Out-of-Distribution (OOD) actions. The Q-values for these OOD actions can be easily overestimated. As a result, the learning policy is biased by using incorrect Q-value estimates. One common approach to avoid Q-value overestimation is to make a pessimistic adjustment. Our key idea is to penalize the Q-values of OOD actions associated with high uncertainty. In this work, we propose Q-Distribution Guided Q-Learning (QDQ), which applies a pessimistic adjustment to Q-values in OOD regions based on uncertainty estimation. This uncertainty measure relies on the conditional Q-value distribution, learned through a high-fidelity and efficient consistency model. Additionally, to prevent overly conservative estimates, we introduce an uncertainty-aware optimization objective for updating the Q-value function. The proposed QDQ demonstrates solid theoretical guarantees for the accuracy of Q-value distribution learning and uncertainty measurement, as well as the performance of the learning policy. QDQ consistently shows strong performance on the D4RL benchmark and achieves significant improvements across many tasks.

Figures (14)

• Figure 1: (a) The maximum of the estimated Q-value often occurs in OOD actions due to the instability of the offline RL backup process and the "distribution shift" problem , so the Q-value of the learning policy (yellow line) will diverge from the behavior policy's action space (blue line) during the training. (b) The red line represents the optimal Q-value within the action space of the dataset, while the blue line depicts the Q-value function of the behavior policy. The gold line corresponds to the Q-value derived from the in-sample Q training algorithm, showcasing a distribution constrained by the behavior policy. On the other hand, the green line illustrates the Q-value resulting from a more conservative Q training process. Although it adopts lower values in OOD actions, the Q-value within in-distribution areas proves excessively pessimistic, failing to approach the optimal Q-value.
• Figure A.1: This exemplifies how the sliding window mechanism operates to augment Q data. Let's consider a sliding window with a width of 50 and a step size of $k=10$. For a specific trajectory, at step 1, we commence with $(s_1,a_1)$ and compute the truncated Q-value utilizing trajectories within the window. At step 2, the sliding window progresses $k$ steps forward, allowing us to compute the truncated Q-value for $(s_{1+k},a_{1+k})$.
• Figure G.1: The derived Q-value distribution when using difference sliding step and same window width to scan over the trajectory's on halfcheetah-medium dataset.The width of the sliding window is set to 200. The Q-value is scaled to facilitate comparison.
• Figure G.2: The Q distribution of the Antmaze-medium-play dataset with varying sliding window widths (100 to 300 steps) is shown in the figure. Widening the sliding window does not change the shape of the Q distribution, even though a larger window covers more information for this sparse reward task with many short trajectories. Instead, enlarging the sliding window decreases the Q value and compresses the size of the derived Q data.
• Figure G.3: The Q-value distribution based on the truncated Q-value v.s. the sample Q-value distribution via the learned consistency model for Gym-MuJoCo tasks.

Theorems & Definitions (22)

• Theorem 4.1: Informal
• Theorem 4.2: Informal
• Theorem 4.3: Informal
• Theorem 4.4: Informal
• Theorem B.1
• Remark B.1
• Remark B.2
• Remark B.3
• Theorem C.1
• Remark C.1