Exclusively Penalized Q-learning for Offline Reinforcement Learning

We consider a Markov Decision Process (MDP) environment denoted as $\mathcal{M}:=(\mathcal{S},\mathcal{A},P,R,\gamma)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $P$ represents the transition probability, $\gamma$ is the discount factor, and $R$ is the bounded reward function. In offline RL, transition samples $d_{t}=(s_{t},a_{t},r_{t},s_{t+1})$ are generated by a behavior policy $\beta$ and stored in the dataset $D$. We can empirically estimate $\beta$ as $\hat{\beta}(a|s)=\frac{N(s,a)}{N(s)}$, where $N$ represents the number of data points in $D$. We assume that $\mathbb{E}_{s\sim D,a\sim\beta}[f(s,a)]\approx\mathbb{E}_{s\sim D,a\sim\hat{\beta}}[f(s,a)]=\mathbb{E}_{s,a\sim D}[f(s,a)]$ for arbitrary function $f$. Utilizing only the provided dataset without interacting with the environment, our objective is to find a target policy $\pi$ that maximizes the expected discounted return, denoted as $J(\pi):=\mathbb{E}_{s_{0},a_{0},s_{1},\cdots\sim\pi}[G_{0}]$, where $G_{t}=\sum^{\infty}_{l=t}\gamma^{l-t}R(s_{l},a_{l})$ represents the discounted return.

In online RL, the optimal policy that maximizes $J(\pi)$ is found through iterative policy evaluation and policy improvement [2, 3]. For policy evaluation, the action value function is defined as $Q^{\pi}(s_{t},a_{t}):=\mathbb{E}_{s_{t},a_{t},s_{t+1},\cdots\sim\pi}[\sum^{\infty}_{l=t}\gamma^{l-t}R(s_{l},a_{l})|s_{t},~{}a_{t}]$. $Q^{\pi}$ can be estimated by iteratively applying the Bellman operator $\mathcal{B}^{\pi}$ to an arbitrary $Q$-function, where $(\mathcal{B}^{\pi}Q)(s,a):=R(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim P(\cdot|s,a),~{}a^{\prime}\sim\pi(\cdot|s^{\prime})}[Q(s^{\prime},a^{\prime})]$. The $Q$-function is updated to minimize the Bellman error using the dataset $D$, given by $\mathbb{E}_{s,a\sim D}\left[\left(Q(s,a)-\mathcal{B}^{\pi}Q(s,a)\right)^{2}\right]$. In offline RL, samples are generated by the behavior policy $\beta$ only, resulting in estimation errors in the $Q$-function for policy actions not present in the dataset $D$. The policy $\pi$ is updated to maximize the $Q$-function, incorporating the estimation error in the policy improvement step. This process accumulates positive bias in the $Q$-function as iterations progress [17].

To mitigate overestimation in offline RL, conservative Q-learning (CQL) [21] penalizes the $Q$-function for the policy actions $a\sim\pi$ and increases the $Q$-function for the data actions $a\sim\hat{\beta}$ while minimizing the Bellman error, where the $Q$-loss function of CQL is given by

where $\alpha\geq 0$ is a penalizing constant. From the value update in (1), the average $Q$-value of data actions $\mathbb{E}_{a\sim\hat{\beta}}[Q(s,a)]$ becomes larger than the average $Q$-value of target policy actions $\mathbb{E}_{a\sim\pi}[Q(s,a)]$ as $\alpha$ increases. As a result, the policy will tend to choose the data actions more from the policy improvement step, effectively reducing overestimation error in the $Q$-function [21].

In this section, we focus on the penalization behavior of CQL, one of the most representative penalty-based offline RL methods, and present an illustrative example to show that unnecessary estimation bias can occur in the $Q$-function due to the penalization. As explained in Section 2.3, CQL penalizes the $Q$-function for policy actions and increases the $Q$-function for data actions in (1). When examining the $Q$-function for each state-action pair $(s,a)$, the $Q$-value increases if $\pi(a|s)>\hat{\beta}(a|s)$; otherwise, the $Q$-value decreases as the penalizing constant $\alpha$ becomes sufficiently large [21].

From the results in Fig. 1, we observe that CQL suffers from unnecessary estimation bias in the $Q$-function for cases (a) and (b). In both cases, the histograms illustrate that policy actions are fully contained in the dataset $\hat{\beta}$, suggesting that the estimation error in the Bellman update is unlikely to occur even without any penalty. However, CQL introduces a substantial negative bias for actions near $0$ where $\pi(0|s_{0})>\hat{\beta}(0|s_{0})$ and a positive bias for other actions. Furthermore, the bias intensifies as the penalty level $\alpha$ increases. In order to mitigate this bias, reducing the penalty level $\alpha$ to zero may seem intuitive in cases like Fig. 1(a) and Fig. 1(b). However, such an approach would be inadequate in cases like Fig. 1(c). In this case, because policy actions close to 0 are rare in the dataset, penalization is necessary to address overestimation caused by estimation errors in offline learning. Furthermore, this problem may become more severe in actual offline learning situations, as the policy continues to change as learning progresses, compared to situations where a fixed policy is assumed.

To address the issue outlined in Section 3.1, our goal is to selectively give a penalty to the $Q$-function in cases like Fig. 1(c), where policy actions are insufficient in the dataset while minimizing unnecessary bias due to the penalty in scenarios like Fig. 1(a) and Fig. 1(b), where policy actions are sufficient in the dataset. To achieve this goal, we introduce a novel exclusive penalty $\mathcal{P}_{\tau}$ defined by

where $f_{\tau}^{\pi,\hat{\beta}}(s)=\mathbb{E}_{a\sim\pi(\cdot|s)}[x_{\tau}^{\hat{\beta}}]$ is a penalty adaptation factor for a given $\hat{\beta}$ and policy $\pi$. Here, $x_{\tau}^{\hat{\beta}}=\min(1.0,\exp(-(\log\hat{\beta}(a|s)-\tau)))$ represents the amount of adaptive penalty that is reduced as log $\hat{\beta}$ exceeds the threshold $\tau$. Thus, the adaptation factor $f^{\pi,\hat{\beta}}_{\tau}$ indicates the average penalty that policy actions should receive. If the probability of estimated behavior policy $\hat{\beta}$ for policy actions exceeds the threshold $\tau$, i.e., policy actions are sufficiently present in the dataset, then $x_{\tau}^{\hat{\beta}}$ will be smaller than 1 and reduce the amount of penalty as much as the amount by which $\hat{\beta}$ exceeds the threshold $\tau$ to avoid unnecessary bias introduced in Section 3.1. Otherwise, it will be 1 due to $\min(1.0,\cdot)$ to maintain the penalty since policy actions are insufficient in the dataset. The latter penalty term $\frac{\pi(a|s)}{\hat{\beta}(a|s)}-1$, positive if $\pi(a|s)>\hat{\beta}(a|s)$ and otherwise negative, imposes a positive penalty on the $Q$-function when $\pi(a|s)>\hat{\beta}(a|s)$, and otherwise, it increases the $Q$-function since the penalty is negative, as the Q-penalization method considered in CQL [21].

To elaborate further on our proposed penalty, Fig. 2(a) depicts the log-probability of $\hat{\beta}$ and the thresholds $\tau$ used for penalty adaptation, with $N$ representing the number of data points. In Fig. 2(a), if the log-probability $\log\hat{\beta}$ of an action $a\in\mathcal{A}$ exceeds the threshold $\tau$, this indicates that the action $a$ is sufficiently represented in the dataset. Thus, we reduce the penalty for such actions. Furthermore, as shown in Fig. 2(a), when the number of actions increase from $N_{1}$ to $N_{2}$, the threshold for determining "enough data" decreases from $\tau_{1}$ to $\tau_{2}$, even if the data distribution remains unchanged.

Furthermore, to explain the role of the threshold $\tau$ in the proposed penalty $\mathcal{P}_{\tau}$, we consider two thresholds, $\tau_{1}$ and $\tau_{2}$. In Fig. 2(b), which illustrates the proposed penalty adaptation factor $f^{\pi,\hat{\beta}}_{\tau_{1}}$ and $f^{\pi,\hat{\beta}}_{\tau_{2}}$ for thresholds $\tau_{1}$ and $\tau_{2}$, $x_{\tau_{1}}^{\hat{\beta}}$ is larger than $x_{\tau_{2}}^{\hat{\beta}}$ because $\tau_{1}>\tau_{2}$. As a result, in the case of $\tau_{1}$, $\mathcal{P}_{\tau_{1}}$ only reduces the penalty for $\pi_{3}$. In other words, $f_{\tau_{1}}^{\pi_{1},\hat{\beta}}=f_{\tau_{1}}^{\pi_{2},\hat{\beta}}=1,$ and $f_{\tau_{1}}^{\pi_{3},\hat{\beta}}<1$. On the other hand, as the number of data samples increases from $N_{1}$ to $N_{2}$, more actions generated by the behavior policy $\beta$ will be stored in the dataset, so policy actions are more likely to be in the dataset. In this case, the threshold should be lowered from $\tau_{1}$ to $\tau_{2}$. As a result, $\hat{\beta}$ exceeds the threshold $\tau_{2}$ in the support of all policies $\pi_{i}$, and $\mathcal{P}_{\tau_{2}}$ reduces the penalty in the support of all policies $\pi_{i}$, i.e., $f_{\tau_{2}}^{\pi_{3},\hat{\beta}}<f_{\tau_{2}}^{\pi_{1},\hat{\beta}}<f_{\tau_{2}}^{\pi_{2},\hat{\beta}}<1$. Thus, even without knowing the exact number of data samples, the proposed penalty $\mathcal{P}_{\tau}$ allows adjusting the penalty level appropriately according to the given number of data samples based on the threshold $\tau$.

Now, we propose exclusively penalized Q-learning (EPQ), a novel offline RL method that minimizes the Bellman error while imposing the proposed exclusive penalty $\mathcal{P}_{\tau}$ on the $Q$-function as follows:

Then, we can prove that the final $Q$-function of EPQ underestimates the true value function $Q^{\pi}$ in offline RL if $\alpha$ is sufficiently large, as stated in the following theorem. This indicates that the proposed EPQ can successfully reduce overestimation bias in offline RL, while simultaneously alleviating unnecessary bias based on the proposed penalty $\mathcal{P}_{\tau}$.

Proof) Proof of Theorem 3.1 is provided in Appendix A.

In order to demonstrate the $Q$-function convergence behavior of the proposed EPQ in more detail, we revisit the previous Pendulum task in Fig. 1. Fig. 4 shows the histogram of $\hat{\beta}$ and the penalty adaptation factor $f_{\tau}^{\pi,\hat{\beta}}(s)$ for Gaussian policy $\pi=N(\mu,0.2)$, where $\mu$ varies from $-2$ to $2$, with varying $\beta$. In Fig. 4(a), $f_{\tau}^{\pi,\hat{\beta}}(s)$ should be less than 1 for any policy mean $\mu$ since all policy actions are sufficient in the dataset. In 4(b), $f_{\tau}^{\pi,\hat{\beta}}(s)$ is less than 1 only if the $\hat{\beta}$ probability near the policy mean $\mu$ is high, and otherwise, $f_{\tau}^{\pi,\hat{\beta}}(s)$ is 1, which indicates the lack of policy action in the dataset. Thus, the result shows that $f_{\tau}^{\pi,\hat{\beta}}(s)$ reflects our motivation in Section 3.1 well. Moreover, Fig. 4 compares the estimation bias curves of CQL and EPQ with $\alpha=10$ in the scenarios presented in Fig. 1. CQL exhibits unnecessary bias for situations in Fig. 4(a) and Fig. 4(b) where no penalty is needed, as discussed in Section 3.1. Conversely, our proposed method effectively reduces estimation bias in these cases while appropriately maintaining the penalty for the scenario in Fig. 4(c) where penalization is required. This experiment demonstrates the effectiveness of our proposed approach, and the subsequent numerical results in Section 4 will numerically show that our method significantly reduces estimation bias in offline learning, resulting in improved performance.

In Section 3.2, EPQ effectively controls the penalty in the scenarios depicted in Fig. 4. However, in cases where the policy is highly concentrated on one side, as shown in Fig. 4, the estimation bias may not be completely eliminated due to the latter penalty term $\frac{\pi}{\hat{\beta}}-1$ in $\mathcal{P}_{\tau}$, as $\pi$ significantly exceeds $\hat{\beta}$. This situation, detailed in Fig. 5, arises when there is a substantial difference in the $Q$-function values among data actions. As the policy is updated to maximize the $Q$-function, the policy shifts towards the data action with a larger $Q$, resulting in a more significant penalty for CQL. To further alleviate the penalty to reduce unnecessary bias in this situation, instead of applying a penalty based on $\hat{\beta}$, we introduce a penalty based on the prioritized dataset (PD) $\hat{\beta}^{Q}\propto\hat{\beta}\exp(Q)$. As shown in Fig. 5(a), which illustrates the difference between the original data distribution $\hat{\beta}$ and the modified data distribution $\hat{\beta}^{Q}$ after applying PD, $\beta^{Q}$ prioritizes data actions with higher $Q$-values within the support of $\hat{\beta}$. According to Fig. 5(a), when the policy $\pi$ focuses on specific actions, the penalty $\frac{\pi}{\hat{\beta}}-1$ increases significantly, as depicted in Fig. 5(b). In contrast, by applying PD, $\hat{\beta}$ is adjusted to approach $\hat{\beta}^{Q}\propto\beta\exp(Q)$, aligning the data distribution more closely with the policy $\pi$. Consequently, we anticipate that the penalty will be significantly mitigated, as the difference between $\pi$ and $\hat{\beta}^{Q}$ is much smaller than the difference between $\pi$ and $\hat{\beta}$. Following this intuition, we modify our penalty using PD as $\mathcal{P}_{\tau,~{}PD}:=f_{\tau}^{\pi,\hat{\beta}}(s)\cdot\left(\frac{\pi(a|s)}{\hat{\beta}^{Q}(a|s)}-1\right)$. It is important to note that the penalty adaptation factor $f_{\tau}^{\pi,\hat{\beta}}(s)$ remains unchanged since we use all data samples in the dataset for $Q$ updates. Additionally, we consider the prioritized dataset for the Bellman update to focus more on data actions with higher $Q$-function values for better performance as considered in [25]. Then, we can derive the final $Q$-loss function of EPQ with PD as

where $w_{s,a}^{Q}=\frac{\hat{\beta}^{Q}(a|s)}{\hat{\beta}(a|s)}=\frac{\exp(Q(s,a))}{\mathbb{E}_{a^{\prime}\sim\hat{\beta}(\cdot|s)}[\exp(Q(s,a^{\prime}))]}$ is the importance sampling (IS) weight, $C$ is the remaining constant term, and the detailed derivation of (4) is provided in Appendix B.1. The ablation study in Section 4.3 will show that EPQ performs better when prioritized dataset $\hat{\beta}^{Q}$ is considered.

Now, we propose the implementation of EPQ based on the value loss function (4). Basically, our implementation follows the setup of CQL [21]. For policy, we utilize the Gaussian policy with a $\textrm{Tanh}(\cdot)$ layer proposed by Haarnoja et al. [4] and update the policy to maximize the $Q$-function with its entropy. Then, the policy loss function is given by

Based on the $Q$-update in (4) and the policy loss function (5), we summarize the algorithm of EPQ as Algorithm 1. More detailed implementation, including the calculation method of the IS weight $w_{s,a}^{Q}$ and redefined loss functions for the parameterized $Q$ and $\pi$, is provided in Appendix B.2.

We compare our algorithm with various constraint-based offline RL methods, including CQL baselines [21] on which our algorithm is based on. For other baseline methods, we consider behavior cloning (BC) and 10% BC, where the latter only utilizes only the top 10% of demonstrations with high returns, TD3+BC [27] that simply combines BC with TD3 [3], Onestep RL [28] that performs a single policy iteration based on the dataset, implicit $Q$-learning (IQL) [29] that seeks the optimal value function for the dataset through expectile regression, mildly conservative $Q$-learning (MCQ) [30] that reduces overestimation by using pseudo $Q$ values for out-of-distribution actions, and MISA [31] that considers the policy constraint based on mutual information. To assess baseline algorithm performance, we utilize results directly from the original papers for CQL (paper) [21] and MCQ [30], as well as reported results from other baseline algorithms according to Ma et al. [31]. For CQL, reproducing its performance is challenging, so we also include reproduced CQL performance labeled as CQL (reprod.) from Ma et al. [31]. Any missing experimental results have been filled in by re-implementing each baseline algorithm. For our algorithm, we explored various penalty control thresholds $\tau\in\{c\cdot\rho,~{}c\in[0,10]\}$, where $\rho$ represents the log-density of $\textrm{Unif}(\mathcal{A})$. For Mujoco tasks, the EPQ penalizing constant is fixed at $\alpha=20.0$, and for Adroit tasks, we consider either $\alpha=5.0$ or $\alpha=20.0$. To ensure robustness, we run our algorithm with four different seeds for each task. Table 1 displays the average normalized returns and corresponding standard deviations for compared algorithms. The performance of EPQ is based on the best hyperparameter setup, with additional results presented in the ablation study in Section 4.3. Further details on the hyperparameter setup are provided in Appendix C.

The results in Table 1 shows that our algorithm significantly outperforms the other constraint-based offline RL algorithms in all considered tasks. In particular, in challenging tasks such as Adroit tasks and AntMaze tasks, where rewards are sparse or intermittent, EPQ demonstrates remarkable performance improvements compared to recent offline RL methods. This is because EPQ can impose appropriate penalty on each state, even if the policy and behavior policy varies depending on the timestep as demonstrated in Section 3.2. Also, we observe that our proposed algorithm shows a large increase in performance in the ‘Hopper-random’, ‘Hopper-medium’, and ‘Halfcheetah-medium’ environments compared to CQL, so we will further analyze the causes of the performance increase in these tasks in the following section. For adroit tasks, the performance of CQL (reprod.) is too low compared to CQL (paper), so we provide the enhanced version of CQL in Appendix E, but the result in Appendix E shows that EPQ still performs better than the enhanced version of CQL.

In Section 4.1, EPQ outperforms CQL baselines significantly across various D4RL tasks based on our proposed penalty in Section 3. To analyze the impact of our overestimation reduction method on performance enhancement, we compare the estimation bias for EPQ and CQL baselines with various penalizing constants $\alpha\in\{0,1,5,10,20\}$ on ‘Hopper-random’, ‘Hopper-medium’, and ‘Halfcheetah-medium’ tasks. In Fig. 6(a), we depict the squared value of estimation bias, obtained from the difference between the $Q$-value and the empirical average return for sample trajectories generated by the policy, to show both overestimation bias and underestimation bias. In the experiment shown in Fig. 6(a), the estimation bias in CQL with $\alpha=0$ became excessively large, causing the gradients to explode and resulting in forced termination of the training. Fig. 6(b) illustrates the corresponding normalized average returns, emphasizing learning progress after $200\mathrm{k}$ gradient steps.

In Fig. 6(a), we observe an increase in estimation bias for CQL as the penalizing constant $\alpha$ rises, attributed to unnecessary bias highlighted in Fig. 1. Reducing $\alpha$ to nearly 0 in CQL, however, fails to effectively mitigate overestimation error, leading to a divergence of the $Q$-function in tasks such as ‘Hopper-random’ and ‘Hopper-medium’, as shown in Fig. 1. Conversely, EPQ demonstrates superior reduction of estimation bias in the $Q$-function compared to CQL baselines for all tasks in Fig. 6(a), indicating its capability to mitigate both overestimation and underestimation bias based on the proposed penalty. As a result, Fig. 6(b) shows that EPQ significantly outperforms all CQL variants on ‘Hopper-random’, ‘Hopper-medium’, and ‘Halfcheetah-medium’ tasks.

To understand the impact of EPQ’s components and hyperparameters, we conduct ablation studies to evaluate each component and the penalty control threshold $\tau$ on the ‘Hopper-random’, ‘Hopper-medium’, and ‘HalfCheetah-medium’ tasks where our proposed method showed a significant performance improvement compared to the baseline CQL.

Component Evaluation: In Section 3, we introduced two variants of the EPQ algorithm: EPQ (w/o PD), which does not incorporate a prioritized dataset as in equation (3), and EPQ (with PD), which leverages a prioritized dataset based on $\hat{\beta}^{Q}$ as in equation (4). In Fig. 7(a), we compare the performance of EPQ (w/o PD), EPQ (with PD), and the CQL baseline to analyze the impact of each component. EPQ (w/o PD) still outperforms CQL, demonstrating that the proposed penalty $\mathcal{P}_{\tau}$ in Section 3.2 enhances performance by efficiently reducing overestimation without introducing unnecessary estimation bias, as discussed in Section 3.2. Additionally, Fig. 7(a) shows that EPQ (with PD) outperforms EPQ (w/o PD) significantly in the HalfCheetah-medium task, indicating that the proposed prioritized dataset contributes to improved performance, as anticipated in Section 3.3.

Penalty Control Threshold $\tau$: As discussed in Section 3.2, EPQ can dynamically control the penalty amount based on the penalty control threshold $\tau$, as illustrated in Fig. 2, even in the absence of knowledge about the exact number of data samples. Fig. 7(b) demonstrates the performance of EPQ with various penalty control thresholds $\tau\in[0.2\rho,0.5\rho,1.0\rho,2.0\rho,5.0\rho,10.0\rho]$, where $\rho$ represents the log-density of Unif$(\mathcal{A})$. Note that $\rho$ is negative, so $\tau=10.0\rho$ is the lowest threshold while $\tau=0.2\rho$ is the highest. The results indicate that in tasks like Hopper-medium, where a variety of actions are not sufficiently sampled, a higher threshold performs better. Conversely, in tasks like Hopper-random, where a broad range of actions is sampled, a lower threshold is more effective. An exception is the HalfCheetah-medium task, which, despite having fewer action variations, visits a diverse range of states. This may result in lower overestimation errors for OOD actions, benefiting from a lower threshold. Furthermore, the performance on the considered tasks appears to be surprisingly less sensitive to changes in $\tau$. We initially expect that performance might be sensitive to $\tau$ since it reflects the fixed number of data samples, but the results indicate that the performance is not significantly affected by variations in $\tau$. Moreover, EPQ algorithms with different $\tau$ consistently outperform the CQL baseline, highlighting the superiority of the proposed method.

In order to reduce the overestimation in offline learning, several constraint-based offline RL methods have been studied. Fujimoto et al. [17] propose a batch-constrained policy to minimize the extrapolation error, and Kumar et al. [19], Wu et al. [20] limits the distribution based on the distance of the distribution, rather than directly constraining the policy. Fujimoto and Gu [27] restricts the policy actions to batch data based on the online algorithm TD3 [3]. Furthermore, Kumar et al. [21], Yu et al. [32] aims to minimize the probability of out-of-distribution actions using the lower bound of the true value. By predicting a more optimistic cost for tuples within the batch data, Xu et al. [22] provides stable training for offline-based safe RL tasks. On the other hand, Ma et al. [31] utilizes mutual information to constrain the policy.

In offline learning setup, the optimality of the dataset greatly impacts the performance [25]. Simply using $n$-$\%$ BC, or applying weighted experiences, [33, 34] which utilize only a portion of the data based on the evaluation results of the given data, fails to exploit the distribution of the data. Based on Haarnoja et al. [35], Reddy et al. [36], Garg et al. [37] uses the Boltzmann distribution for offline learning, training the policy to follow actions with higher value in the imitation learning domain [38, 39]. Kostrikov et al. [29] and Xiao et al. [40] argue that the optimality of data can be improved by using expectile regression and in-sample SoftMax, respectively. Additionally, methods that learn the value function from the return of the data in a supervised manner have been proposed [41, 28, 42, 43].

In offline RL, imposing constraints on the policy can decrease the performance, thus Kumar et al. [21], Lyu et al. [44] impose penalties on out-of-distribution actions by structuring the learned value function as a lower bound to the actual values. Additionally, Fakoor et al. [45] addresses the issue by imposing a policy constraint based on divergence and suppressing overly optimistic estimations on the value function, thereby preventing excessive expansion of the value function. Moreover, Wu et al. [46] predicts the instability of actions through the variance of the value function, imposing penalties on the out-of-distribution actions, while Lyu et al. [30] replaces the $Q$ values for out-of-distribution actions with pseudo $Q$-values and Agarwal et al. [47], An et al. [48], Bai et al. [49], Lee et al. [50] mitigates the instability of learning the value function by applying ensemble techniques. In addition, Ghosh et al. [51] interprets the changes in MDP from a Bayesian perspective through the value function, thereby conducting adaptive policy learning.