Towards Understanding Distributional Reinforcement Learning: Regularization, Optimization, Acceleration and Sinkhorn Algorithm

1 INTRODUCTION . The intrinsic characteristics of classical reinforcement learning ( RL ) algorithms , such as temporaldifference ( TD ) learning ( Sutton & Barto , 2018 ) and Q-learning ( Watkins & Dayan , 1992 ) , are based on the expectation of discounted cumulative rewards that an agent observes while interacting with the environment . In stark contrast to the classical expectation-based RL , a new branch of algorithms called distributional RL estimates the full distribution of total returns and has demonstrated stateof-the-art performance in a wide range of environments ( Bellemare et al. , 2017a ; Dabney et al. , 2018b ; b ; Yang et al. , 2019 ; Zhou et al. , 2020 ; Nguyen et al. , 2020 ) . Meanwhile , distributional RL has also enjoyed further benefits in risk-sensitive control , policy exploration settings ( Mavrin et al. , 2019 ; Rowland et al. , 2019 ) and robsutness ( Sun et al. , 2021 ) . Despite the existence of numerous algorithmic variants of distributional RL with remarkable empirical success , theoretical studies of advantages of distributional RL over expectation-based RL are less established . Existing works include ( Lyle et al. , 2019 ) that proved in many tabular and linear approximation settings , distributional RL behaves exactly the same as expectation-based RL . Lyle et al . ( 2021 ) investigated the impact of distributional RL from the perspective of representation dynamics . Martin et al . ( 2020 ) recently mapped distributional RL problems to a Wasserstein gradient flow problem , treating the distributional Bellman residual as a potential energy functional . Offline distributional RL ( Ma et al. , 2021 ) has also been proposed to investigate the efficacy of distributional RL in both risk-neutral and risk-averse domains . Recent works have tended towards closing the gap between theory and practice in distributional RL . From an algorithmic perspective , the Sinkhorn loss ( Sinkhorn , 1967 ) can be used to tractably approximate the Wasserstein distance and has been successfully applied in numerous crucial machine learning developments , including the Sinkhorn-GAN ( Genevay et al. , 2018 ) and Sinkhorn-based adversarial training ( Wong et al. , 2019 ) . Inspired by the distributional RL literature , Martin et al . ( 2020 ) argues for the use of a second-order stochastic dominance relation to select among a multiplicity of competing solutions via Sinkhorn iteration ( Sinkhorn , 1967 ) , which can be useful to manage stochastic uncertainty in RL paradigms . However , a Sinkhorn-based distributional RL algorithm has not yet to be formally proposed and investigated . In this paper , we theoretically illuminate the superiority of distributional RL over expectation-based RL from the perspectives of regularization , optimization , acceleration , and representation . Specifically , we simplify distributional RL as a neural Z-fitted iteration , within which we establish an equivalence between distributional RL and a form of entropy-regularized maximum likelihood estimation ( MLE ) . We also demonstrate that the resulting novel cross entropy regularization correlates strongly with the behavior of maximum entropy RL . By incorporating a histogram distributional loss , we further achieve stable optimization and guaranteed generalization of distributional RL , attributable to desirable smoothness properties of the distribution loss . We further characterize the effect of acceleration on distributional RL and discuss when distributional RL algorithms are effective in various environments . After gaining insights from the theoretical advantages of distributional RL , we further propose a novel distributional RL algorithm based on Sinkhorn loss that interpolates between the Wasserstein distance and maximum mean discrepancy ( MMD ) . Our approach allows us to find a trade-off that simultaneously leverages the geometry of the Wasserstein distance and the favorable high-dimensional sample complexity and unbiased gradient estimates of MMD . In summary , our analysis opens the door to a deeper understanding of theoretical advantage of distributional RL . 2 PRELIMINARY KNOWLEDGE . In the standard RL setting , an agent interacts with an environment via a Markov decision process ( MDP ) , a 5-tuple ( S , A , R , P , γ ) , where S and A are the state and action spaces , respectively . P is the environment transition dynamics , R is the reward function and γ ∈ ( 0 , 1 ) is the discount factor . State-value function vs. state-value distribution . Given a policy π , the discounted sum of future rewards is a random variable Zπ ( s , a ) = ∑∞ t=0 γ tR ( st , at ) , where s0 = s , a0 = a , st+1 ∼ P ( ·|st , at ) , and at ∼ π ( ·|st ) . In the control setting , expectation-based RL focuses on the state-value function Qπ ( s , a ) , which is the expectation of Zπ ( s , a ) , i.e. , Qπ ( s , a ) = E [ Zπ ( s , a ) ] . Distributional RL , on the other hand , focuses on the state-value distribution , the full distribution of Zπ ( s , a ) . Leveraging knowledge on the entire distribution can better capture the uncertainty of returns in the MDP beyond the expectation of return ( Dabney et al. , 2018a ; Mavrin et al. , 2019 ) . Bellman operators vs. distributional Bellman operators . For the policy evaluation in expectation-based RL , the value function is updated via the Bellman operator T πQ ( s , a ) = E [ R ( s , a ) ] + γEs′∼p , π [ Q ( s′ , a′ ) ] . In distributional RL , the state-value distribution of Zπ ( s , a ) is updated via the distributional Bellman operator Tπ TπZ ( s , a ) = R ( s , a ) + γZ ( s′ , a′ ) , ( 1 ) where s′ ∼ P ( ·|s , a ) and a′ ∼ π ( ·|s′ ) . From a theoretical perspective , both the Bellman operator T π in the policy evaluation setting and the Bellman optimality operator T in the control setting are contractive in the stationary policy case . In contrast , the distributional Bellman operator Tπ is contractive under certain distribution divergence metrics , but the distributional Bellman optimality operator T can only converge to a set of optimal non-stationary value distributions in a weak sense ( Elie & Arthur , 2020 ) . 3 EFFECT OF REGULARIZATION ON DISTRIBUTIONAL RL . Although the theoretical framework of distributional RL in the tabular setting has been basically established mentioned in Section 2 , the theoretical understanding of its advantages over expectationbased RL has been less studied . In this section , we attribute the superiority of distributional RL into its regularization effect . 3.1 DISTRIBUTIONAL RL : NEURAL Z-FITTED ITERATION . Neural Q-Fitted Iteration . In the function approximation setting , Deep Q Learning ( Mnih et al. , 2015 ) can be simplified into Neural Q-Fitted Iteration ( Fan et al. , 2020 ) under tricks of experience replay and the target network Qθ∗ , where we update parameterized Qθ ( s , a ) in each iteration k : Qk+1θ = argmin Qθ 1 n n∑ i=1 [ yi −Qkθ ( si , ai ) ] 2 , ( 2 ) where the target yi = r ( si , ai ) + γmaxa∈AQkθ∗ ( s ′ i , a ) is fixed within every Ttarget steps to update target network Qθ∗ by letting θ∗ = θ and the experience buffer induces independent samples { ( si , ai , ri , s′i ) } i∈ [ n ] . In an ideal case that neglects the non-convexity and TD approximation errors , we have Qk+1θ = T Qkθ , which is exactly the updating under Bellman optimality operator . Under the two conditions , the optimization problem in Eq . 2 can be viewed as Least Square Estimation ( LSE ) in a neural network parametric regression problem between the updating of target network Qθ∗ . Neural Z-Fitted Iteration . Analogous to neural Q-fitted iteration , we can also simplify value-based distributional RL methods based on a parameterized Zθ into a Neural Z-fitted Iteration as Zk+1θ = argmin Zθ 1 n n∑ i=1 dp ( Yi , Z k θ ( si , ai ) ) , ( 3 ) where the target Yi = R ( si , ai ) + γZkθ∗ ( s ′ i , πZ ( s ′ ) ) with πZ ( s′ ) = argmaxa′ E [ Zkθ∗ ( s ′ , a′ ) ] is fixed within every Ttarget steps to update target network Zθ∗ , and dp is a divergence metric between two distributions . Notably , the options of representation manner on Zθ and the metric dp are pivotal for the empirical success of distributional RL algorithms . For instance , QR-DQN ( Dabney et al. , 2018b ) approximates Wasserstein distance Wp , which leverages quantiles to represent the distribution of Zθ . C51 ( Bellemare et al. , 2017a ) represents Zθ via a categorical distribution under the convergence of Cramér distance ( Bellemare et al. , 2017b ; Rowland et al. , 2018 ) , a special case with p = 2 of the ` p distance ( Elie & Arthur , 2020 ) , while Moment Matching ( Nguyen et al. , 2020 ) learns deterministic samples to represent the distribution of Zθ based on Maximum Mean Discrepancy ( MMD ) . Contractive properties under typical metrics dp can be summarized as follows • T π is γ-contractive under the supreme form of Wassertein distance Wp . • T π is γ1/p-contractive under the supreme form of ` p distance . • T π is γα/2-contractive under MMD∞ with the kernel kα ( x , y ) = −‖x− y‖α , ∀α ∈ R. For the completeness , the definition of mentioned distances and the proof of contraction are provided in Appendix A . Although the widely used Kullback–Leibler ( KL ) divergence is not a contraction ( Morimura et al. , 2012 ) , we show in Proposition 1 that the KL divergence still enjoys desirable properties in distributional RL context , which can be reasonable for the theoretical analysis . We assume Zθ is absolutely continuous and has joint supports , under which the KL divergence is welldefined . Proof of Proposition 1 and the definition of supreme DKL are provided in Appendix B . Proposition 1 . Denote the supreme of DKL as D∞KL , we have : ( 1 ) Tπ is a non-expansive operator underD∞KL , i.e. , D ∞ KL ( T πZ1 , T πZ2 ) ≤ D∞KL ( Z1 , Z2 ) , ( 2 ) D∞KL ( Zn , Z ) → 0 impliesWp ( Zn , Z ) → 0 , ( 3 ) the expectation of Zπ is still γ-contractive , i.e. , ‖ETπZ1 − ETπZ2‖∞ ≤ γ ‖EZ1 − EZ2‖∞ . 3.2 DISTRIBUTIONAL RL : A NOVEL ENTROPY-REGULARIZED MLE . The reasonable properties of KL divergence in Proposition 1 allows us to leverage it to conduct the theoretical analysis . To separate the impact of additional distribution information from the expectation of Zπ , we leverage the variant technique of gross error model from robust statistics ( Huber , 2004 ) , similar to the technique to analyze Label Smoothing ( Müller et al. , 2019 ) and Knowledge Distillation ( Hinton et al. , 2015 ) . Specifically , we denote the one-dimensional full distribution of Zπ as F , and the distribution on the remaining support getting rid of E [ Zπ ] as Fµ . Hence , we can obtain the distribution decomposition for Zπ ( s , a ) as F s , a ( x ) = ( 1− ) 1 { x≥E [ Zπ ( s , a ) ] } ( x ) + F s , aµ ( x ) , ( 4 ) where controls the proportion of F s , aµ ( x ) and the indicator function 1 { x≥E [ Zπ ( s , a ) ] } = 1 if x ≥ E [ Zπ ( s , a ) ] , otherwise 0 . After taking derivatives on both sides , we attain the relationship of their density functions as ps , a ( x ) = ( 1 − ) δ { x=E [ Zπ ( s , a ) ] } ( x ) + µs , a ( x ) , where µs , a ( x ) is the density function related to Zπ ( s , a ) on remaining supports removing E [ Zπ ( s , a ) ] . It is worth noting that the existence of µs , a ( x ) can be simply guaranteed by directly computing µs , a ( x ) = ps , a ( x ) / − ( 1 − ) δ { x=E [ Zπ ( s , a ) ] } / as long as ps , a ( x ) and the expectation of Zπ ( s , a ) exist . Next , we use ps , a ( x ) and qs , aθ ( x ) to denote the density distributions behind { Yi } i∈ [ n ] and Zkθ ( s , a ) in neural Zfitted iteration via Eq . 3 , respectively . Therefore , we can derive the following result in Proposition 2 . Proposition 2 . Let H ( P , Q ) as the cross entropy , i.e. , H ( P , Q ) = − ∫ x∈X P ( x ) logQ ( x ) dx . Let α be a positive constant , and based on the decomposition in Eq . 4 and DKL as dp , Neural Z-fitted iteration in Eq . 3 can be reformulated as Zk+1θ = argmin Zθ 1 n n∑ i=1 H ( δ { x=E [ Zπ ( si , ai ) ] } , q si , ai θ ) + αH ( µ si , ai , qsi , aiθ ) . ( 5 ) We provide the proof in Appendix C. For the uniformity of notation , we still use s , a in the following analysis instead of si , ai in Eq . 5 . Importantly , the first term in Eq . 5 can be further simplified as − ∫ x∈X log q s , a θ ( E [ Z ( s , a ) ] ) . Minimizing this first term can be viewed as a variant of Maximum Likelihood Estimation ( MLE ) on the expectation E [ Z ( s , a ) ] rather than the traditional MLE directly on observed samples . The cross entropy regularization in the second term pushes qs , aθ to approximate the distribution µs , a in order to fully utilize the additional distributional information while learning , serving as the key to the superiority of distributional RL . This novel cross entropy regularization regarding µs , a and qs , aθ is different from the classical entropy regularization used in RL , which we further analyze their connection and discrepancy in Section 3.3 . In summary , distributional RL can be simplified as a novel entropy-regularized MLE within neural Z-fitted iteration framework in stark contrast to the Least-Square estimation of expectation-based RL in the neural Q-fitted iteration .

Summary: Distributional reinforcement learning (DRL) is a family of RL algorithms that estimates distributions of the value function rather than the expectation. The first part of the paper interprets distributional RL as adding a cross-entropy regularizer to the traditional objective. The second part discusses several advantages of DRL over expectation-based RL, including more stable optimization, better sample efficiency, and better state representations. The last part proposes a new distributional RL algorithm based on Sinkhorn loss, aiming to leverage the nice geometric interpretation of Wasserstein distance and sample complexity of Maximum Mean Discrepancy.