Distributional Reinforcement Learning with Monotonic Splines

1 INTRODUCTION . A fundamental problem in traditional value-based RL is to estimate the expectation of future returns ( Mnih et al. , 2015 ; Van Hasselt et al. , 2016 ) . Distributional RL differs from this by also taking into account the intrinsic randomness of returns within MDPs ( Morimura et al. , 2010 ; Bellemare et al. , 2017 ) . To do so , distributional RL algorithms characterize the total return as a random variable and estimate its underlying distribution . In contrast , traditional value-based RL algorithms focus only on the mean of the random variable . Distributional RL offers several advantages over value-based RL that computes only expected returns . The distribution of returns enables risk-sensitive RL by facilitating the optimization of other statistics than just the mean of the returns ( Dabney et al. , 2018a ; Martin et al. , 2020 ) . Even when we stick to maximizing the mean of the returns , the distribution offers a more reliable and robust way of computing the expectation , which has led to a series of records on the Atari benchmark among value-based non-distributed RL techniques ( Bellemare et al. , 2017 ; Dabney et al. , 2018a ; b ; Hessel et al. , 2018 ; Yang et al. , 2019 ; Zhou et al. , 2020 ; Nguyen et al. , 2021 ) . Intuitively , while it is sufficient to represent an expected return by a single mean value , errors due to finite samples and function approximations can be reduced by ” canceling ” each other when multiple sample returns or quantile values are used . This is similar to the benefits of ensemble learning techniques although , strictly speaking , distributional RL is not an ensemble RL technique . In fact , distributional RL has been combined with ensemble learning and truncated critic predictions to mitigate overestimation bias in continuous control ( Kuznetsov et al. , 2020 ) . One key aspect of distributional RL algorithms is the parameterization of return distributions . In Categorical DQN ( C51 ) ( Bellemare et al. , 2017 ) , the return distributions are limited to categorical distributions over a fixed set of discrete values . It is also shown that the distributional Bellman operator is a contraction under the maximal form of the Wasserstein metric , but in practice , C51 optimizes the cross-entropy loss with a Cramér-minimizing projection ( Rowland et al. , 2018 ) . To bridge the gap between theoretical analysis and algorithmic implementation , quantile regression ( QR ) -based distributional RL algorithms ( Dabney et al. , 2018a ; b ; Yang et al. , 2019 ; Zhou et al. , 2020 ) estimate a finite number of quantile values instead of the distribution of returns since quantile regression can easily use the Wasserstein metric as the objective . In fact , the Wasserstein metric is approximately minimized by optimizing the quantile Huber loss ( Huber , 1992 ) between the Bellman updated distribution and the current return distribution . Although with an infinite number of quantiles , the step quantile function in those quantile regression based methods will approximate the full quantile function arbitrarily closely , in practice , it is infeasible to have infinite quantiles in most existing architectures . In addition , the quantile crossing issue , recently pointed out and solved by ( Zhou et al. , 2020 ) , was ignored by previous distributional RL techniques . The issue is that if no global constraint is applied , the quantile values estimated by a neural network at different quantile levels are not guaranteed to satisfy monotonicity , which can distort policy search and affect exploration during training ( Zhou et al. , 2020 ) . In this work , we propose to learn a continuous representation for quantile functions based on monotonic rational-quadratic splines ( Gregory & Delbourgo , 1982 ) . The monotonic property of these splines naturally solves the quantile crossing issue described above . Furthermore , unlike step functions or piecewise linear functions that provide a crude approximation in each bin , monotonic splines provide a more flexible and smooth approximation . With sufficiently many knots , splines can approximate any quantile function arbitrarily closely . We compare empirically our spline-based technique with other quantile-based methods in stochastic environments . We demonstrate that our method offers greater accuracy in terms of quantile approximation , faster convergence during training and higher rewards at test time . 2 DISTRIBUTIONAL REINFORCEMENT LEARNING . In standard RL settings , agent-environment interactions are modeled as a MDP , represented as a tuple ( S , A , R , P , γ ) ( Puterman , 2014 ) . S and A denote state and action spaces . P ( ·|s , a ) defines the transition . R is the state and action dependent reward , and γ ∈ ( 0 , 1 ) is a discount factor . For a policy π , the discounted sum of returns is denoted as a random variable Zπ ( s , a ) =∑∞ t=0 γ tR ( st , at ) , where s0 = s , a0 = a , st+1 ∼ P ( ·|st , at ) , and at ∼ π ( ·|st ) . The Q-value ( state-action value ) is defined as Qπ ( s , a ) = E [ Zπ ( s , a ) ] . The optimal Q-value , Q∗ ( s , a ) = maxπ Q π ( s , a ) , is the unique fixed point of the Bellman optimality operator T ( Bellman , 1966 ) Q∗ ( s , a ) = T Q∗ ( s , a ) : = E [ R ( s , a ) ] + γEP max a′ Q∗ ( s′ , a′ ) ( 1 ) In most deep RL studies , Q is approximated by a neural network . To update Q , Q-learning trains the network iteratively to minimize the squared temporal difference ( TD ) error L2t = [ rt + γmax a′ Qφ− ( s ′ , a′ ) −Qφ ( s , a ) ] 2 , ( 2 ) where φ− is the target network which is updated periodically with the most recent φ . Instead of learning the scalar Q ( s , a ) , distributional RL considers the distribution over returns ( the law of Z ) to capture the aleatoric uncertainty ( intrinsic stochasticity in the environment ) . A similar distributional Bellman operator for Z can be derived as ( Bellemare et al. , 2017 ) T πZ ( s , a ) D= R ( s , a ) + γZ ( S′ , A′ ) , ( 3 ) with S′ ∼ P ( ·|s , a ) , A′ ∼ π ( ·|S′ ) , and X D= Y indicates that random variables X and Y follow the same distribution . In theory , Bellemare et al . ( 2017 ) proved the distributional Bellman operator is a contraction in the p-Wasserstein metric Wp ( X , Y ) = ( ∫ 1 0 |F−1X ( ω ) − F −1 Y ( ω ) | pdω ) 1/p , ( 4 ) where F−1 is the quantile function ( inverse cumulative distribution function ) . Following this theory , a series of distributional RL algorithms have been proposed based on quantile regression to estimate F−1 at precisely chosen quantile fractions , such that the Wasserstein metric is minimized . 2.1 QUANTILE REGRESSION FOR DISTRIBUTIONAL RL . In QR-DQN ( Dabney et al. , 2018b ) , the random return is approximated by a uniform mixture of N Diracs Zθ ( s , a ) = 1 N N∑ i=1 δθi ( s , a ) , ( 5 ) with each θi set to a fixed quantile fraction , τ̂i = τi−1+τi 2 for 1 ≤ i ≤ N , and τi = i/N . The quantile estimation is performed by minimizing the quantile Huber loss , with threshold η 1 N N∑ i=1 N∑ j=1 ρητ̂i ( δij ) ( 6 ) on the pairwise TD error δij = r + γθj ( s′ , a′ ) − θi ( s , a ) , where ρητ ( δ ) = |τ − Iδ < 0|Lη ( δ ) , with Lη ( δ ) = { 1 2δ 2 , |δ| ≤ η η ( |δ| − 12η ) , otherwise . ( 7 ) Based on QR-DQN , Dabney et al . ( 2018a ) proposed to sample quantile fractions from a base distribution , e.g . τ ∈ U ( [ 0 , 1 ] ) rather than fixing them . They built an implicit quantile network ( IQN ) to learn mappings from sampled probability embeddings to corresponding quantile values . FQF ( Yang et al. , 2019 ) further improves IQN by learning a function to propose τ ’ s . However , the quantile values generated by neural networks may not satisfy the non-decreasing property of F−1 ( known as the quantile crossing issue ) . This was recently solved by NC-QR-DQN ( Zhou et al. , 2020 ) , by applying a softmax , followed by a cumulative sum of the output logits of the neural network Ω , and then rescaling by multiplying a non-negative factor α ( s , a ) and adding an offset β ( s , a ) : θi ( s , a ) = α ( s , a ) × ιi , a + β ( s , a ) , with ιi , a = i∑ j=0 χj , a , and χj , a = softmax ( Ω ( s ) ) j , a ( 8 ) One recent method NDQFN ( Zhou et al. , 2021 ) further combines the ideas of NC-QR-DQN and IQN to learn a monotonic function for F−1 by connecting the neighboring two monotonic quantile data points with line segments . Different from NC-QR-DQN , NDQFN generates monotonic quantile values by first learning a baseline value and then adding non-negative increments . 2.2 OTHER DISTRIBUTIONAL METHODS . Other recent methods investigate different metrics for the distributional Bellman operator . Moment matching , generally parameterized as the maximum mean discrepancy ( MMD ) between two sample sets in a reproducing Hilbert kernel space ( Gretton et al. , 2012 ) , is adopted by Nguyen et al . ( 2021 ) to propose moment matching DQN ( MM-DQN ) . The MMD loss with kernel κ is derived as : d2κ ( { υi } , { ψi } ) = 1 N2 ∑ i , j κ ( υi , υj ) + 1 M2 ∑ i , j κ ( ψi , ψj ) − 2 NM ∑ i , j κ ( υi , ψj ) , ( 9 ) where { υi } Ni=1 ∼ Z ( s , a ) and { ψi } Mi=1 ∼ R ( s , a ) + γZ ( S′ , A′ ) . It is worth noting that the theoretical analysis by Nguyen et al . ( 2021 ) shows the distributional Bellman operator under MMD is not a contraction with commonly used Gaussian kernels or expprod kernels . It is a contraction only when the kernel function is shift invariant and scale sensitive . Categorical distributional RL was also combined with policy gradient to obtain the Reactor algorithm ( Gruslys et al. , 2018 ) for discrete control and the Distributed Distributional Deep Deterministic Policy Gradient ( D4PG ) algorithm ( Barth-Maron et al. , 2018 ) for continuous control . Subsequently , Singh et al . ( 2020 ) replaced categorical return distributions by samples in Sample-based Distributional Policy Gradient ( SDPG ) , yielding improved sample efficiency . The return distribution can also be represented by a generative network trained by adversarial training ( in the same way as GANs ) to minimize temporal differences between sampled returns ( Doan et al. , 2018 ; Freirich et al. , 2019 ) . While most distributional RL techniques compute state-action return distributions , Li & Faisal ( 2021 ) proposed the Bayesian Distributional Policy Gradient ( BDPG ) algorithm that computes state return distributions and uses inference to derive a curiosity bonus . In another line of work , Tessler et al . ( 2019 ) introduced the Distributional Policy Optimization ( DPO ) framework in which an agent ’ s policy evolves towards a distribution over improving actions .

Summary: The paper studies the problem of distributional reinforcement learning. It augments the traditional quantile-based algorithms with monotonic rational-quadratic splines. Such augmentation provides a natural solution to the quantile crossing issue, which exists for many other quantile-based algorithms. Extensive experiment results are performed to verify the effectiveness of the proposed method.