Density Estimation for Conservative Q-Learning

Abstract: Batch Reinforcement Learning algorithms aim at learning the best policy from a batch of data without interacting with the environment. Within this setting, one difficulty is to correctly assess the value of state-action pairs that are far from the dataset. Indeed, the lack of information may provoke an overestimation of the value function, leading to non-desirable behaviors. A compromise between enhancing the behaviour policy's performance and staying close to it must be found. To alleviate this issue, most existing approaches introduce a regularization term to favor state-action pairs from the dataset. In this paper, we refine this idea by estimating the density of these state-action pairs to distinguish neighbourhoods. The resulting regularization guides the policy toward meaningful unseen regions, improving the learning process. We hence introduce Density Conservative Q-Learning (D-CQL), a batch-RL algorithm with strong theoretical guarantees that carefully penalizes the value function based on the amount of information collected in the state-action space. The performance of our approach is outlined on many classical benchmark in batch-RL.

Reviews and author responses:

 We thank the reviewer for his detailed review. We performed a global answer to the reviewers main concerns. We answer the remaining open questions in the following paragraphs.

**Difference with CQL:** The difference between CQL and D-CQL goes beyond from transforming $\zeta_\nu(a|s) = 1 - \frac{\hat{\pi}_\beta(a|s)}{\mu(a|s)}$ to $\zeta_\nu(a|s) = \max \left( \nu,  1 - \frac{\hat{\rho}_\beta(s,a)}{\hat{\rho}_\beta(s,a_\mathcal{D})}\right)$. Indeed, $\zeta_\nu$ is integrated in the algorithm in the exponential in the practical algorithm in section $4.2$ increasing its weight. In contrast, in CQL, the maximization of the $Q$-values under $\hat{\pi}_\beta$ might be negligible compared to the minimization of the soft-maximum $\log \left( \sum_a \exp\left( Q(s,a) \right) \right)$.

**The lower-bound is loosened depending on $\mu$:** In D-CQL, note the point-wise lower-bound is kept for any choice on $\mu$, so the resulting lower bound holds for both the value function $V^{CQL}$ for all states and for all $Q^{CQL}$-values for all state-action pairs. Regarding CQL, the lower bound in the value function only holds for $\mu = \pi$, not for the commonly used uniform distribution.

 We thank the reviewer for his detailed review. We performed a global answer to the reviewers main concerns. We answer the remaining open questions in the following paragraphs.

**Results when $\nu = 1$:** It indeed implies that all actions are penalized equally whether they are ID or OOD. It seems forcing the $Q$-values to have a reasonable scale is sufficient for the Bellman update to work when the data set comes from a near-expert policy. However, as we see on Figure ($1$) and Figure ($5$), helping the agent distinguishing between ID and OOD actions is of clear interest for data sets gathered with medium policies.


**A high $\pi_\beta(a|s)$ implies a low $\zeta_\nu(a|s)$  that in turns further augments the Q-values in this area:** First, let's explain the first part: if an action $a$ is likely to be chosen by the behavioural policy, it is likely to appear in the data set. It implies that the density $\rho_\beta(s,a)$ will be high and therefore $\zeta_\nu(a|s)$ would be low. Regarding the second part, our regularization is equivalent to minimizing the KL-Divergence between $\pi_{DB}(\cdot|s)$ (with $\pi_{DB}(a|s) = \frac{\exp\left( \zeta_\nu(a|s) \ Q(s,a) \right)}{\sum_a \exp\left( \zeta_\nu(a|s) \ Q(s,a) \right)}$) and $\pi_\beta(\cdot|s)$. This Divergence is low when $\pi_{DB}(a|s)$ and $\pi_\beta(a|s)$ are close for all actions. Thus, a high $\pi_\beta(a|s)$ implies a high $\pi_{DB}(a|s)$. However, as we just saw, a high $\pi_\beta(a|s)$ also implies a low $\zeta_\nu(a|s)$. These $Q$-values must then compensate the low $\zeta_\nu(a|s)$ for $\pi_\beta$ to be close to $\pi_{DB}$.


**Comparison in the $Q$-values for CQL:** D-CQL is The D-CQL can be seen indeed a soft penalization compared to the hard one in CQL. However, our soft penalization is directly integrated in the exponential in Equation ($7$) which deeply increase its effects. This tremendous effect is highlighted on Figure ($5$) since a small change of $\nu$ has drastic effects on the agent's performance and difference in $Q$-values.

 We thank the reviewer for his detailed review. We performed a global answer to the reviewers main concerns. We answer the remaining open questions in the following paragraphs.

**Why is improving the bound important:**  CQL's lower bound is under a $\mu$-expectation, and $\mu$ is often chosen to be close to a uniform distribution. Thus, the lower bound can be highly loosened when the action space is high-dimensional and continuous. Besides, having a point-wise lower bound on the true $Q$-values is of high interest for safety considerations. For instance, if a practitioner wants to check if a policy can be played into a real system by comparing its value to a baseline, he would want to check a point-wise lower-bound especially on safety-critical systems.

**Difference between $(2)$ and $(3)$:** As the reviewer pointed out, these two equations are very similar and have the same goal: penalizing the OOD actions according to their probability of appearing in the data set. In fact, $(2)$ was our first objective. However, building a distribution $\kappa$ of the OOD actions was challenging. For example, $\zeta_\nu$ can be seen as an un-normalized distribution but building the expectation of the $Q$-values with the distribution $\kappa = \frac{1}{Z} \zeta_\nu$ (with $Z$ the unknown normalizing factor) was cumbersome. We tried MCMC methods but it was too time-consuming as the expectation must be re-computed at each training step. It was simpler to treat $\zeta_\nu$ as a weighting scheme that stays inside the $\mu$-expectation. In addition, it allows to have the weighting scheme in the exponential in Equation ($7$) increasing its importance.


**Why weighting scheme should not be $0$ on ID actions:** Intuitively, ID actions should not be penalized as their consequences are described by the dataset. We explained this phenomenon in Section $3.2$ and in Section $6$: the minimum of the resulting $Q$-values is unbounded and eventually corrupts learning. Hence, the key idea is to penalize all actions while keeping the order between actions depending on their probability of appearing in the data set intact. Since ID actions will be penalized as well, a stable solution will be found: the D-CQL penalization will minimize the $Q$-values associated to ID actions while the classic Bellman update will maximize them. This behavior is supported by our experiments (Figure 1 and 5) where setting $\nu$ to $0$ leads to poor results.

**Why there are no experiments on HalfCheetah and Walker with expert data:** Our algorithm was not designed to be used on data sets gathered from expert data, as this problem has been dealt with intensively: most of the existing batch algorithms (Lee et al., 2021; Kostrikov et al., 2021;Kumar et al., 2020) already extract expert policies out of these data sets. In fact, our experiments on Hopper showed no distinction between OOD and ID actions had to be done to have the best performance. It was sufficient to force the $Q$-values to have a reasonable scale for the Bellman update to work. This is why we focused on medium data sets where a trade off between staying close to the data set and discovering new actions must be found. It also seems more realistic to assume medium datasets for industrial applications, as there would be little interest in enhancing already expert policies.


**Criticism towards the empirical distribution:** The KDE technique was not used in the original paper (Kumar et al., 2020). In fact, one of their main point was that it can be difficult to estimate the behavioural policy when the latter is complex, so they simply estimate the expectation $\mathbb{E}_{a\sim\hat{\pi}_\beta(\cdot|s)}\left[ Q(s,a) \right]$ with the given sample(s). This technique leads to having a Dirac for $\hat{\pi}_\beta$ if only one action $a$ is present for a given state $s$ in the data set which is common when the state-action space is continuous. In this work, we argue highly efficient techniques exist to estimate densities so this drawback is alleviated. Hence, we proposed to use NF to estimate any distribution and validate this technique in Appendix $B.2$.

 We thank the reviewer for his detailed review. We performed a global answer to the reviewers main concerns. We answer the remaining open questions in the following paragraphs.

**About Theorem 1:** Indeed, assuming no sampling errors leads to a strong equality everywhere when we use the vanilla Bellman operator, as expected. However, we removed the sampling error to emphasize our main motivation: clearly penalizing $Q$-values associated to Out-Of-Distribution (OOD) actions and let the ones associated to In-Distribution (ID) actions intact. Before diving into the training and which estimator to use, we wanted to show the objective was well motivated. The sampling error can easily be integrated in the theorem with concentration qualities as done in (Kumar et al. (2020)) but do not bring additional information on the OOD actions. Nevertheless, we can include an analysis with the sampling error in the rebuttal if necessary.

**Experiments:** Each experiments has been done with $4$ seeds to take this variance into account which is the common setting. We can see some environments (Walker2d-HalfCheetah) are stable as opposed to Hopper, which is known to be quite challenging. The benchmark on (Kostrikov et al., 2021) in Figure($3$) clearly shows the difficulty of this environment. Yet, D-CQL reached a near-expert policy on *hopper-medium* for more than $200$ epochs. We then observed an unlearning effect, which is an active problem in Reinforcement Learning  (Aviral Kumar, 2021; Kumar et al., 2021). We aim to focus on this problem on future works.

 We thank the reviewers for their detailed feedback, they are deeply appreciated. Since most of the reviewers have the same concerns, we give a global answer and replied to specific concerns personally.

We identified $5$ main issues.

**Justification of $\zeta_\nu$**  Many $\zeta_\nu$ can indeed be used. 
- As proposed, the RBF function is intuitive, but is too restrictive. In particular, it assumes a symmetry between the origin, namely $a_{\mathcal{D}}$, and the selected action $a$ which may be misleading. On the other hand, the density-based weighting scheme captures the geometry of the problem.
- Another suggestion is to use the conditional density $\pi_\beta(a|s)$ instead of $\frac{\rho_\beta(s,a)}{\rho_\beta(s,a_{\mathcal{D}})}$. While this proposition captures the interesting areas and the geometry of the problem, it might be too harsh because NFs tend to have unreasonable scale for likely areas. In our experiments, the density went up to $10^{10}$. Our weighting scheme divides the joined density with $\rho_\beta(s,a_{\mathcal{D}})$ that acts as a normalization allowing for a smoother $\zeta_\nu$. This normalization is enough to keep the weighting scheme between $0$ and $1$ most of the times, as Figure~($3$) attests while keeping relevant information. In addition, learning the joined density instead of the conditional one is much simpler. 
    
   
**Justification of the NF technique** First, we would like to emphasize our main contribution is the design of a new regularization based of the density induced by the state-action pairs of the data set. We focused on Normalizing Flows because of their ability to generalize from complex high dimensional distributions, but practitioners may use any technique of their choice especially if a prior or the behavioural policy is available. Besides, some of the reviewers were concerned regarding their ability to work on high dimensional environments, but this technique is known to work well even on images (Dinh et al., 2017; Winkler et al., 2019; Kobyzev et al., 2020). Thus, Normalizing Flows should be sufficient for a plethora of data sets. We also discuss its advantages compared with fitting a Gaussian or a Mixture of Gaussian distributions in Appendix $B.2$. Finally, we would like to point out that, even if the learned density possesses small mistakes, they should be absorbed by the main update - the Bellman one - as the weighting scheme is acting inside the regularizer. 




**Justification of section $4.3$** This section brings another view to our regularization that differs from learning a lower-bound on the true $Q$-values. This point of view strengthens the role of the weighting scheme $\zeta_\nu$ as we recover the fact that it helps distinguishing the $Q$-values associated with ID from the ones associated with OOD actions. We believe this was an interesting result worth mentioning.


**Algorithm too close to CQL** Our goal was to answer to an existing issue: how to take the most out of the data set to safely learn from relevant neighbourhoods and enhance the learned policy as much as possible. Hence, we focused on CQL as it has state-of-the-art performances but has a too harsh regularization on such data sets. We also believe the introduced theoretical lower-bound could be improved. In addition, our algorithm supplies the hyper-parameter $\nu$ which allows to choose the trade-off between distinguishing ID from OOD actions. We plan to continue our experiments on environments with behavioural policies with medium performances, and provide a detailed analysis.

**Not enough baselines** We chose to compare our algorithm to CQL as it was a direct enhancement its regularization. It allows a direct comparison between both regularization, and a detailed analysis regarding the penalization of OOD actions. For instanced, we noticed data sets gathered with expert policies don't need to distinguish ID from OOD actions but only to keep the $Q$-values to a reasonable scale for the Bellman update to work. On the other hand, the problem gets interesting for data sets gathered with medium policies where a trade-off must be found.

This paper proposed a conservative Q learning approach based on density and uncertainty estimation, D-CQL. It argues it is a more effective regularization than previous conservative Q learning approaches. The performance of the proposed method is compared with other batch RL methods in several classical benchmarks. Strength:
1. The proposed method overcomes some not-so-justified designs in the original CQL. D-CQL tries to regularize the Q value in out-of-distribution actions instead of simply penalizing an arbitrary distribution that is different from behavior policy. 
2. The density estimation method might be useful for many batch RL algorithms using similar concepts/modules.

Weakness
1. The theory in this paper is meaningless. Theorem 1 shows that the fixed point solution lower bounds expected value Q^\pi assuming no sampling error (and is tight under some conditions). But assuming no sampling error the vanilla Bellman equation's solution is tight everywhere. A big part of the motivation in conservative updates is due to high sampling error in some OOD regions. I'd rather see it removed instead of a meaningless theorem.
2. The core algorithmic design is vague. The paper use uncertainty in the motivation and description of the algorithm choice. The specific form of uncertainty \zeta as a function of the density is not explained/justified at all. It's not clear to me why that function can be viewed as "uncertainty".
3. The empirical validation is not clear either. The performance is so close given the large variance of RL methods. It is not explained how the promised benefit in previous sections has helped or not. One of the four experiments shows a clearly decreasing curve and does not converge. The performance is only compared with CQL and there are many recently proposed batch RL algorithms. There is some interesting idea of regularizing OOD action's value, but the significance of both theoretical and empirical contributions is vague.

The paper aims to improve CQL by providing a better compromise between learning a good policy and staying close to the behavior policy. The authors introduce a weighting scheme into the CQL algorithm, which assigns small weights to those state-action pairs close to the dataset and large weight to the OOD actions. The effect is the gap between the underestimated values of those in-distribution state-action pairs and their true values are closer (i.e., the lower bound is tighter). The calculation of the weighting scheme requires a density estimation, and the authors propose to use normalizing flow to learn the density function. Then, the authors connect their weighted updating rule and the one by penalizing the KL divergence between learned policy and behavior policy. Experiments are presented to show the effect of their version of the CQL method. Strength: 

1. CQL concerns an important problem in offline RL — distribution shift, and the paper focuses on improving CQL, so it is an important topic. 

2. The paper has a reasonably clear logic flow. 

Weaknesses:

I feel the motivation of the work is confusing. I can understand the authors want to improve CQL somehow further. But it is never made clear: 
1) what existing problems are and why they matter; 
Is it the lower bound on the exiting CQL is too loose? Why is improving the bound important? 

2) what is the effect you want to achieve? 
Is it an offline algorithm that can learn a better policy from data generated by a poor behavior policy? 

The contribution is incremental, and I doubt the significance. As the paper cited in section 2.2, Kumar et al. (2020) propose to penalize the actions not described by the dataset, which enables a general definition of \mu. Note that the additional weighting scheme can be essentially thought of as a new type of \mu. I don’t see a clear difference between (2) and (3). Both can be considered as the specially designed action sampling distribution. 

Why is theorem 1 useful? If I understand correctly, the key thing you want to say is the additional weighting can provide a tighter bound for OOD state-action pairs or those not close to the dataset? But the first step should be figuring out the effects of having a tight/loose bound. Does it hurt optimality/convergence rate/generalization…? Even partially answering this question can better motivate the reweighting approach. I believe the proof of the theorem is a simple modification from the existing CQL work. 

The choice of the weighting scheme lacks justification. An intuitive choice is the RBF function. Furthermore, according to theorem 1, the proposed weighting is useful only when the action is OOD, i.e., the weight is 1; when the action is ID, weight should be 0, but your weighting scheme does not give zero? 

Section 4.1, the proposed method comes out suddenly. Is there any reason to choose normalizing flows? Of course, normalizing flows is a good method enabling both efficient sampling and density evaluation. In your algorithm, you only need to evaluate the density but not to sample. There should be plenty of other choices. When testing ideas, it is more natural to start with some simple methods. 

What is the purpose of section 4.3? 

I expect to see how various concrete choices of the weighting scheme can affect the distance (e.g., KL divergence) between the learned policy and the behavior policy. 

The experiments. 1. more random seeds should be tested (figure 1) — it is hard to distinguish algorithms from the current learning curves. Readers cannot see a clear message from them. 2. I expect more baselines to be compared and more domains to be tested. As I mentioned, the choices of the weighting and the way of learning density functions are not strongly motivated. In this case, I have to ask for stronger empirical results: baselines with other design choices and more domains. 3. The experiments in Fig 2 are incomplete. Why are there no experiments for half cheetah and walker with expert data? 4. Please provide reproducing details.  

The abstract says, “… with a strong theoretical guarantee.” I don’t think there is any strong theory in the paper.  

Page 3. Last paragraph. The criticism towards using empirical dataset distribution for \hat{pi}_\beta does not make sense to me. When the state/action is continuous, the empirical estimation should be kernel density estimation, which is a consistent estimator for estimating empirical distributions. The kernel can be chosen as smooth, so the KDE should have support everywhere. 

Minor: 

1. there is a nontrivial number of grammar issues/typos. Please double-check. 

2. Many sentences are confusing or logically disconnected. 

e.g., in the abstract, “A compromise between enhancing …. To alleviate this issue, … ” what issue? 

“Improving the learning process.” In what sense? Higher sample efficiency? 

“Indeed, the lack of information … ” what information? Why does it provoke overestimation? 

I believe saying “based on the amount of information collected” is inaccurate because the paper does not really introduce any information measurement. 
 Please see the main review. 

This paper focuses on the trade-off between extrapolation error and over-penalization for batch reinforcement learning. Unlike the CQL to penalize uniformly any actions  that do not belong to the dataset, this paper proposes to relax the strong penalization imposed on the out-of-distribution actions. Specifically, the proposed D-CQL considers the uncertainty on OOD actions, and only supresses the Q-values of actions which belong to OOD data with large probability. Thus the D-CQL is considered to have a better exploration when the dataset includes some suboptimal behaviors. strengths:
- the motivation and idea are intuitive and natural;
- the proposed solution (weighting scheme according to the density estimation) also make sense;
 
weakness:
- the proposed method critically depends on the accuracy of the density estimator of the in-distribution trajectory datasets, however, in this aspects, the paper does not make any significant contribution and actually this is the most difficult part of the whole pipeline as density estimation in high dimensional state-action space is very challenging. I think that an ablation study should be presented to verify that the underlying estimator really works in complicated environments.

- the empirical results in Fig.2 do not show a significant performance improvement than CQL, in fact, only in halfcheetah-medium-v0, the proposed D-CQL can achieve a better performance than CQL in the final;
- About the experiments, this paper only compares with the original algorithm, CQL, however, it's necessary to include more baselines(e.g., BAIL, BCQ, etc) to verify the effectiveness of D-CQL;

Some questions:
1. About the ablation study in Figure.1, the right figure shows the D-CQL with $\nu=1$ has a similar performance with  CQL. According to the definition in Eq.5 and Eq.7, $\nu=1$ means the final objective will penalize all the actions from  OOD or ID, it can't constrict the leanred policy to behave similarily with $\pi_\beta$, so why does it achieve such a good performance?
2. I'm confused with the explanation about Eq.9, why a high $\pi_\beta(a|s)$ implies a low $\zeta(a|s)$ that in turns further augments the Q-values in this area?
3. The D-CQL can be seen a soft penalization instead of the hard one in CQL that constrains its policy to be extremely close to the behavioural policy. So it should have a slight difference between the final Q-value, however the Fig.2 shows D-CQL often has a larger difference between ID and OOD actions. Theses results confused me.
 The idea of this paper is reasonable but to implement this idea, the problem of density estimation in high dimensional state-action space should be addressed first, which is unfortunately not convincing in this paper.

This paper presents an improved version of CQL (conservative q-learning), by using $\zeta$ from estimated data density, which was $\zeta(a|s)=1-\frac{\hat{\pi}_\beta(a|s)}{\mu(a|s)}$ in original CQL paper. The paper shows that proposed D-CQL, which uses $ \zeta = \max ( \nu , 1- \frac { \eta _ \beta (s,a) } { \eta _ \beta ( s , a _ { \mathcal { D } } ) } )$ with normalizing flow density estimator $\eta_\beta$, shows better performance than CQL in number of domains. 
Strengths:
- The paper is well written and easy to follow.
- It is a good point that original CQL actually does not satisfy $\text{supp} \mu \subset \text{supp} \hat{\pi}_\beta$.

Weaknesses:
- This paper criticizes various points of original CQL paper, such as:
   - Lower bound of CQL is loosened in the common setting where $\mu$ is a uniform distribution
   - CQL (Theorem 3.1 of CQL paper) requires $\text{supp} \mu \subset \text{supp} \hat{\pi}_\beta$, which does not hold when $\hat{\pi}_\beta$ is an empirical distribution. 
   - The regularization is equally penalizing the actions that are not appearing in the dataset, no matter their closeness with the dataset
   - I understand that by learning the density estimator can address the last point, but I am not sure whether the paper addresses the first two points.
- Newly introduced OOD weighting scheme: $ \zeta = \max ( \nu , 1- \frac { \eta _ \beta (s,a) } { \eta _ \beta ( s , a _ { \mathcal { D } } ) } )$ is not well rationalized in the paper. The constraints we had is that it should be state independent and to be bounded between 0 and 1. Why can't we simply estimate $\hat{\pi}_\beta$ and use $\max (\nu, 1-\hat{\pi}_\beta$)? Is there any reason on learning joint density rather than a conditional density? 
- I think the paper lacks concrete backgrounds for the suggestion of $\zeta$, and at least needs a thorough empirical analysis of various versions of $\zeta$, e.g. $\zeta$ that is bounded between 0,1 vs that is not bounded.
- Most of the contributions of the paper seems to be modified kinds of stuff from CQL paper, and overall the paper looks like a souped-up version of CQL.
 Overall, I believe that the contributions of the paper are only marginally significant. The arguments criticizing CQL are made but do not seem to be properly addressed by the proposed method. The results have shown that the performance has increased, but the introduced hyperparameter $\nu$ does not have a proper way of optimizing it, which is different from $\alpha$ of CQL that can be tuned to prevent the divergence of Q values. We usually expect to have a performance improvement from having another freedom on hyperparameter, and I cannot credit the performance improvement of this paper much.

Meta-review: The authors introduce a modification to CQL to use a weighting based on density estimates. In an idealized setting, they show that the estimate Q-values bound the true Q-values. Finally, they evaluate their proposed approach on a few benchmark offline RL tasks.

Generally, all reviewers felt that the results were too incremental. The theoretical result follows with light modifications from the CQL paper and even then, the implications of the result are unclear. The experimental results showed small improvements or comparable performance while requiring training a density estimator and introducing an additional hyperparameter. Furthermore, the set of tasks evaluated was limited and no comparisons to other methods than CQL were shown. 

While I appreciate the effort the authors took to investigate this improvement, at this time, the paper falls below the bar and I recommend rejection.

Decision: Reject