Year: 2021

Abstract:
Distributional Reinforcement Learning (RL) differs from traditional RL by estimating the distribution over returns to capture the intrinsic uncertainty of MDPs. One key challenge in distributional RL lies in how to parameterize the quantile function when minimizing the Wasserstein metric of temporal differences. Existing algorithms use a step functions or piecewise linear functions. In this paper, we propose to learn smooth continuous quantile functions represented by monotonic rational-quadratic splines, which also naturally solve the quantile crossing problem. Experiments in stochastic environments show that a dense estimation for quantile functions enhances distributional RL in terms of faster empirical convergence and higher rewards in most cases.

heading: 1 INTRODUCTION
text: 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.

heading: 2 DISTRIBUTIONAL REINFORCEMENT LEARNING
text: 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.

heading: 2.1 QUANTILE REGRESSION FOR DISTRIBUTIONAL RL
text: 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.

heading: 2.2 OTHER DISTRIBUTIONAL METHODS
text: 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.

heading: 3 OUR ALGORITHM
text: Our method is originally motivated by NC-QR-DQN, where a special architecture is designed for the last layer of the neural network to satisfy the monotonicity of F−1. The output represents the estimated quantile values at chosen quantile fractions. One drawback of discretization is that a precise approximation for F−1 may need infinite fractions. But in practice, one can only use finite
quantile fractions to estimate quantile values for decision making. In this work, we propose to learn a dense approximation for F−1 using monotonic rational-quadratic splines (Gregory & Delbourgo, 1982) as a building block.

heading: 3.1 QUANTILE APPROXIMATION WITH MONOTONIC RATIONAL QUADRATIC SPLINES
text: Monotonic splines produce a monotonic interpolant to a set of monotonic data points (called knots) {(xk, yk)}Kk=0, which has been recently used as a transformation function in normalizing flows (Durkan et al., 2019; Dolatabadi et al., 2020). Furthermore, denote {dk}Kk=0, a set of positive numbers, as the derivative of each knot. The monotonic rational-quadratic splines aim to find rational-quadratic functions with form fk(x) =
Ok(x) Pk(x) to fit the points and derivatives in each interval (called bin) [xk, xk+1], where Ok and Pk are quadratic functions (with form ax2 + bx+ c).
Gregory & Delbourgo (1982) suggested to construct Ok and Pk as follows. Denote gk = (yk+1 − yk)/(xk+1 − xk) and hk(x) = (x− xk)/(xk+1 − xk) for x ∈ [xk, xk+1]. The expressions for the quadratic Ok(hk(x)) and Pk(hk(x)) for the kth bin is defined by (use hk for short of hk(x)):
Ok(hk) = gkyk+1h 2 k + (ykdk+1 + yk+1)hk(1− hk) + gkyk(1− hk)2 Pk(hk) = gk + (dk+1 + dk − 2gk)hk(1− hk) (10)
Then, the rational-quadratic function for the kth bin is computed by the quotient of Ok and Pk
fk(hk) = Ok(hk)
Pk(hk) = yk + (yk+1 − yk)[gkh2k + dkhk(1− hk)] gk + (dk+1 + dk − 2gk)hk(1− hk) . (11)
Equation 11 is proven to be monotonic and continuously differentiable, while passing through the knots and satisfying the given derivatives at the knots (Gregory & Delbourgo, 1982).
The monotonicity of the above splines fits the non-decreasing property of F−1. Let F−1Z(s,a)(τ) be the quantile function for the random variable of the discounted total return Z(s, a) with τ ∈ [0, 1]. Given the number of bins K, the aims of the spline approximator for F−1Z(s,a) are threefold. First, propose a partition for the domain of definition [0, 1] with τ0 < ... < τk < ... < τK . Here τ0 = 0 and τK = 1. Second, estimate the corresponding quantile values q0 < ... < qk < ... < qK . Third, assess the derivatives at those points with d0, ..., dk, ...dK . We give a small fixed positive value for d0 and dK as they are assigned with endpoints. After the generation of these three sets of statistics, the monotonic spline of each bin is given by Equation 11 (by replacing xk by τk and yk by qk).

heading: 3.2 MODEL IMPLEMENTATION
text: We now show how to learn the monotonic splines for quantile functions in distributional RL by neural networks, and we name the technique spline DQN (SPL-DQN). As shown in Figure 1, the SPL-DQN consists of three major components, including a Feature Extractor which extracts latent features from a state, a Knots Logit Network which, for each action, generates the logits of the widths and heights for K bins, and derivatives for K − 1 inner knots, and a Bin Scale Network which recovers the heights in [0, 1] to the original quantile range. Here we describe the model for a
discrete action space of size |A|. To use monotonic splines in continuous control, the model can be modified by taking state-action pairs as input and only producing knots for that state-action pair.
The Feature Extractor F is usually made up of multiple convolutional layers with subsequent fullyconnected layers for image-like inputs or stacked fully-connected layers for non-image inputs. It produces the feature embedding F(s) ∈ Rd of state s. Then the Knots Logit NetworkW maps F(s) to unconstrained logits v with dimension |A| × (3K − 1) using a fully-connected layer. The vector va for each action a is partitioned as va = [vWa , v H a , v D a ], where v W a and v H a have length K, and v D a has length K − 1. Instead of directly learning τk,a and qk,a associated with each monotonic knot, we propose to learn the normalized width and height of each bin. Here, vectors vWa and v H a are each passed through a softmax function and are interpreted as the normalized widths and heights. Vector vDa is regarded as the derivatives, and is passed through a softplus function to satisfy monotonicity.
With the width and height of each bin, τk,a and qk,a of each knot can be easily calculated by a cumulative sum. Since the values of vWa and v H a fall into [0, 1], each τk,a (k > 0) is computed by
τk,a = k∑ i=1 vWi,a, k = 1, ...,K; a = 1, ..., |A| (12)
without rescaling as the domain of a quantile function is [0, 1] (τ0,a = 0 by definition). To compute each qk,a, another transformation is required to rescale vHa to a range corresponding to the true quantile values. Inspired by NC-QR-DQN, we introduce the Bin Scale Network to generate two adaptive scale factors α and β by applying a fully connected layer C: Rd → R|A|×2 to the state embedding F(s). We compute the exponential of α to ensure the total bin height is positive. Then q0,a = βa and for k > 0, qk,a is computed by
qk,a = exp(αa)× k∑ i=1 vHi,a + βa, k = 1, ...,K; a = 1, ..., |A| (13)

heading: 3.3 APPROXIMATE WASSERSTEIN METRIC MINIMIZATION
text: When using continuous approximations of the quantile functions, there are several choices to compute the integral of the Wasserstein metric between two quantile functions. We can try to calculate the integral directly, but this is not straightforward for rational-quadratic functions since the integral rarely has a closed form. Alternatively, we can calculate the Riemann integral, but this leads to a loss function analogous to the L1-norm, which may cause instability in training. Thus, in this work, we perform quantile regression (Koenker & Hallock, 2001) in a projected space to approximately minimize the Wasserstein metric.
Let τ̃ = (τ̃0, ..., τ̃N ) be a fixed sequence of non-decreasing quantile fractions (note that the set τ̃ is different from the x-values of knots in Section 3.1 to partition the [0, 1] domain, which are learned by the neural network. In our experiments, we let τ̃ be uniformly fixed), we project the monotonic spline quantile function f to a quantile distribution space ZQ by computing
Zq(s, a) = N∑ i=1 (τ̃i − τ̃i−1)δq̂i(s,a), (14)
where each q̂i(s, a) is the corresponding quantile value at the quantile fraction τ̂i = τ̃i−1+τ̃i
2 given by f(τ̂i) with 1 ≤ i ≤ N . To compute f(τ̂i), we first search which bin τ̂i lies in. Then the value is returned by the corresponding spline function given hk(τ̂i) as input. In this case, the optimal value distribution Z is achieved by minimizing the 1-Wasserstein metric with Zq
W1(Z(s, a), Zq(s, a)) = N∑ i=1 ∫ τ̃i τ̃i−1 |F−1Z(s,a)(ω)− q̂i(s, a)|dω, (15)
which is equivalent to finding a projection operator ΠW1 such that
ΠW1Z := arg min Zq∈ZQ (Z,Zq). (16)
Furthermore, Dabney et al. (2018b) shows that the unique minimizer of this operator is given by
F−1Z(s,a)(τ̂i) = q̂i(s, a), τ̂i = τ̃i−1 + τ̃i
2 (17)
Proposition 1. [Proposition 2 in Dabney et al. (2018b)] Let ΠW1 be the quantile projector defined above. When applied to value distributions, it gives a projection for each state-value distribution. For any two value distributions Z1, Z2 ∈ Z for an MDP with countable state and action spaces,
d∞(ΠW1T πZ1,ΠW1T πZ2) ≤ γd∞(Z1, Z2), (18) where dp(Z1, Z2) = sups,aWp(Z1(s, a), Z2(s, a)) and Z is the space of action-value distributions with finite moments.
Proposition 1 suggests that after projecting f to Zq , the operator ΠW1T π is a γ-contraction under the measure d∞ and the repetition of this operator converges to a fixed point in space ZQ.
Based on Proposition 1, the ultimate goal is to estimate quantile values in Equation 17 for F−1Z(s,a) using quantile regression in each training batch. In our implementation, we uniformly fix τ̃ = (τ̃0, ..., τ̃N ) to be consistent with QR-DQN and NC-QR-DQN, which leads to the same quantile Huber loss as shown in Equation 6 and 7. However, the advantage of our method over QR-DQN and NC-QR-DQN is that we can freely enrich the density of τ̃ to get a better estimation of the quantile function without increasing the size of the model architecture, while QR-DQN and NC-QR-DQN must enlarge the output dimension of their models to get more quantile estimates. Since we can freely query quantile values at any quantile fraction, quantile fraction embedding as done in IQN and FQF is no longer necessary in our method.
Remark: Although one recent method NDQFN also learns continuous monotonic quantile functions, our method is different from NDQFN in three aspects. First, the x-values of those monotonic knots, i.e., τ0, ..., τK , are uniformly fixed in NDQFN, while they are trainable in our method. Second, by also learning the derivatives at each knot, we get a smooth interpolant over the entire domain, while NDQFN connects those knots with line segments, which has limited approximation ability. Third, to get the increments of y-values of those knots, i.e., q0, ..., qK , NDQFN learns a function taking the quantile fraction embeddings, i.e., the embeddings of corresponding τs, as input, while we do not calculate increments but use a scale network as discussed in Section 3.2.
To demonstrate the monotonicity and approximation strength of our method in stochastic environments, we plot the quantile functions learned by SPL-DQN, NC-QR-DQN, NDQFN, and QR-DQN in a variant of the classic Windy Gridworld domain (Sutton & Barto, 2018). In Figure 2a, the agent starts at the yellow circle state and makes standard moves in a gridworld to reach the red flag. A reward of−1 is earned at each step. Some columns are affected by some wind blowing from bottom to top. The orange line shows the optimal trajectory without stochasticity. We set each state transition to have probability 0.1 of moving in a random direction without any wind effect, otherwise the transition is affected by the wind, which pushes the agent northward. All methods here use the same training settings and similar network architectures as discussed in Appendix A.1. We compute the ground truth value distribution for an optimal policy (learned by policy iteration) at each state by performing one thousand Monte-Carlo (MC) rollouts and recording the observed returns as an empirical distribution. Then we transform the empirical distribution to the quantile function as the baseline. Here we show case the learned quantile functions at cyan square state and yellow circle state (start state) as shown in Figure 2b and 2c.
All these four methods eventually learn the optimal policy, however their quantile approximations are quite different. Without constraints, quantile functions given by QR-DQN clearly violate the monotonic property, which is known as the quantile crossing issue (Zhou et al., 2020). Although
NC-QR-DQN applies monotonic constraints, the estimated quantile range is biased towards smaller values according to the quantile functions given by MC, and we observe that the quantile functions learned by NC-QR-DQN are straight lines for some states, e.g. cyan square state, which means that it fails to learn the value distribution in those states, and in turn this leads to a biased estimation for the start state. The reason for this biased estimation is that in NC-QR-DQN, when rescaling the quantile range in Equation 8, a ReLU function is imposed to the coefficient α to ensure it is non-negative. However, this often sets α to zero and the quantile distribution will only depend on the shift parameter β (so this leads to a straight line). In this case, the value distribution cannot be precisely captured. For NDQFN, its quantile approximation at the goal nearing state (cyan square state) is close to SPL-DQN , but it overestimates the quantile range at the start state. We also observe the overestimation issue of NDQFN at another state in the middle of the orange line trajectory as shown in Figure 6 in the appendix (and more empirical results). Though still exhibiting estimation errors, the quantile functions learned by SPL-DQN are often the closest to the ground truth.
4 EXPERIMENTS
While most previous distributional RL algorithms were evaluated with Atari games from the Arcade Learning Environment (ALE), it was noted that the ALE is deterministic (Bellemare et al., 2017) and therefore questionable as a benchmark to evaluate distributional algorithms that are designed to capture environment stochasticity when there is none. However, we note that sticky actions can be used in Atari games to introduce stochasticity in policies (Machado et al., 2018) and this regime was used to evaluate IQN and FQF (Yang et al., 2019). When the environment is deterministic, value distributions still arise due to stochastic policies, stochastic approximations and random parameter initialization, but the resulting value distributions tend to be simple and close to deterministic. It is also well-known that deterministic environments possess optimal
policies that are open-loop and therefore ignore observations (Machado et al., 2018; Koul et al., 2019). In practice, it is often desirable to train controllers with simulators in which noise is injected to increase the robustness of the learned policies in case of discrepancies between the simulator and the real world. Hence, in this work, we modify several robotics environments by adding stochasticity, including one discrete environment from OpenAI Gym (Brockman et al., 2016) and nine continuous environments from PyBulletGym (Ellenberger, 2018–2019). We compare our method with QR-DQN, IQN, FQF, NC-QR-DQN, MM-DQN, and NDQFN. For MM-DQN, we used the unrectified Kernel κα(x, y) = −||x− y||α with α = 1 (parameter taken from Nguyen et al. (2021)) instead of the Gaussian kernel recommended by the authors when they tested on Atari games since the unrectified kernel gave better results in the robotics benchmarks used in this paper. A summary of how different QR-based methods compute the QR loss is provided in Appendix B. For a fair comparison, we made sure the same Feature Extractor architecture was used in different models. To simplify acronyms, we omit -DQN when referring to a method in what follows.

heading: 4.1 DISCRETE CONTROL IN CARTPOLE
text: We begin our experimental results in a stochastic environment with a discrete action space modified from Cartpole (Florian, 2007). The system is controlled by a force of +1 or−1 applied to the cart. A reward of +1 is returned if the pole remains upright. We set each state transition to have probability 0.05 of moving to a neighboring state to make the environment stochastic. The QR-based methods use N = 8 quantiles to compute the QR loss. MM-DQN uses N = M = 8 samples. More training details are provided in Appendix A.2.
As the episode rewards may vary significantly due to stochasticity, to better reflect the training process, we define the running score as a soft update of episode rewards:
running score = 0.99× running score+ 0.01× episode rewards (19) Figure 3 shows the running score curves for stochastic Cartpole. In general, SPL learns much faster (faster empirical convergence) than its counterparts. As discussed before, SPL can freely increase
the number of quantiles when performing quantile regression without enlarging the output dimension of the model. We further increase the number of quantiles to 24 to compute QR loss while keeping the number of bins unchanged (K = 8 and N = 24), yielding the curve labeled ’SPL1’ in Figure 3. This curve shows that approximately minimizing the Wasserstein metric with more quantiles leads to better quantile approximations and increases the learning speed and performance of SPL. As NDQFN also learns continuous quantile functions, we do the same experiment (K = 8 andN = 24) for NDQFN, whose training curve is labeled by ’NDQFN1’ in Figure 7 in the appendix. Tough its training performance improves, SPL with N = 24 is still better.

heading: 4.2 CONTINUOUS CONTROL IN PYBULLETGYM
text: PyBulletGym provides RoboSchool1, which is a free port of MuJoCo2. The state of these environments contains joint information of a robot and an action is a multi-dimensional continuous vector. We take nine environments from RoBoSchool and make them stochastic by introducing Gaussian noiseN (µ, σ) to both the location and velocity of each part of the robot, with µ = 0 and σ varying in different environments. We choose a reasonable σ for each environment such that robots won’t exhibit unrealistic motion. That is, for noise sensitive environments, such as Walker2D and Humanoid, we use a smaller σ, and for relatively easy tasks, like InvertedPendulumSwingup, we choose a bigger one. The noise setting for different environments is shown in Table 3 in the appendix.
To evaluate on continuous control tasks, we combine distributional RL with DDPG (Lillicrap et al., 2016) by modifying the critic, as done by Zhang & Yao (2019). Instead of learning Q, the critic learns the distribution Z directly. To handle continuous actions, the critic takes state-action pairs as input. As an exception, for the Humanoid environment, we combine distributional RL with SAC (Haarnoja et al., 2018) due to the fact that DDPG is not as good as SAC for this environment. To update the actor in DDPG and SAC, the expectation ofQ values is computed as the expectation of quantile samples given by the distributional critic. We refer to the original papers for hyperparameter settings, which are discussed in Appendix A.3. We also include raw DDPG and SAC as baselines.
Figure 4 shows the running score curves given by Equation 19 for these stochastic environments. Generally, the training performance varies among different approaches in different environments, however, in most cases, the quantile regression based methods who learn monotonic quantile representations are better than those whose quantile representations have no monotonicity guarantee, which clarifies that the quantile crossing issue can distort policy learning as pointed out by Zhou et al. (2020). Especially, for SPL, apart from Reacher and InvertedPendulumSwingup, it always converges faster and performs better during training. For InvertedPendulumSwingup, SPL performs comparably to NC-QR. Although NDQFN also learns continuous monotonic quantile functions, its performance is even worse than NC-QR in most cases, because NDQFN queries linear functions for quantile samples when computing QR loss, but the approximation ability of piecewise linear function is very limited. For methods with no monotonic quantile guarantee, we notice that although IQN is the best in Reacher, it performs worse in InvertedPendulum and Humanoid during training.
To further demonstrate the ability of our method to handle uncertainty of the environment, we slightly increase the noise in HalfCheetah toN (0, 0.008) (labeled by HalfCheetah1) andN (0, 0.01) (labeled by HalfCheetah2). The training curves in these two environments are shown in Figure 5. On average, QR and MM for DDPG behave poorly in these three HalfCheetah variants. The enhanced randomness of environments degrades the training performance of SPL, but SPL is generally faster and better than NC-QR and IQN, thanks to more precise quantile approximations.
During training, Ornstein-Uhlenheck noise OU(µ′, σ′) (Uhlenbeck & Ornstein, 1930) is utilized when selecting actions to induce exploration in DDPG. At the evaluation stage, the methods are executed with only exploitation (without action noise). We test the best models we get after training for each method, and the testing score across different environments are shown in Table 6 in Appendix A.3. Apart from Reacher, SPL outperforms its counterparts in all other domains. For the first two environments, although the training performances vary significantly among different methods, the testing scores of their best models are close to each other. For most remaining environments, the testing scores of SPL and NC-QR are significantly better than other methods. Appendix C provides
1https://openai.com/blog/roboschool/ 2http://www.mujoco.org
additional results for SPL and NDQFN when trained with uniformly spaced quantile fractions and random quantile fractions sampled from U([0, 1]).
Limitations: Our technique assumes full state observability. This is a limitation that could be removed in future work by extending it to partially observable domains. Similar to previous work in distributional RL, there is a computational cost associated with representing the value distribution. Hence our technique does not scale as well as non-disributional RL techniques. While we demonstrated our spline approximation in combination with DQN, DDPG and SAC, our distributional RL technique is limited to RL techniques that include a critic.

heading: 5 CONCLUSION
text: Based on previous works in distributional RL, we propose a more general and precise approximation for quantile functions using monotonic rational-quadratic splines when minimizing the Wasserstein metric. With a monotonic continuous representation of the quantile function, the quantile value at every quantile level is accessible during training, yielding greater accuracy. One direction for future work is to investigate other monotonic function approximators for quantile estimation.
ETHICS STATEMENT
We confirm that our work has no ethics issue.
REPRODUCIBILITY STATEMENT
To ensure reproducibility, we have discussed the detailed experiment settings in Appendix A. The code for the main experiments is submitted in the supplementary material.

heading: A EXPERIMENT DETAILS
text: A.1 WINDY GRIDWORLD
Training hyperparameters We train SPL-DQN with NC-QR-DQN, NDQFN, and QR-DQN in the stochastic Windy Gridworld. The common parameter settings are summarized in Table 1. The model input is the coordinates of the state. The -greedy parameter decreases by a half every two thousand episodes. To make the results comparable, we use the same Feature Extractor for the methods.
Specifically, for SPL-DQN and NDQFN, the number of bins is K = 30.
Additional results The learned quantile functions at the green square state is shown in Figure 6. NDQFN overestimates the quantile range.
To further demonstrate the approximation ability of SPL compared with NDQFN, we train SPL and NDQFN with 10 random seeds in windy gridworld, and check the quantile functions learnt by these two methods for the states on the orange line trajectory (apart from the goal state). There are 15 states on the trajectory times 10 seeds, which yield 150 quantile functions. For each state, we can use the shortest path to determine an upper bound on the return for that state. Whenever a qunatile function goes above that upper bound, we have a clear overestimation. For NDQFN, 36.7% (55/150) of the quantile functions yield an overestimation. In contrast, for SPL, only 4% (6/150) of the quantile functions yield an overestimation.
A.2 CARTPOLE
Training hyperparameters We train SPL-DQN with QR-DQN, IQN, FQF, NC-QR-DQN, MMDQN, and NDQFN in the stochastic Cartpole. The common parameter settings are summarized in Table 2. The model input is a vector of length 4, which contains the cart position, the cart velocity, the pole angle, and the pole angular velocity. The -greedy parameter decreases by 0.00005 every time step. To make the results comparable, we use the same Feature Extractor for the methods.
Specifically, for SPL-DQN and NDQFN, the number of bins is K = 8. For FQF, it contains a quantile fraction proposal network, whose learning rate is 2.5e−9, and the optimizer is RMSProp (Yang et al., 2019).
Additional results As NDQFN also learns continuous quantile functions, the number of quantile samples can be enlarged when computing QR loss. Here we train NDQFN with bin number K = 8 and quantile number N = 24. The training curve is labeled by ’NDQFN1’ as shown in Figure 7. The curves labeled by ’SPL1’, ’SPL’, and ’NDFQN’ are taken from Figure 3.
A.3 PYBULLETGYM
Noise setting for different environments We introduce different level of noise to the environments such that robots won’t exhibit unrealistic motion. The noise setting is shown in Table 3.
Training hyperparameters Hyperparameters for DDPG and DDPG based models are summarized in Table 4. The critic also uses an L2 weight decay of 10−2. The soft target update coefficient
is 0.001. Ornstein-Uhlenheck noise (OU(µ′, σ′)) (Uhlenbeck & Ornstein, 1930) is combined with actions for exploration in DDPG, where we use µ′ = 0 and σ′ = 0.1.
Hyperparameters for SAC and SAC based models are summarized in Table 5. The soft target update coefficient is 0.005.
For the critic implemented by SPL-DQN and NDQFN, the number of bins is K = 32. For the critic implemented by FQF, the learning rate for quantile fraction network is 2.5e−9, and the corresponding optimizer is RMSProp.
Additional results We test the best models given by different methods and perform four random seeds. The results are shown in Table 6. We test all DDPG based agents without Ornstein-Uhlenheck noise for 0.125 million frames, and SAC based agents for 2.5 thousand episodes. We treat DDPG and SAC scores as baselines and scale other methods’ scores by them, i.e.
method scaled test score = method raw test score
DDPG/SAC raw test score (20)

heading: B COMPUTING QR LOSS
text: We summarize how different QR-based distributional RL methods sample quantile values when computing the QR loss. Since the QR loss is computed in a TD manner, we will need N current quantile samples ({q1i }, i = 1, ..., N ) and N ′ target quantile samples ({q2i }, i = 1, ..., N ′) corresponding to two quantile fraction sets ({τ1i }, i = 1, ..., N ) and ({τ2i }, i = 1, ..., N ′). Without loss of generalization, we consider N = N ′. Here we discuss the case with discrete actions, and denote the action space by |A|. For discrete quantile approximations, including QR-DQN, NC-QR-DQN, IQN, and FQF, in order to get N quantile samples for each action, the output dimension of the model is |A|×N for an input state. For QR-DQN and NC-QR-DQN, {τ1i } and {τ2i } are assumed to be uniformly spaced. For
IQN, {τ1i } and {τ2i } are independently drawn from a uniform distribution U([0, 1]). For FQF, {τ1i } and {τ2i } are proposed by a quantile fraction network. For methods that learn a continuous approximation of the quantile function, including SPL-DQN and NDQFN, the output (for an input state) consists of knots with shape |A| × (K + 1) when the domain is divided intoK bins. For SPL-DQN, it leans the x, y values, and derivatives of those knots. A smooth continuous function with closed form is obtained in each bin. When sampling quantile values to compute the QR loss, SPL-DQN uniformly fixes {τ1i } and {τ2i }, and {q1i } and {q2i } are obtained by querying the closed form with {τ1i } and {τ2i } as inputs. For NDQFN, it only learns the y-values of those knots, and the x-values of the knots are uniformly spaced. The continuous function is constructed by connecting neighboring knots with linear functions. When sampling quantile values to compute the QR loss, NDQFN draws {τ1i } and {τ2i } from a uniform distribution U([0, 1]) independently, and {q1i } and {q2i } are obtained by querying the linear functions in each bin.

heading: C COMPARISON OF SPL AND NDQFN WITH TWO TRAINING REGIMES
text: Figure 8 shows additional results that compare SPL and NDQFN when trained with uniformly spaced points versus points sampled uniformly at random. In this figure, SPL corresponds to training with uniformly spaced quantile fractions as described in the main paper. NDQFN corresponds to training with quantile fractions sampled from U([0, 1]) as described in (Zhou et al., 2021). To iron out this difference in the training, we also report SPL-rnd where training is done with quantile fractions sampled from U([0, 1]) and NDQFN-uni where training is done with uniformly spaced quantile fractions. We observe that SPL outperforms NDQFN in both training regimes.

heading: D DDPG AND SAC WITH DISTRIBUTIONAL CRITIC
text: We summarize the algorithm of using distributional critic for DDPG and SAC in this section. For FQF critic, it has to update the quantile fraction proposal network separately, so it is individually described. Similarly, for MM critic, it use moment matching instead of quantile regression, and it is individually described as well. For SAC based methods, readers can refer to the original paper for how to compute gradient for policy and state value.
Algorithm 1 DDPG with QR-based distributional critic (apart from FQF) Require: Initialize critic network Z(s, a|ψ) with weights ψ; Initialize actor network µ(s|θ) with
weights θ; Initialize target network Z ′ and µ′ with weights ψ′ ← ψ, θ′ ← θ; Initialize reply buffer D
1: for each episode do 2: Initialize random process OU for action exploration 3: for each time step t do 4: at ∼ µ(st|θ) +OU t 5: st+1 ∼ p(st+1|st, at) 6: D ← D ∪ {(st, at, r(st, at), st+1)} 7: Sample minibatch of N transitions (si, ai, ri, si+1) from D 8: Choose current quantile fractions {τ} according to critic’s strategy 9: Compute corresponding current quantiles {qi} ← Z(si, ai|ψ)
10: Choose target quantile fractions {τ̄} according to critic’s strategy 11: Compute corresponding target quantiles {q̄i+1} ← ri + γZ ′(si+1, µ′(si+1|θ′)|ψ′) 12: Update critic by minimizing QR loss 13: Compute expectation of quantiles Qs,a ← E[Z(s, a|ψ)] 14: Update actor by 15: ∇θJ ≈ 1N ∑ i∇aQs,a|s=si,a=µ(si)∇θµ(s|µ)|si 16: Update target networks 17: ψ′ ← σψ + (1− σ)ψ′ 18: θ′ ← σθ + (1− σ)θ′ 19: end for 20: end for
Algorithm 2 DDPG with QR-based distributional critic (FQF) Require: Initialize critic value network Z(s, a|ψ) with weights ψ; Initialize actor network µ(s|θ)
with weights θ; Initialize target network Z ′ and µ′ with weights ψ′ ← ψ, θ′ ← θ; Initialize reply buffer D 1: Initialize critic fraction proposal network P (s, a|ω) with weights ω and its target P ′(s, a|ω′) with weights ω′ 2: for each episode do 3: Initialize random process OU for action exploration 4: for each time step t do 5: at ∼ µ(st|θ) +OU t 6: st+1 ∼ p(st+1|st, at) 7: D ← D ∪ {(st, at, r(st, at), st+1)} 8: Sample minibatch of N transitions (si, ai, ri, si+1) from D 9: Compute current quantile fractions {τ} ← P (si, ai|ω)
10: Compute corresponding current quantiles {qi} ← Z(si, ai|ψ) 11: Choose target quantile fractions {τ̄} ← P (si+1, µ′(si+1|θ′)|ω) 12: Compute corresponding target quantiles {q̄i+1} ← ri + γZ ′(si+1, µ′(si+1|θ′)|ψ′) 13: Update critic fraction proposal network by 14: ∂ω∂τi = 2Z[τi]− Z[τ̂i]− Z[τ̂i−1], τ̂ = τi+τi+1 2 15: Update critic value network by minimizing QR loss 16: Compute expectation of quantiles Qs,a ← E[Z(s, a|ψ)] 17: Update actor by 18: ∇θJ ≈ 1N ∑ i∇aQs,a|s=si,a=µ(si)∇θµ(s|µ)|si 19: Update target networks 20: ψ′ ← σψ + (1− σ)ψ′ 21: θ′ ← σθ + (1− σ)θ′ 22: ω′ ← σω + (1− σ)ω′ 23: end for 24: end for
Algorithm 3 DDPG with MM distributional critic Require: Initialize critic network Z(s, a|ψ) with weights ψ; Initialize actor network µ(s|θ) with
weights θ; Initialize target network Z ′ and µ′ with weights ψ′ ← ψ, θ′ ← θ; Initialize reply buffer D
1: for each episode do 2: Initialize random process OU for action exploration 3: for each time step t do 4: at ∼ µ(st|θ) +OU t 5: st+1 ∼ p(st+1|st, at) 6: D ← D ∪ {(st, at, r(st, at), st+1)} 7: Sample minibatch of N transitions (si, ai, ri, si+1) from D 8: Compute current Q samples {qi} ← Z(si, ai|ψ) 9: Compute target Q samples {q̄i+1} ← ri + γZ ′(si+1, µ′(si+1|θ′)|ψ′)
10: Update critic by minimizing MMD loss 11: Compute expectation of Q values Qs,a ← E[Z(s, a|ψ)] 12: Update actor by 13: ∇θJ ≈ 1N ∑ i∇aQs,a|s=si,a=µ(si)∇θµ(s|µ)|si 14: Update target networks 15: ψ′ ← σψ + (1− σ)ψ′ 16: θ′ ← σθ + (1− σ)θ′ 17: end for 18: end for
Algorithm 4 SAC with QR-based distributional critic (apart from FQF) Require: The learning rates λπ , λZ , and λV for functions πθ, Zw, and Vψ; Initialize parameters θ,
w, ψ, ψ̄; Initialize reply buffer D 1: for each iteration do 2: for each time step t do 3: at ∼ πθ(st) 4: st+1 ∼ p(st+1|st, at) 5: D ← D ∪ {(st, at, r(st, at), st+1)} 6: end for 7: for each gradient update step do 8: Choose current and target quantile fractions according to critic’s strategy 9: Compute current and target quantile values (for computing JZ)
10: ψ ← ψ − λV∇ψJV (ψ) 11: wi ← wi − λZ∇wiJZ(wi) for i ∈ {1, 2} (JZ is QR loss) 12: θ ← θ − λπ∇θJπ(θ) 13: ψ̄ ← σψ + (1− σ)ψ̄ 14: end for 15: end for
Algorithm 5 SAC with QR-based distributional critic (FQF) Require: The learning rates λπ , λZ , and λV for functions πθ, Zw, and Vψ; Initialize parameters
θ, w, ψ, ψ̄; Initialize reply buffer D; The learning rate λP for fraction proposal network Pφ; Initialize parameter φ
1: for each iteration do 2: for each time step t do 3: at ∼ πθ(st) 4: st+1 ∼ p(st+1|st, at) 5: D ← D ∪ {(st, at, r(st, at), st+1)} 6: end for 7: for each gradient update step do 8: Compute current and target quantile fractions using Pφ 9: Compute current and target quantile values (for computing JZ)
10: ψ ← ψ − λV∇ψJV (ψ) 11: φ← φ− λφ∇φJP (φ) (JP is quantile fraction loss) 12: wi ← wi − λZ∇wiJZ(wi) for i ∈ {1, 2} (JZ is QR loss) 13: θ ← θ − λπ∇θJπ(θ) 14: ψ̄ ← σψ + (1− σ)ψ̄ 15: end for 16: end for
Algorithm 6 SAC with QR-based distributional critic (MM) Require: The learning rates λπ , λZ , and λV for functions πθ, Zw, and Vψ; Initialize parameters θ,
w, ψ, ψ̄; Initialize reply buffer D 1: for each iteration do 2: for each time step t do 3: at ∼ πθ(st) 4: st+1 ∼ p(st+1|st, at) 5: D ← D ∪ {(st, at, r(st, at), st+1)} 6: end for 7: for each gradient update step do 8: Compute current and target Q value samples (for computing JZ) 9: ψ ← ψ − λV∇ψJV (ψ)
10: wi ← wi − λZ∇wiJZ(wi) for i ∈ {1, 2} (JZ is MMD loss) 11: θ ← θ − λπ∇θJπ(θ) 12: ψ̄ ← σψ + (1− σ)ψ̄ 13: end for 14: end for

Review summaries:
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.  

This paper proposes SPL-DQN, a quantile-based distributional RL method that uses monotonic rational-quadratic splines to approximate the quantile function of the cumulative return (starting from each state-action pair), and the features of using this technique are: (i) The approximated quantile function is ensured to be monotonic; (ii) The monotonic rational-quadratic splines are continuously differentiable and can provide a more flexible class of smooth approximators for the quantile function (compared to piecewise linear or step functions). Moreover, the feature (i) resolves the quantile crossing problem of several existing distributional RL methods (e.g. QR-DQN), under which the quantiles output from the neural network may not be monotonic. By conducting experiments with a toy example, this paper also empirically discovers a few potential issues with the relevant prior works that enforce monotonicity (e.g., NC-QR-DQN and NDQFN). Experimental results in both tasks with discrete (Cartpole) and continuous action spaces (Roboschool) are provided to demonstrate the performance of SPL-DQN.

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.