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