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: The authors propose a DICE-family method for solving constrained offline reinforcement learning problems. To do this, they adapt ideas from OptiDICE and find a reduction from a nested constrained optimization problem to a single unconstrained optimization problem that can be efficiently represented with a neural network. Additionally, they draw on ideas from CoinDICE to estimate a confidence interval over the cost, which makes their method better at obeying constraints. The  authors compare their method to a number of baselines on both random grid worlds and continuous environments, and find that COptiDICE achieves both good performance and better constraint satisfaction than alternative methods.