Offline Policy Optimization with Variance Regularization

1 INTRODUCTION . Offline batch reinforcement learning ( RL ) algoithms are key towards scaling up RL for real world applications , such as robotics ( Levine et al. , 2016 ) and medical problems . This is because offline RL provides the appealing ability for agents to learn from fixed datasets , similar to supervised learning , avoiding continual interaction with the environment , which could be problematic for safety and feasibility reasons . However , significant mismatch between the fixed collected data and the policy that the agent is considering can lead to high variance of value function estimates , a problem encountered by most off-policy RL algorithms ( Precup et al. , 2000 ) . A complementary problem is that the value function can become overly optimistic in areas of state space that are outside the visited batch , leading the agent in data regions where its behavior is poor Fujimoto et al . ( 2019 ) . Recently there has been some progress in offline RL ( Kumar et al. , 2019 ; Wu et al. , 2019b ; Fujimoto et al. , 2019 ) , trying to tackle both of these problems . In this work , we study the problem of offline policy optimization with variance minimization . To avoid overly optimistic value function estimates , we propose to learn value functions under variance constraints , leading to a pessimistic estimation , which can significantly help offline RL algorithms , especially under large distribution mismatch . We propose a framework for variance minimization in offline RL , such that the obtained estimates can be used to regularize the value function and enable more stable learning under different off-policy distributions . We develop a novel approach for variance regularized offline actor-critic algorithms , which we call Offline Variance Regularizer ( OVR ) . The key idea of OVR is to constrain the policy improvement step via variance regularized value function estimates . Our algorithmic framework avoids the double sampling issue that arises when computing gradients of variance estimates , by instead considering the variance of stationary distribution corrections with per-step rewards , and using the Fenchel transformation ( Boyd & Vandenberghe , 2004 ) to formulate a minimax optimization objective . This allows minimizing variance constraints by instead optimizing dual variables , resulting in simply an augmented reward objective for variance regularized value functions . We show that even with variance constraints , we can ensure policy improvement guarantees , where the regularized value function leads to a lower bound on the true value function , which mitigates the usual overestimation problems in batch RL The use of Fenchel duality in computing the variance allows us to avoid double sampling , which has been a major bottleneck in scaling up variance-constrained actor-critic algorithms in prior work A . & Ghavamzadeh ( 2016 ) ; A . & Fu ( 2018 ) . Practically , our algorithm is easy to implement , since it simply involves augmenting the rewards with the dual variables only , such that the regularized value function can be implemented on top of any existing offline policy optimization algorithms . We evaluate our algorithm on existing offline benchmark tasks based on continuous control domains . Our empirical results demonstrate that the proposed variance regularization approach is particularly useful when the batch dataset is gathered at random , or when it is very different from the data distributions encountered during training . 2 PRELIMINARIES AND BACKGROUND . We consider an infinite horizon MDP as ( S , A , P , γ ) where S is the set of states , A is the set of actions , P is the transition dynamics and γ is the discount factor . The goal of reinforcement learning is to maximize the expected return J ( π ) = Es∼dβ [ V π ( s ) ] , where V π ( s ) is the value function V π ( s ) = E [ ∑∞ t=0 γ tr ( st , at ) | s0 = s ] , and β is the initial state distribution . Considering parameterized policies πθ ( a|s ) , the goal is maximize the returns by following the policy gradient ( Sutton et al. , 1999 ) , based on the performance metric defined as : J ( πθ ) = Es0∼ρ , a0∼π ( s0 ) [ Qπθ ( s0 , a0 ) ] = E ( s , a ) ∼dπθ ( s , a ) [ r ( s , a ) ] ( 1 ) where Qπ ( s , a ) is the state-action value function , since V π ( s ) = ∑ a π ( a|s ) Qπ ( s , a ) . The policy optimization objective can be equivalently written in terms of the normalized discounted occupancy measure under the current policy πθ , where dπ ( s , a ) is the state-action occupancy measure , such that the normalized state-action visitation distribution under policy π is defined as : dπ ( s , a ) = ( 1 − γ ) ∑∞ t=0 γ tP ( st = s , at = a|s0 ∼ β , a ∼ π ( s0 ) ) . The equality in equation 1 holds and can be equivalently written based on the linear programming ( LP ) formulation in RL ( see ( Puterman , 1994 ; Nachum & Dai , 2020 ) for more details ) . In this work , we consider the off-policy learning problem under a fixed dataset D which contains s , a , r , s′ tuples under a known behaviour policy µ ( a|s ) . Under the off-policy setting , importance sampling ( Precup et al. , 2000 ) is often used to reweight the trajectory under the behaviour data collecting policy , such as to get unbiased estimates of the expected returns . At each time step , the importance sampling correction π ( at|st ) µ ( at|st ) is used to compute the expected return under the entire trajectory as J ( π ) = ( 1 − γ ) E ( s , a ) ∼dµ ( s , a ) [ ∑T t=0 γ tr ( st , at ) ( ∏T t=1 π ( at|st ) µ ( at|st ) ] . Recent works ( Fujimoto et al. , 2019 ) have demonstrated that instead of importance sampling corrections , maximizing value functions directly for deterministic or reparameterized policy gradients ( Lillicrap et al. , 2016 ; Fujimoto et al. , 2018 ) allows learning under fixed datasets , by addressing the over-estimation problem , by maximizing the objectives of the form maxθ Es∼D [ Qπθ ( s , πθ ( s ) ] . 3 VARIANCE REGULARIZATION VIA DUALITY IN OFFLINE POLICY OPTIMIZATION . In this section , we first present our approach based on variance of stationary distribution corrections , compared to importance re-weighting of episodic returns in section 3.1 . We then present a derivation of our approach based on Fenchel duality on the variance , to avoid the double sampling issue , leading to a variance regularized offline optimization objective in section 3.2 . Finally , we present our algorithm in 1 , where the proposed regularizer can be used in any existing offline RL algorithm . 3.1 VARIANCE OF REWARDS WITH STATIONARY DISTRIBUTION CORRECTIONS . In this work , we consider the variance of rewards under occupancy measures in offline policy optimization . Let us denote the returns as Dπ = ∑T t=0 γ tr ( st , at ) , such that the value function is V π = Eπ [ Dπ ] . The 1-step importance sampling ratio is ρt = π ( at|st ) µ ( at|st ) , and the T-steps ratio can be denoted ρ1 : T = ∏T t=1 ρt . Considering per-decision importance sampling ( PDIS ) ( Precup et al. , 2000 ) , the returns can be similarly written as Dπ = ∑T t=0 γ trtρ0 : t. The variance of episodic returns , which we denote by VP ( π ) , with off-policy importance sampling corrections can be written as : VP ( π ) = Es∼β , a∼µ ( ·|s ) , s′∼P ( ·|s , a ) [ ( Dπ ( s , a ) − J ( π ) ) 2 ] . Instead of importance sampling , several recent works have instead proposed for marginalized importance sampling with stationary state-action distribution corrections ( Liu et al. , 2018 ; Nachum et al. , 2019a ; Zhang et al. , 2020 ; Uehara & Jiang , 2019 ) , which can lead to lower variance estimators at the cost of introducing bias . Denoting the stationary distribution ratios as ω ( s , a ) = dπ ( s , a ) dµ ( s , a ) , the returns can be written as Wπ ( s , a ) = ω ( s , a ) r ( s , a ) . The variance of marginalized IS is : VD ( π ) = E ( s , a ) ∼dµ ( s , a ) [ ( Wπ ( s , a ) − J ( π ) ) 2 ] = E ( s , a ) ∼dµ ( s , a ) [ Wπ ( s , a ) 2 ] − E ( s , a ) ∼dµ ( s , a ) [ Wπ ( s , a ) ] 2 ( 2 ) Our key contribution is to first consider the variance of marginalized IS VD ( π ) itself a as risk constraints , in the offline batch optimization setting . We show that constraining the offline policy optimization objective with variance of marginalized IS , and using the Fenchel-Legendre transformation on VD ( π ) can help avoid the well-known double sampling issue in variance risk constrained RL ( for more details on how to compute the gradient of the variance term , see appendix B ) . We emphasize that the variance here is solely based on returns with occupancy measures , and we do not consider the variance due to the inherent stochasticity of the MDP dynamics . 3.2 VARIANCE REGULARIZED OFFLINE MAX-RETURN OBJECTIVE . We consider the variance regularized off-policy max return objective with stationary distribution corrections ωπ/D ( which we denote ω for short for clarity ) in the offline fixed dataset D setting : max πθ J ( πθ ) : = Es∼D [ Qπθ ( s , πθ ( s ) ) ] − λVD ( ω , πθ ) ( 3 ) where λ ≥ 0 allows for the trade-off between offline policy optimization and variance regularization ( or equivalently variance risk minimization ) . The max-return objective under Qπθ ( s , a ) has been considered in prior works in offline policy optimization ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ) . We show that this form of regularizer encourages variance minimization in offline policy optimization , especially when there is a large data distribution mismatch between the fixed dataset D and induced data distribution under policy πθ . 3.3 VARIANCE REGULARIZATION VIA FENCHEL DUALITY . At first , equation 3 seems to be difficult to optimize , especially for minimizing the variance regularization w.r.t θ . This is because finding the gradient of V ( ω , πθ ) would lead to the double sampling issue since it contains the squared of the expectation term . The key contribution of OVR is to use the Fenchel duality trick on the second term of the variance expression in equation 2 , for regularizing policy optimization objective with variance of marginalized importance sampling . Applying Fenchel duality , x2 = maxy ( 2xy − y2 ) , to the second term of variance expression , we can transform the variance minimization problem into an equivalent maximization problem , by introducing the dual variables ν ( s , a ) . We have the Fenchel conjugate of the variance term as : V ( ω , πθ ) = max ν { − 1 2 ν ( s , a ) 2 + ν ( s , a ) ω ( s , a ) r ( s , a ) + E ( s , a ) ∼dD [ ω ( s , a ) r ( s , a ) 2 ] } = max ν E ( s , a ) ∼dD [ − 1 2 ν ( s , a ) 2 + ν ( s , a ) ω ( s , a ) r ( s , a ) + ω ( s , a ) r ( s , a ) 2 ] ( 4 ) Regularizing the policy optimization objective with variance under the Fenchel transformation , we therefore have the overall max-min optimization objective , explicitly written as : max θ min ν J ( πθ , ν ) : = Es∼D [ Qπθ ( s , πθ ( s ) ) ] −λE ( s , a ) ∼dD [ ( − 1 2 ν2+ν ·ω ·r+ω ·r2 ) ( s , a ) ] ( 5 )

Summary: This paper proposes a novel algorithm for offline policy optimization. The main idea is to prevent overestimation bias by regularizing against the variance of the importance weighted value estimate. There are two key modifications: (1) using an importance weight from the stationary distribution and (2) using Fenchel duality to introduce a min-max problem to avoid double sampling when estimating the gradient of the variance regularization term. The theory section motivates the use of variance regularization and the experiments show improvements over BCQ when adding the proposed variance regularization algorithm.