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: This paper proposes a new neural network design to represent quantile functions for distributional reinforcement learning, based on smooth rational-quadratic splines. This representation has the advantage of being continuously differentiable. The loss is computed by evaluating the quantile loss on a set of uniformly spaced quantile levels. The benefits are shown empirically on standard continuous control environments, modified with noise to make them stochastic, by using the proposed distributional representation as critic in DDPG/SAC algorithms.