Hellinger Distance Constrained Regression

1 INTRODUCTION . Policy gradient algorithms are methods of model-free reinforcement learning that optimize policy through differentiating expected discounted return . Despite the simplicity , to converge , these methods should stay on-policy because of the first-order approximation of state visitation frequencies . This issue makes agents learn through trial-and-error , using data only once . To make policy gradient updates more off-policy , we can add a constraint on the update to decrease the step size if the current policy is too far from the sampling policy . One of the first methods was to add total variation distance squared as a constraint for mixture policies . Later it was proven that there is a lower bound for new policy ’ s expected discounted return ( Kakade & Langford , 2002 ) . Recently it was proven that this lower bound exists for all types of updates ( Schulman et al. , 2015 ) . Next , total variation distance squared was replaced by Kullback-Leibler divergence that is greater than or equal to the previous one ( Pinker ’ s inequality ( Levin & Peres , 2017 ) ) , so that the lower bound was decreased ( Schulman et al. , 2015 ) . Using Lagrangian , have been derived off-policy method called Advantage-Weighted Regression ( Peng et al. , 2019 ) , which also used KL as a constraint . This article proposes a new method whose lower bound of expected discounted return is greater than or equal to the bound with KL . We achieve this by replacing total variation distance by Hellinger distance , which decreases lower bound . Therefore strictness stays the same . Then we derive an offpolicy method called Hellinger Distance Constrained Regression using the new constraint . It can be used on discrete and continuous action spaces since derivation uses Lebesgue integrals rather than a summation or Riemann integrals . 2 PRELIMINARIES . To better present the problem , we start from basic definitions , go through the history of improvements , and then describe the disadvantages of using KL divergence as a constraint . We consider an infinite-horizon discounted Markov decision process ( MDP ) , defined by the tuple ( S , A , P , r , ρ0 , γ ) , where S is a set of states ( finite or infinite ) , A is a set of actions ( finite or infinite ) , P : S × A × S → R is the transition probability distribution , r : S → R is the reward function , ρ0 : S → R is the distribution of the initial state s0 , and γ ∈ ( 0 , 1 ) is the discount factor . Let π denote a stochastic policy π : S × A → [ 0 , 1 ] , and then its expected discounted return is : η ( π ) = Es0 , a0 , ... [ ∞∑ t=0 γtr ( st ) ] , where s0 ∼ ρ0 ( · ) , at ∼ π ( ·|st ) , st+1 ∼ P ( ·|st , at ) . ( 1 ) This paper uses state-action value function Qπ , state value function Vπ , and advantage function Aπ with the following definitions : Qπ ( st , at ) = Est+1 , at+1 , ... [ ∞∑ l=0 γlr ( st+l ) ] Vπ ( st ) = Eat , st+1 , at+1 , ... [ ∞∑ l=0 γlr ( st+l ) ] Aπ ( st , at ) = Qπ ( st , at ) − Vπ ( st ) . ( 2 ) Let ρπ ( s ) be unnormalized visitation frequencies of state s where actions are chosen according to π : ρπ ( s ) = ∞∑ t=0 γtP ( st = s ) . ( 3 ) Following identity expresses the expected return of another policy π̃ in terms of the advantage over π , accumulated over states ( see Schulman et al . ( 2015 ) for proof ) : η ( π̃ ) = η ( π ) + ∫ ρπ̃ ( s ) ∫ π̃ ( a|s ) Aπ ( s , a ) da ds . ( 4 ) In approximately optimal learning , we replace state visitation frequency ρπ̃ by ρπ , since this drastically decrease optimization complexity : Lπ ( π̃ ) = η ( π ) + ∫ ρπ ( s ) ∫ π̃ ( a|s ) Aπ ( s , a ) da ds . ( 5 ) Let πold denote current policy , then the lower bound for the expected discounted return for the new policy πnew will be ( see Schulman et al . ( 2015 ) for proof ) : η ( πnew ) ≥ Lπold ( πnew ) − 4 γ ( 1− γ ) 2 α2 where = max s , a |Aπ ( s , a ) | , α = max s DTV ( πold ( ·|s ) ||πnew ( ·|s ) ) , DTV ( πold ( ·|s ) ||πnew ( ·|s ) ) = 1 2 ∫ |πold ( a|s ) − πnew ( a|s ) | da . ( 6 ) Theoretical Trust-Region Policy Optimization algorithm relays on Pinsker ’ s inequality ( see ( Tsybakov , 2009 ) for proof ) : DKL ( πold ( ·|s ) ||πnew ( ·|s ) ) ≥ DTV ( πold ( ·|s ) ||πnew ( ·|s ) ) 2 where DKL ( πold ( ·|s ) ||πnew ( ·|s ) ) = ∫ πold ( a|s ) log πold ( a|s ) πnew ( a|s ) da . ( 7 ) To retain strictness and decrease calculation complexity , total variation distance squared was replaced with Kullback-Leibler divergence DKL ( πold ( ·|s ) ||πnew ( ·|s ) ) : η ( πnew ) ≥ Lπold ( πnew ) − CDmaxKL ( πold||πnew ) where = max s , a |Aπ ( s , a ) | , C = 4 γ ( 1− γ ) 2 , DmaxKL ( πold||πnew ) = max s DKL ( πold ( ·|s ) ||πnew ( ·|s ) ) . ( 8 ) However , this replacement greatly decreases the lower bound for the expected discounted return for the new policy . Moreover , Kullback-Leibler divergence has no upper bound . Therefore we have no policy-independent lower bound for this type of update . 3 HELLINGER DISTANCE IN POLICY OPTIMIZATION . We can improve lower bound ( compared to KL ) by replacing DTV ( πold ( · | s ) || πnew ( · | s ) ) with Hellinger distance H ( πold ( ·|s ) || πnew ( · | s ) ) : H ( πold ( ·|s ) || πnew ( ·|s ) ) 2 = 1− ∫ √ πold ( a|s ) πnew ( a|s ) da ( 9 ) Theorem 1 ( see Appendix A or ( Tsybakov , 2009 , section 2.4 ) for proof ) proves correctness and improvement ( compared to KL ) to the lower bound . Let p ( v ) and q ( v ) be two probability density functions then : DTV ( p ( · ) || q ( · ) ) 2 ≤ 2H ( p ( · ) || q ( · ) ) 2 ≤ DKL ( p ( · ) || q ( · ) ) ( 10 ) Replacing p ( v ) and q ( v ) with πold ( · | s ) and πnew ( · | s ) respectively , new lower bound follows : η ( πnew ) ≥ Lπold ( πnew ) − 8 γ ( 1− γ ) 2 α2 where = max s , a |Aπ ( s , a ) | , α = max s H ( πold ( ·|s ) ||πnew ( ·|s ) ) ( 11 ) It is worth to note that H ( πold ( · | s ) || πnew ( · | s ) ) ≤ 1 . 4 HELLINGER DISTANCE CONSTRAINED REGRESSION ( HDCR ) . We could use presented lower bound as in TRPO , but instead , we derive an offline regression algorithm by introducing the following optimization problem where µ is the sampling policy : arg max π ∫ ρµ ( s ) ∫ π ( a|s ) Aµ ( s , a ) da ds s.t . ∫ ρµ ( s ) H ( π ( ·|s ) ||µ ( ·|s ) ) ds ≤ , ∫ π ( a|s ) da = 1 , ∀s ∈ S. ( 12 ) Constructing Lagrangian , differentiating it with respect to π , and solving for π gives us the following optimal policy ( see Appendix B for derivation ) : π∗ ( a|s ) = µ ( a|s ) β 2 ( β − 2Aµ ( s , a ) ) 2 , where β is a Lagrangian multiplier . ( 13 ) Constructing a regression problem of KL divergence between optimal policy π∗ and current policy π and simplifying gives us the following supervised regression problem ( see Appendix B for derivation ) : arg max π Es∼ρµ ( · ) Ea∼µ ( ·|s ) log π ( a|s ) 1 ( β −Aµ ( s , a ) ) 2 ( 14 ) Using notation of ABM paper ( Siegel et al. , 2020 ) ” advantage-weighting ” function is f ( A ( s , a ) ) = 1 ( β−A ( s , a ) ) 2 . If we use HDCR with Experience Replay , in equation 14 , we replace µ ( ·|s ) in expectation Ea ∼ µ ( · | s ) and advantage function Aµ ( s , a ) . Let Π = { πi , πi+1 , ... , πi+N } be a set of sampling policies from which actions were sampled , w ( πi ) probability of selecting policy πi , then : µ ( s , a ) = ∫ Π w ( π ) ρπ ( s ) π ( a|s ) dπ Aµ ( s , a ) = ∫ Π w ( π ) ρπ ( s ) ( Qπ ( s , a ) − Vπ ( s ) ) dπ∫ Π w ( π ) ρπ ( s ) dπ V ( s ) = ∫ Π w ( π ) ρπ ( s ) Vπ ( s ) dπ∫ Π w ( π ) ρπ ( s ) dπ ( 15 ) The proof will repeat the proof for AWR with Experience Replay ( Peng et al. , 2019 ) . Practically we simply sample uniformly from the replay buffer and using a value function estimator . This type of sampling provides us an approximation of expectation and state value function . Let D denote a set of stored trajectories ( replay buffer ) , A ( s , a ; φ ) is an advantage function parameterized by vector φ . The most popular method to obtain A ( s , a ; φ ) for offline learning is to use state-action function estimator Q ( s , a ; φ ) parameterized by φ : A ( s , a ; φ ) = Q ( s , a ; φ ) − ∫ π ( a′|s ; θk ) Q ( s , a′ ; φk ) da′ ( 16 ) Despite Q ( s , a ; φ ) being closer to the expectation of discounted return following policy π ( because of ” taking ” the first action according to π rather then µ ) , we found Monte-Carlo return more efficient on tiny offline datasets . Greater performance using MC return can be explained by a lack of experience ” produced ” by certain actions . Monte-Carlo estimation of A ( s , a ; φ ) can be described as follows where RDst , at = ∑T l=0 γ lrt+l and V ( s ; φ ) is a state value function estimator parameterized by vector φ : A ( s , a ; φ ) = RDs , a − V ( s ; φ ) ( 17 ) Also , we can use Generalized Advantage Estimation ( Schulman et al. , 2016 ) , where γ ∈ [ 0 , 1 ] and δφkt is a one-step temporal difference for state st calculated using old vector φk : δφkt = rt + γV ( st+1 ; φk ) − V ( st ; φk ) Â ( st , at ; φk ) = T∑ l=0 ( γλ ) lδφkt+l A ( st , at ; φ ) = Â ( st , at ; φk ) + V ( st ; φk ) − V ( st ; φ ) ( 18 ) Finally , we propose the following reinforcement learning algorithm : Algorithm 1 : Hellinger Distance Constrained Regression θ1 ← random initial weights D ← ∅ for iteration k = 1 , ... , kmax do add trajectories { τi } sampled via πθk to D φk+1 ← arg minφ Es , a∼DA2 ( s , a ; φ ) θk+1 ← arg maxθ Es , a∼D [ log πθ ( a|s ) 1 ( β−A ( s , a ; φk ) ) 2 ] end

Summary: The paper provides a new metric - Hellinger distance to be combined with trust region ideas in policy optimization. The major difference from prior work is the change of this distance metric. The paper shows that with this distance metric, along with Lagrangian relaxation, one could show analytic results of improved policies. The paper also shows similar lower bound improvement results and compared with baselines on offline rl tasks.