COptiDICE: Offline Constrained Reinforcement Learning via Stationary Distribution Correction Estimation

1 INTRODUCTION . Reinforcement learning ( RL ) has shown great promise in a wide range of domains , such as complex games ( Mnih et al. , 2015 ; Silver et al. , 2017 ) and robotic control ( Lillicrap et al. , 2016 ; Haarnoja et al. , 2018 ) . However , the need to interact with the environment during learning hinders its widespread application to many real-world problems for which executing exploratory behavior in the environment is costly or dangerous . Offline RL ( also known as batch RL ) ( Lange et al. , 2012 ; Levine et al. , 2020 ) algorithms sidestep this problem and perform policy optimization solely from a set of pre-collected data . The use of existing ( offline ) data can make offline reinforcement learning applicable to real world systems and and has led to a sharp increase in interest in this paradigm . Recent works on offline RL , however , mostly assume that the environment is modeled as a Markov decision process ( MDP ) , and standard offline RL algorithms focus on reward-maximization only ( Fujimoto et al. , 2019 ; Wu et al. , 2019 ; Kumar et al. , 2019 ; Siegel et al. , 2020 ; Wang et al. , 2020 ; Kumar et al. , 2020 ) . In contrast , in real-world domains it is common that the behavior of the agent is subject to additional constraints beyond the reward . Consider , for example , autonomous driving or an industrial robot in a factory . Some behaviors may damage the agent itself or its surroundings , and safety constraints should thus be considered as part of the objective of a suitable reinforcement learning system . One of the ways to mathematically characterize a constrained RL problem is through the formalism of constrained Markov Decision Processes ( CMDP ) ( Altman , 1999 ) . In CMDPs taking an action incurs a cost as well as a reward , and the goal is to maximize the expected long-term reward while satisfying a bound on the expected long-term cost . In this work , we aim to solve the constrained decision making problem in the offline RL setting , to enable deployment in various safety-critical domains where direct learning interactions are infeasible . Offline constrained RL inherits the difficulties of offline unconstrained RL , while introducing additional challenges . First , since the target policy being optimized deviates from the data-collection policy with no further data collection , distribution shift becomes the central difficulty . To mitigate the distributional shift , existing offline RL methods frequently adopt the pessimism principle : either by explicit policy ( and critic ) regularization that penalizes deviation from the data-collection policy ( Jaques et al. , 2019 ; Wu et al. , 2019 ; Kumar et al. , 2019 ; Siegel et al. , 2020 ; Wang et al. , 2020 ; Lee et al. , 2020 ; Kostrikov et al. , 2021 ) or by reward penalty to the uncertain state-action regions ( Kumar et al. , 2020 ; Kidambi et al. , 2020 ; Yu et al. , 2020 ) . Second , in offline constrained RL , the computed policy should satisfy the given cost constraints when it is deployed to the real environment . Unfortunately off-policy policy evaluation inherently has estimation errors , and it is therefore difficult to ensure that a policy estimated from a finite dataset will satisfy the constraint when executed in the environment . In addition , constrained policy optimization usually involves an additional optimization for the Lagrange multiplier associated with the cost constraints . Actor-critic-based constrained RL algorithms thus have to solve triple ( i.e . critic , actor , Lagrange multiplier ) intertwined optimization problems ( Borkar , 2005 ; Tessler et al. , 2019 ) , which can be very unstable in practice . In this paper , we present an offline constrained RL algorithm that optimizes the state-action stationary distribution directly , rather than the Q-function or the policy . We show that such treatment obviates the need for multiple estimators for the value and the policy , yielding a single optimization objective that is practically solvable . Still , naively constraining the cost value may result in severe constraint violation in the real environment , as we demonstrate empirically . We thus propose a method to constrain the upper bound of the cost value , aiming to compute a policy more robust in constraint violation , where the upper bound is computed in a way motivated by a recent advance in off-policy confidence interval estimation ( Dai et al. , 2020 ) . Our algorithm , Offline Constrained Policy Optimization via stationary DIstribution Correction Estimation ( COptiDICE ) , estimates the stationary distribution corrections of the optimal policy that maximizes rewards while constraining the cost upper bound , with the goal of yielding a cost-conservative policy for better actual constraint satisfaction . COptiDICE computes the upper bound of the cost efficiently by solving an additional minimization problem . Experimental results show that COptiDICE attains a better policy in terms of constraint satisfaction and reward-maximization , outperforming several baseline algorithms . 2 BACKGROUND . A Constrained Markov Decision Process ( CMDP ) ( Altman , 1999 ) is an extension of an MDP , formally defined by a tuple M = 〈S , A , T , R , C = { Ck } 1 .. K , ĉ = { ĉk } 1 .. K , p0 , γ〉 , where S is the set of states , A is the set of actions , T : S × A → ∆ ( S ) is a transition probability , R : S × A → R is the reward function , Ck : S × A → R is the k-th cost function with its corresponding threshold ĉk ∈ R , p0 ∈ ∆ ( S ) is the initial state distribution , and γ ∈ ( 0 , 1 ] is the discount factor . A policy π : S → ∆ ( A ) is a mapping from state to distribution over actions . We will express solving CMDPs in terms of stationary distribution . For a given policy π , the stationary distribution dπ is defined by : dπ ( s , a ) = ( 1− γ ) ∞∑ t=0 γt Pr ( st = s , at = a ) if γ < 1 , lim T→∞ 1 T+1 T∑ t=0 Pr ( st = s , at = a ) if γ = 1 , ( 1 ) where s0 ∼ p0 , at ∼ π ( st ) , st+1 ∼ T ( st , at ) for all timesteps t ≥ 0 . We define the value of the policy as VR ( π ) : = E ( s , a ) ∼dπ [ R ( s , a ) ] ∈ R and VC ( π ) : = E ( s , a ) ∼dπ [ C ( s , a ) ] ∈ RK . Constrained RL aims to learn an optimal policy that maximizes the reward while bounding the costs up to the thresholds by interactions with the environment : max π VR ( π ) s.t . VCk ( π ) ≤ ĉk ∀k = 1 , . . . , K ( 2 ) Lagrangian relaxation is typically employed to solve Eq . ( 2 ) , leading to the unconstrained problem : min λ≥0 max π VR ( π ) − λ > ( VC ( π ) − ĉ ) ( 3 ) where λ ∈ RK is the Lagrange multiplier and ĉ ∈ RK is the vector-valued cost threshold . The inner maximization in Eq . ( 3 ) corresponds to computing an optimal policy that maximizes scalarized rewards R ( s , a ) − λ > C ( s , a ) , due to the linearity of the value function with respect to the reward function . The outer minimization corresponds to balancing the cost penalty in the scalarized reward function : if the current policy is violating the k-th cost constraint , λk increases so that the cost is penalized more in the scalarized reward , and vice versa . In the offline RL setting , online interaction with the environment is not allowed , and the policy is optimized using the fixed offline dataset D = { ( s0 , s , a , r , c , s′ ) i } Ni=1 collected with one or more ( unknown ) data-collection policies . The empirical distribution of the dataset is denoted as dD , and we will abuse the notation dD for s ∼ dD , ( s , a ) ∼ dD , ( s , a , s′ ) ∼ dD . We abuse ( s0 , s , a , s′ ) ∼ dD for s0 ∼ p0 , ( s , a , s′ ) ∼ dD . We denote the space of data samples ( s0 , s , a , s′ ) as X . A naive way to solve ( 3 ) in an offline manner is to adopt an actor-critic based offline RL algorithm for maxπ VR−λ > C ( π ) while jointly optimizing λ . However , the intertwined training procedure of offpolicy actor-critic algorithms often suffers from instability due to the compounding error incurred by bootstrapping out-of-distribution action values in an offline RL setting ( Kumar et al. , 2019 ) . The instability would be exacerbated when the nested optimization for λ is added . 3 OFFLINE CONSTRAINED RL VIA STATIONARY DISTRIBUTION CORRECTION ESTIMATION . In this section , we present our offline constrained RL algorithm , Constrained Policy Optimization via stationary DIstribution Correction Estimation ( COptiDICE ) . The derivation of our algorithm starts by augmenting the standard linear program for CMDP ( Altman , 1999 ) with an additional f -divergence regularization1 : max d E ( s , a ) ∼d [ R ( s , a ) ] − αDf ( d||dD ) ( 4 ) s.t . E ( s , a ) ∼d [ Ck ( s , a ) ] ≤ ĉk ∀k = 1 , . . . , K ( 5 ) ∑ a′ d ( s′ , a′ ) = ( 1− γ ) p0 ( s′ ) + γ ∑ s , a d ( s , a ) T ( s′|s , a ) ∀s′ ( 6 ) d ( s , a ) ≥ 0 ∀s , a , ( 7 ) where Df ( d||dD ) : = E ( s , a ) ∼dD [ f ( d ( s , a ) dD ( s , a ) ) ] is the f -divergence between the distribution d and the dataset distribution dD , and α > 0 is the hyperparameter that controls the degree of pessimism , i.e . how much we penalize the distribution shift , a commonly adopted principle for offline RL ( Nachum et al. , 2019b ; Kidambi et al. , 2020 ; Yu et al. , 2020 ; Lee et al. , 2021 ) . We assume that dD > 0 and f is a strictly convex and continuously differentiable function with f ( 1 ) = 0 . Note that when α = 0 , the optimization ( 4-7 ) reduces to the standard linear program for CMDPs . The Bellman-flow constraints ( 6-7 ) ensure that d is the stationary distribution of a some policy , where d ( s , a ) can be interpreted as a normalized discounted occupancy measure of ( s , a ) . Thus , we seek the stationary distribution of an optimal policy that maximizes the reward value ( 4 ) while bounding the cost values ( 5 ) . Once the optimal solution d∗ has been estimated , its corresponding optimal policy is obtained by π∗ ( a|s ) = d ∗ ( s , a ) ∑ a′ d ∗ ( s , a′ ) . Now , consider the Lagrangian for the constrained optimization problem ( 4-7 ) : min λ≥0 , ν max d≥0 E ( s , a ) ∼d [ R ( s , a ) ] − αDf ( d||dD ) − K∑ k=1 λk ( E ( s , a ) ∼d [ Ck ( s , a ) ] − ĉk ) − ∑ s′ ν ( s′ ) [ ∑ a′ d ( s′ , a′ ) − ( 1− γ ) p0 ( s′ ) − γ ∑ s , a d ( s , a ) T ( s′|s , a ) ] ( 8 ) where λ ∈ RK+ is the Lagrange multiplier for the cost constraints ( 5 ) , and ν ( s ) ∈ R is the Lagrange multiplier for the Bellman flow constraints ( 6 ) . Solving ( 8 ) in its current form requires evaluation of T ( s′|s , a ) for ( s , a ) ∼ d , which is not accessible in the offline RL setting . To make the optimization tractable , we rearrange the terms so that the direct dependence on d is eliminated , while introducing new optimization variables w that represent stationary distribution corrections : min λ≥0 , ν max d≥0 E ( s , a ) ∼d s′∼T ( s , a ) [ R ( s , a ) − λ > C ( s , a ) + ν ( s′ ) − ν ( s ) ] − αE ( s , a ) ∼dD [ f ( d ( s , a ) dD ( s , a ) ) ] + ( 1− γ ) Es0∼p0 [ ν ( s0 ) ] + λ > ĉ = min λ≥0 , ν max w≥0 E ( s , a ) ∼dD [ w ( s , a ) eλ , ν ( s , a ) − αf ( w ( s , a ) ) ] + ( 1− γ ) Es0∼p0 [ ν ( s0 ) ] + λ > ĉ ( 9 ) 1For simplicity , we only consider discounted CMDPs ( γ < 1 ) here , but the result can be generalized to the case of average reward-cost CMDPs ( γ = 1 ) ( Appendix A ) . where eλ , ν ( s , a ) : = R ( s , a ) − λ > C ( s , a ) + γEs′∼T ( s , a ) [ ν ( s′ ) ] − ν ( s ) is the advantage function by regarding ν as a state value function , and w ( s , a ) : = d ( s , a ) dD ( s , a ) is the stationary distribution correction . Every term in Eq . ( 9 ) can be estimated from samples in the offline dataset D : min ν , λ≥0 max w≥0 E ( s0 , s , a , s′ ) ∼dD [ w ( s , a ) êλ , ν ( s , a , s ′ ) − αf ( w ( s , a ) ) + ( 1− γ ) ν ( s0 ) ] + λ > ĉ ( 10 ) where êλ , ν ( s , a , s′ ) : = R ( s , a ) −λ > C ( s , a ) +γν ( s′ ) −ν ( s ) is the advantage estimate using a single sample . As a consequence , ( 10 ) can be optimized in a fully offline manner . Moreover , exploiting the strict convexity of f , we can further derive a closed-form solution for the inner maximization in ( 9 ) as follows . All the proofs can be found in Appendix B . Proposition 1 . For any ν and λ , the closed-form solution of the inner maximization problem in ( 9 ) is given by : w∗λ , ν ( s , a ) = ( f ′ ) −1 ( 1 αeλ , ν ( s , a ) ) + where x+ = max ( 0 , x ) ( 11 ) Finally , by plugging the closed-form solution ( 11 ) into ( 9 ) , we obtain the following convex minimization problem : min λ≥0 , ν L ( λ , ν ) = E ( s , a ) ∼dD [ w∗λ , ν ( s , a ) eλ , ν ( s , a ) − αf ( w∗λ , ν ( s , a ) ) ] + ( 1− γ ) Es0∼p0 [ ν ( s0 ) ] + λ > ĉ ( 12 ) To sum up , by operating in the space of stationary distributions , constrained ( offline ) RL can in principle be solved by solving a single convex minimization ( 12 ) problem . This is in contrast to existing constrained RL algorithms that manipulate both Q-function and policy , and thus require solving triple optimization problems for the actor , the critic , and the cost Lagrange multiplier with three different objective functions . Note also that when λ is fixed and treated as a constant , ( 12 ) reduces to OptiDICE ( Lee et al. , 2021 ) for unconstrained RL with the scalarized rewards R ( s , a ) − λ > C ( s , a ) , without considering the cost constraints . In order to meet the constraints , λ should also be optimized , and the procedure of ( 12 ) can be understood as joint optimization of : ν ← arg min ν L ( λ , ν ) ( OptiDICE for R− λ > C ) ( 13 ) λ← arg min λ≥0 λ > ( ĉ− E ( s , a ) ∼dD [ w∗λ , ν ( s , a ) C ( s , a ) ] ︸ ︷︷ ︸ ≈VC ( π ) ) ( Cost Lagrange multiplier ) Once the optimal solution of ( 12 ) , ( λ∗ , ν∗ ) , is computed , w∗λ∗ , ν∗ ( s , a ) = dπ ∗ ( s , a ) dD ( s , a ) is also derived by ( 11 ) , which is the stationary distribution correction between the stationary distribution of the optimal policy for the CMDP and the dataset distribution .

Summary: This paper studied the policy optimization problem in the offline constrained MDP setting. Compare with previous works, in which policy gradient based approaches are widely adopted, this paper solves the problem via policy visitation distribution that rooted from the primal-dual formulation of Bellman operator, which is novel. In order to guarantee the constraints are always satisfied, this paper provides a novel approach based on CoinDICE to efficiently estimate an upper bound of constraint violation. The author also provide sufficient empirical verifications to support their proposed algorithm.