Density Estimation for Conservative Q-Learning

1 INTRODUCTION . Transposing the recent successes of Reinforcement Learning ( RL ) such as recommendation systems ( Rojanavasu et al. , 2005 ; Zheng et al. , 2018 ) , video games ( Mnih et al. , 2013 ) , go ( Silver et al. , 2017 ) to real-world systems is not possible without facing many challenges ( Dulac-Arnold et al. , 2021 ) . One of those challenges is that in many real-world applications , direct access to the system can be limited , and sometimes even forbidden . This can be due to various reasons , an important one being that a learning controller may incorrectly assess the implications of its actions and damage the system . Batch , or offline , Reinforcement Learning ( Lange et al. , 2012 ; Levine et al. , 2020 ) provides a framework to address those RL problems when no interaction with the system is allowed . In place , the learner is given a dataset collected under a ( possibly unknown ) behavioural policy and has to derive the most efficient policy out of this dataset . One of the main problematic encountered within this setting is the value-function over-estimation problem ( Fujimoto et al. , 2019 ; Levine et al. , 2020 ; Kumar et al. , 2019a ; Wu et al. , 2019 ) . Indeed , when the dataset only covers a small subset of the state-action space , the agent typically wrongly extrapolates the value functions related to pairs far from the dataset , denoted as Out-Of-Distribution ( OOD ) . This error is then used as a target in the learning process , leading to highly increasing estimates which may lead to a disastrous learned policy . This problem has been extensively studied in the traditional online Reinforcement Learning setting ( Sutton & Barto , 1998 ) . However in this case , this issue is quite naturally alleviated since when a value function becomes over-optimistic , it will drive the agent to visit the related state-action pair . Therefore , the agent will have the chance to directly check the consequences of such pairs and its estimation can be corrected . Some modifications , mostly relying on reducing these errors with Ensemble Learning ( van Hasselt , 2010 ; van Hasselt et al. , 2016 ; Anschel et al. , 2017 ; Lee et al. , 2021b ) , can also be used to enhance learning . In Batch RL , since no inspection can be done to investigate the accuracy of the estimates , those methods can not be used and hence classical deep Off-Policy methods dramatically fail in this setting ( Fujimoto et al. , 2019 ; Levine et al. , 2020 ) . Additional parts must be introduced to build robust and efficient agents . An important family of state-of-the art algorithms addressing this problem focus on constraining the learned policy to stay close to the dataset either by minimizing its distance to the behavioral policy ( Siegel et al. , 2020a ; Wu et al. , 2019 ; Kumar et al. , 2019b ) or by penalizing unseen state-action pairs ( Luo et al. , 2019 ; Kumar et al. , 2020 ; Yu et al. , 2021 ) . Especially , Kumar et al . ( 2020 ) propose to optimize a lower bound on the value functions and use this lower bound as a proxy in the policy optimization process . This lower bound should be tight when state-action pairs are contained in the dataset and loosened otherwise . Following this purpose , they proposed Conservative Q-Learning ( CQL ) that introduces a penalization to under-estimate value functions associated to OOD actions . On top of driving the agent to stay close to the dataset transitions , this method gives a chance to standard Off-Policy learning as estimates will have a reasonable scale during learning . Nevertheless , given the lack of information regarding the behavioural policy , their regularization remains abrupt and might be problematic in practice . First , the resulting lower bound is loosened with the recommended distributions . Second , their approach does not consider actions ’ neighbourhoods and could therefore over-penalize interesting action space ’ s regions . In this paper , we tackle these problems by introducing a novel penalization based on an estimation of the dataset ’ s probability density function instead of the behavioural policy . Besides being easier to learn that the behavioural policy , this density is able to provide information on which regions are near the dataset and thus safe to learn from , and which are the ones far from it and prone to over-estimation errors . On top of instigating those relevant information , we show that this new regularization leads to a more appropriate lower bound on the value function . We show how to integrate this penalization in a CQL-like algorithm and derive a novel algorithm called Density Conservative Q-Learning ( D-CQL ) . We finally investigate empirically the relevance of our approach . 2 PRELIMINARIES AND MOTIVATIONS . 2.1 NOTATIONS AND POLICY ITERATION . The agent-environment framework is modeled as a Markov Decision Process ( MDP ) ( S , A , r , P , γ ) defined by a state space S , an action space A , a transition kernel P : S × A → ∆ ( S ) , a reward function r : S × A → [ Rmin , Rmax ] and a discount factor γ ∈ ] 0 , 1 [ . A policy π : S → ∆ ( A ) is a decision rule mapping a state over a distribution of actions . The value of a policy is measured through the value function V π ( s ) = EP [ ∑ t=0 γ tr ( st , at ) | s0 = s , at ∼ π ( ·|st ) ] and its associated Q-value function Qπ ( s , a ) = EP [ ∑ t=0 γ tr ( st , at ) | s0 = s , a0 = a , at ∼ π ( ·|st ) ∀t ≥ 1 ] . The goal is to find the policy maximizing these value functions . Let r , Q be matrices associated to all state-action pairs and Pπ be the transition matrix induced by the policy π : Pπ ( s , s′ ) = Ea∼π ( ·|s ) P ( s′|s , a ) . Value functions can be learned by iterating the Bellman operator defined for the Q-function as BπQ = r + γPπQ where PπQ ( s , a ) = Es′∼P ( ·|s , a ) , a′∼π ( ·|s ) [ Q ( s′ , a′ ) ] . This operator is a γ-contraction ( Puterman , 1994 ) , hence having a unique fixed point Qπ . Many classic recent algorithms rely on the Policy Iteration scheme ( Lillicrap et al. , 2016 ; Fujimoto et al. , 2018 ; Haarnoja et al. , 2018 ) , where the agent alternates between Policy Evaluation ( PE ) with the computation of Qπ and Policy Improvement ( PI ) by maximizing the learned Q-values Qk+1 ← arg min Q Bπ k Qk −Q , ( policy evaluation ) πk+1 ← arg max π Ea∼π [ Qk+1 ( · , a ) ] . ( policy improvement ) This iterative process converges towards the optimal policy ( Sutton & Barto , 1998 ; Santos & Rust , 2004 ) . Since no knowledge is assumed on the environment , these steps are commonly solved using samples from a dataset D = { ( si , ai , ri ) Ni=1 } . The expectation under P ( ·|s , a ) is now estimated using sample and leads to the empirical Bellman operator B̂ . We slightly abuse notations and consider D is also a distribution on S ×A× S. The procedure is now under the dataset expectation Qk+1 ← arg min Q Es , a , s′∼D [ ( B̂π k Qk ( s , a ) −Q ( s , a ) ) 2 ] , ( approximate policy evaluation ) πk+1 ← arg max π Es∼D , a∼π [ Qk+1 ( s , a ) ] . ( approximate policy improvement ) The Q-values and the policy π are commonly estimated using Neural Networks that are trained with gradient optimization methods . In the approximate policy evaluation step , an important aspect is the expectation of actions from π ( ·|s′ ) in the bootstrapped target . This is where OOD actions may appear and become a decisive topic : extrapolation errors emerge , are then back-propagated to eventually lead to highly over-optimistic estimates . In the Online setting , the dataset is constantly updated with samples gathered with the learned policy , coping with this problem . Note there are no OOD states in the Bellman update as the empirical operator only depends on a seen state s′ . They might appear on model-based algorithms ( Yu et al. , 2021 ) but are for now out of the scope of model-free algorithms . An interesting future direction would be to consider them as well . In this work , we focus on the Batch setting where the agent can not interact with the environment . The agent dataset is now fixed and has been gathered with an unknown behavioural policy πβ . No knowledge is assumed on πβ as the dataset can come from different sources : optimal control , human or a mixture of them ( Fu et al. , 2020 ; Gülçehre et al. , 2020 ) . Actions that belong to the dataset distribution given the state s are denoted as In-Distribution ( ID ) and the ones far from it are Out-Of-Distribution ( OOD ) . A major challenge is to find a good trade-off between staying close to the dataset to avoid extrapolation errors and taking some liberty to overcome the suboptimality of the behaviour policy . 2.2 CONSERVATIVE Q-LEARNING . This approach addresses the Q-values over-estimation problem exacerbated in Batch RL . As a reminder , the Policy Evaluation step relies on a bootstrapped target that back-propagates any extrapolation error during learning . This error may be of high importance when the consequences of a given state-action pair are unknown , and could lead to highly optimistic estimates . This problem is intensified in Batch RL . First , the dataset often describes a small subset of the state-action space S ×A thus OOD actions are very likely to appear . Second , the agent never gets the chance the visit the state-action pairs related to high Q-values and can not correct its possibly wrong estimations . At this step , Kumar et al . ( 2020 ) propose to penalize the actions not described by the dataset while keeping intact actions from the dataset with proper information . This greatly stabilizes learning : the agent does not have to deal with drastically high estimates and therefore gives the agent a chance to learn accurate Q-values thanks to the empirical Bellman operator . It also implicitly drives the agent to favor state-action pairs described in the dataset . This penalization translates into minimizing the Q-values on an arbitrary distribution µ ( e.g . uniform ) and maximizing them on πβ formally expressed as follows Q̂k+1 ← arg min Q αEs∼D [ ( Ea∼µ ( ·|s ) [ Q ( s , a ) ] − Ea∼πβ ( ·|s ) [ Q ( s , a ) ] ) ] +Es , a , s′∼D [ ( B̂πkQ̂k ( s , a ) −Q ( s , a ) ) 2 ] . ( 1 ) Assuming the absence of sampling errors and suppµ ⊂ suppπβ , this update provides a lowerbound on the expected Q-values on the distribution µ : Ea∼µ [ QCQL ( s , a ) ] ≤ Ea∼µ [ Qπ ( s , a ) ] for all s ∈ S with α > 0 . When µ is the current policy π , the lower bound is in the value function , that is V CQL ( s ) ≤ V π ( s ) . This property may be highly desirable on applications where the agent must check if the current value function is above a certain threshold such as Constrained Policy Optimization ( Achiam et al. , 2017 ) and Conservative Exploration ( Garcelon et al. , 2020 ) . However , for the common setting where µ is a uniform distribution over the action space , the lower bound is loosened . While a certain control in the Q-values is kept , both the value and point-wise inequalities are lost . Another point of attention is the maximization term relying on the knowledge of the behavioral policy . In the general setting , it is out of reach as the dataset may be gathered in various ways . Kumar et al . ( 2020 ) propose to use the empirical dataset distribution π̂β ( a|s ) = ∑ s′ , a′∈D 1 [ s′=s , a′=a ] ∑ s′∈D 1 [ s′=s ] , that may become a Dirac over the dataset actions especially when the action space A is continuous . Not only does this breaks the assumption suppµ ⊂ supp π̂β ( as in practice suppπβ 6⊂ supp π̂β ) , it also might be problematic on the penalization itself : the distance between the selected actions and the dataset is no longer considered . In other words , actions that are not appearing in the dataset will be equally penalized no matter their closeness with the dataset D. Thus , the regularization pushes the agent to reproduce the behavioural policy , not enhance it .

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