Conservative Offline Distributional Reinforcement Learning

1 Introduction . In many applications of reinforcement learning , actively gathering data through interactions with the environment can be risky and unsafe . Offline ( or batch ) reinforcement learning ( RL ) avoids this problem by learning a policy solely from historical data ( called observational data ) [ 9 , 22 , 23 ] . A shortcoming of most existing approaches to offline RL [ 11 , 46 , 20 , 21 , 48 , 18 ] is that they are designed to maximize the expected value of the cumulative reward ( which we call the return ) of the policy . As a consequence , they are unable to quantify risk and ensure that the learned policy acts in a safe way . In the online setting , there has been recent work on distributional RL algorithms [ 7 , 6 , 27 , 38 , 17 ] , which instead learn the full distribution over future returns . They can use this distribution to plan in a way that avoids taking risky , unsafe actions . Furthermore , when coupled with deep neural network function approximation , they can learn better state representations due to the richer distributional learning signal [ 4 , 26 ] , enabling them to outperform traditional RL algorithms even on the risk-neutral , expected return objective [ 4 , 7 , 6 , 47 , 14 ] . We propose Conservative Offline Distributional Actor-Critic ( CODAC ) , which adapts distributional RL to the offline setting . A key challenge in offline RL is accounting for high uncertainty on out-of-distribution ( OOD ) state-action pairs for which observational data is limited [ 23 , 20 ] ; the value estimates for these state-action pairs are intrinsically high variance , and may be exploited by the policy without correction due to the lack of online data gathering and feedback . We build on conservative Q-learning [ 21 ] , which penalizes Q values for OOD state-action pairs to ensure that 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the learned Q-function lower bounds the true Q-function . Analogously , CODAC uses a penalty to ensure that the quantiles of the learned return distribution lower bound those of the true return distribution . Crucially , the lower bound is data-driven and selectively penalizes the quantile estimates of state-actions that are less frequent in the offline dataset ; see Figure 1 . We prove that for finite MDPs , CODAC converges to an estimate of the return distribution whose quantiles uniformly lower bound the quantiles of the true return distribution ; in addition , this data-driven lower bound is tight up to the approximation error in estimating the quantiles using finite data . Thus , CODAC obtains a uniform lower bound on all integrations of the quantiles , including the standard RL objective of expected return , the risk-sensitive conditional-value-at-risk ( CVaR ) objective [ 35 ] , as well as many other risk-sensitive objectives . We additionally prove that CODAC expands the gap in quantile estimates between in-distribution and OOD actions , thus avoiding overconfidence when extrapolating to OOD actions [ 11 ] . Our theoretical guarantees rely on novel techniques for analyzing the distributional Bellman operator , which is challenging since it acts on the infinite-dimensional function space of return distributions ( whereas the traditional Bellman operator acts on a finite-dimensional vector space ) . We provide several novel results that may be of independent interest ; for instance , our techniques can be used to bound the error of the fixed-point of the empirical distributional Bellman operator ; see Appendix A.6 . Finally , to obtain a practical algorithm , CODAC builds on existing distributional RL algorithms by integrating conservative return distribution estimation into a quantile-based actor-critic framework . In our experiments , we demonstrate the effectiveness of CODAC on both risk-sensitive and riskneutral RL . First , on two novel risk-sensitive robot navigation tasks , we show that CODAC successfully learns risk-averse policies using offline datasets collected purely from a risk-neutral agent , a challenging task that all our baselines fail to solve . Next , on the D4RL Mujoco suite [ 10 ] , a popular offline RL benchmark , we show that CODAC achieves state-of-art results on both the original risk-neutral version as well a modified risk-sensitive version [ 43 ] . Finally , we empirically show that CODAC computes quantile lower-bounds and gap-expanded quantiles even on high-dimensional continuous-control problems , validating our key theoretical insights into the effectiveness of CODAC . Related work . There has been growing interest in offline ( or batch ) RL [ 22 , 23 ] . The key challenge in offline RL is to avoid overestimating the value of out-of-distribution ( OOD ) actions rarely taken in the observational dataset [ 40 , 44 , 20 ] . The problem is that policy learning optimizes against the value estimates ; thus , even if the estimation error is i.i.d. , policy optimization biases towards taking actions with high variance value estimates , since some of these values will be large by random chance . In risk-sensitive or safety-critical settings , these actions are exactly the ones that should be avoided . One solution is to constrain the learned policy to take actions similar to the ones in the dataset ( similar to imitation learning ) —e.g. , by performing support matching [ 46 ] or distributional matching [ 20 , 12 ] . However , these approaches tend to perform poorly when data is gathered from suboptimal policies . An alternative is to regularize the Q-function estimates to be conservative at OOD actions [ 21 , 48 , 18 ] . CODAC builds on these approaches , but obtains conservative estimates of all quantile values of the return distribution rather than just the expected return . Traditionally , the literature on off-policy evaluation ( OPE ) [ 32 , 16 , 39 , 25 , 37 ] aims to estimate the expected return of a policy using precollected offline data ; CODAC proposes an OPE procedure amenable to all objectives that can be expressed as integrals of the return quantiles . Consequently , our fine-grained approach not only enables risk-sensitive policy learning , but also improves performance on risk-neutral domains . In particular , CODAC builds on recent works on distributional RL [ 4 , 7 , 6 , 47 ] , which parameterize and estimate the entire return distribution instead of just a point estimate of the expected return ( i.e. , the Q-function ) [ 29 , 28 ] . Distributional RL algorithms have been shown to achieve state-of-art performance on Atari and continuous control domains [ 14 , 3 ] ; intuitively , they provide richer training signals that stabilize value network training [ 4 ] . Existing distributional RL algorithms parameterize the policy return distribution in many different ways , including canonical return atoms [ 4 ] , the expectiles [ 36 ] , the moments [ 31 ] , and the quantiles [ 7 , 6 , 47 ] . CODAC builds on the quantile approach due to its suitability for risk-sensitive policy optimization . The quantile representation provides an unified framework for optimizing different objectives of interest [ 7 ] , such as the riskneutral expected return , and a family of risk-sensitive objectives representable by the quantiles ; this family includes , for example , the conditional-value-at-risk ( CVaR ) return [ 35 , 5 , 38 , 17 ] , the Wang measure [ 45 ] , and the cumulative probability weighting ( CPW ) metric [ 42 ] . Recent work has provided theoretical guarantees on learning CVaR policies [ 17 ] ; however , their approach can not provide bounds on the quantiles of the estimated return distribution , which is significantly more challenging since there is no closed-form expression for the Bellman update on the return quantiles . Finally , there has been some recent work adapting distributional RL to the offline setting [ 1 , 43 ] . First , REM [ 1 ] builds on QR-DQN [ 7 ] , an online distributional RL algorithm ; however , REM can be applied to regular DQN [ 29 ] and does not directly utilize the distributional aspect of QR-DQN . The closest work to ours is ORAAC [ 43 ] , which uses distributional RL to learn a CVaR policy in the offline setting . ORAAC uses imitation learning to avoid OOD actions and stay close to the data distribution . As discussed above , imitation learning strategies can perform poorly unless the dataset comes from an optimal policy ; in our experiments , we find that CODAC significantly outperforms ORAAC . Furthermore , unlike ORAAC , we provide theoretical guarantees on our approach . 2 Background . Offline RL . Consider a Markov Decision Process ( MDP ) [ 33 ] ( S , A , P , R , ) , where S is the state space , A is the action space , P ( s0 | s , a ) is the transition distribution , R ( r | s , a ) is the reward distribution , and 2 ( 0 , 1 ) is the discount factor , and consider a stochastic policy ⇡ ( a | s ) : S ! ( A ) . A rollout using ⇡ from state s using initial action a is the random sequence ⇠ = ( ( s0 , a0 , r0 ) , ( s1 , a1 , r1 ) , ... ) such that s0 = s , a0 = a , at ⇠ ⇡ ( · | st ) ( for t > 0 ) , rt ⇠ R ( · | st , at ) , and st+1 ⇠ P ( · | st , at ) ; we denote the distribution over rollouts by D⇡ ( ⇠ | s , a ) . The Q-function Q⇡ : S ⇥ A ! R of ⇡ is its expected discounted cumulative return Q⇡ ( s , a ) = ED⇡ ( ⇠|s , a ) [ P 1 t=0 trt ] . Assuming the rewards satisfy rt 2 [ Rmin , Rmax ] , then we have Q⇡ ( s , a ) 2 [ Vmin , Vmax ] ✓ [ Rmin/ ( 1 ) , Rmax/ ( 1 ) ] . A standard goal of reinforcement learning ( RL ) , which we call risk-neutral RL , is to learn the optimal policy ⇡⇤ such that Q⇡ ⇤ ( s , a ) Q⇡ ( s , a ) for all s 2 S , a 2 A and all ⇡ . In offline RL , the learning algorithm only has access to a fixed dataset D : = { ( s , a , r , s0 ) } , where r ⇠ R ( · | s , a ) and s0 ⇠ P ( · | s , a ) . The goal is to learn the optimal policy without any interaction with the environment . Though we do not assume that D necessarily comes from a single behavior policy , we define the empirical behavior policy to be ⇡̂ ( a | s ) : = P s0 , a02D ( s 0=s , a0=a ) P s02D ( s 0=s ) . With slight abuse of notation , we write ( s , a , r , s0 ) ⇠ D to denote a uniformly random sample from the dataset . Also , in this paper , we broadly refer to actions not drawn from ⇡̂ ( · | s ) ( i.e. , low probability density ) as out-of-distribution ( OOD ) . Fitted Q-evaluation ( FQE ) [ 9 , 34 ] uses Q-learning for offline RL , which leverages the fact that Q⇡ = T ⇡Q⇡ is the unique fixed point of the Bellman operator T ⇡ : R|S||A| ! R|S||A| defined by T ⇡Q ( s , a ) = ER ( r|s , a ) [ r ] + · EP⇡ ( s0 , a0|s , a ) [ Q ( s0 , a0 ) ] , where P⇡ ( s0 , a0|s , a ) = P ( s0 | s , a ) ⇡ ( a0 | s0 ) . In the offline setting , we do not have access to T ⇡ ; instead , FQE uses an approximation T̂ ⇡ obtained by replacing R and P in T ⇡ with estimates R̂ and P̂ based on D. Then , we can estimate Q⇡ by starting from an arbitrary Q̂0 and iteratively computing Q̂k+1 : = argmin Q L ( Q̂ , T̂ ⇡Q̂k ) where L ( Q , Q0 ) = ED ( s , a ) ⇥ ( Q ( s , a ) Q0 ( s , a ) ) 2 ⇤ . Assuming we search over the space of all possible Q ( i.e. , do not use function approximation ) , then the minimizer is Q̂k+1 = T̂ ⇡Q̂k , so Q̂k = ( T̂ ⇡ ) kQ0 . If T̂ ⇡ = T ⇡ , then limk ! 1 Q̂k = Q⇡ . Distributional RL . In distributional RL , the goal is to learn the distribution of the discounted cumulative rewards ( i.e. , returns ) [ 4 ] . Given a policy ⇡ , we denote its return distribution as the random variable Z⇡ ( s , a ) = P 1 t=0 trt , which is a function of a random rollout ⇠ ⇠ D⇡ ( · | s , a ) ; note that Z⇡ includes three sources of randomness : ( 1 ) the reward R ( · | s , a ) , ( 2 ) the transition P ( · | s , a ) , and ( 3 ) the policy ⇡ ( · | s ) . Also , note that Q⇡ ( s , a ) = ED⇡ ( ⇠|s , a ) [ Z⇡ ( s , a ) ] . Analogous to Q-function Bellman operator , the distributional Bellman operator for ⇡ is T ⇡Z ( s , a ) D : = r + Z ( s0 , a0 ) where r ⇠ R ( · | s , a ) , s0 ⇠ P ( · | s , a ) , a0 ⇠ ⇡ ( · | s0 ) , ( 1 ) where D= indicates equality in distribution . As with Q⇡ , Z⇡ is the unique fixed point of T ⇡ in Eq . 1 . Next , let FZ ( s , a ) ( x ) : [ Vmin , Vmax ] ! [ 0 , 1 ] be the cumulative density function ( CDF ) for return distribution Z ( s , a ) , and FR ( s , a ) be the CDF of R ( · | s , a ) Then , we have the following equality , which captures how the distributional Bellman operator T ⇡ operates on the CDF FZ ( s , a ) [ 17 ] : FT ⇡Z ( s , a ) ( x ) = X s0 , a0 P⇡ ( s0 , a0 | s , a ) Z FZ ( s0 , a0 ) ✓ x r ◆ dFR ( s , a ) ( r ) . ( 2 ) Let X and Y be two random variables . Then , the quantile function ( i.e. , inverse CDF ) F 1X of X is F 1X ( ⌧ ) : = inf { x 2 R | ⌧ FX ( x ) } , and the p-Wasserstein distance between X and Y is Wp ( X , Y ) = ( R 1 0 |F 1 Y ( ⌧ ) F 1 X ( ⌧ ) |pd⌧ ) 1/p . Then , the distributional Bellman operator T ⇡ is a -contraction in the Wp [ 4 ] —i.e. , letting d̄p ( Z1 , Z2 ) : = sups , aWp ( Z1 ( s , a ) , Z2 ( s , a ) ) be the largest Wasserstein distance over ( s , a ) , and Z = { Z : S ⇥A ! P ( R ) | 8 ( s , a ) . E [ |Z ( s , a ) |p ] < 1 } be the space of distributions over R with bounded p-th moment , then d̄p ( T ⇡Z1 , T ⇡Z2 ) d̄p ( Z1 , Z2 ) ( 8Z1 , Z2 2 Z ) . ( 3 ) As a result , Z⇡ may be obtained by iteratively applying T ⇡ to an initial distribution Z . As before , in the offline setting , we can approximate T ⇡ by T̂ ⇡ using D. Then , we can compute Z⇡ ( represented as F 1Z ( s , a ) ; see below ) by starting from an arbitrary Ẑ 0 , and iteratively computing Ẑk+1 = argmin Z Lp ( Z , T̂ ⇡Ẑk ) where Lp ( Z , Z 0 ) = ED ( s , a ) [ Wp ( Z ( s , a ) , Z 0 ( s , a ) ) p ] . ( 4 ) We call this procedure fitted distributional evaluation ( FDE ) . One distributional RL algorithmic framework is quantile-based distributional RL [ 7 , 6 , 47 , 27 , 38 , 43 ] , where the return distribution Z is represented by its quantile function F 1Z ( s , a ) ( ⌧ ) : [ 0 , 1 ] ! R. Given a distribution g ( ⌧ ) over [ 0 , 1 ] , the distorted expectation of Z is g ( Z ( s , a ) ) = R 1 0 F 1 Z ( s , a ) ( ⌧ ) g ( ⌧ ) d⌧ , and the corresponding policy is ⇡g ( s ) : = argmaxa g ( Z ( s , a ) ) [ 7 ] . If g = Uniform ( [ 0 , 1 ] ) , then Q⇡ ( s , a ) = g ( Z ( s , a ) ) ; alternatively , g = Uniform ( [ 0 , ⇠ ] ) corresponds to the CVaR [ 35 , 5 , 6 ] objective , where only the bottom ⇠-percentile of the return is considered . Additional risk-sensitive objectives are also compatible . For example , CPW [ 42 ] amounts to g ( ⌧ ) = ⌧ / ( ⌧ + ( 1 ⌧ ) ) 1 , and Wang [ 45 ] has g ( ⌧ ) = FN ( F 1N ( ⌧ ) + ) , where FN is the standard Gaussian CDF . A drawback of FDE is that it does not account for estimation error , especially for pairs ( s , a ) that rarely appear in the given dataset D ; thus , Ẑk ( s , a ) may be an overestimate of Zk ( s , a ) [ 12 , 20 , 21 ] , even in distributional RL ( since the learned distribution does not include randomness in the dataset ) [ 14 , 3 ] . Importantly , since we act by optimizing with respect to Ẑk ( s , a ) , the optimization algorithm will exploit these errors , biasing towards actions with higher uncertainty , which is the opposite of what is desired . In Section 3 , we propose and analyze a penalty designed to avoid this issue .

Summary: At its core, the paper introduces an algorithm called conservative distributional evaluation (CDE). In the tabular settings, the CDE algorithms theoretically guarantees to converge to a fixed point whose quantile function lower bounds that of the true return distribution, also its fixed point is more conservative for rare actions in the dataset (the property that is called gap-expanding). These are desired properties for offline scenarios with full coverage (see the comments bellow). CDE is then used in an actor-critic framework, which is called Conservative Offline Distributional Actor Critic (CODAC). This gets a min-max objective: the inner loop chooses the current policy to maximize the CDE objective, and the outer loop minimizes the CDE objective for this policy. Some experimental results are followed to support the theory.