D ISTRIBUTIONAL R EINFORCEMENT L EARNING WITH M ONOTONIC S PLINES
A BSTRACT
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 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.
I NTRODUCTION
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 sufﬁcient to represent an expected return by a single mean value, errors due to ﬁnite 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 beneﬁts 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 ﬁxed 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, C optimizes the cross-entropy loss with a Cram ´ er-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 ﬁnite 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 inﬁnite 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 inﬁnite 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 ﬂexible and smooth approximation.
With sufﬁciently 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.
D ISTRIBUTIONAL R EINFORCEMENT L EARNING
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 ) deﬁnes 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 ) = (cid:80) ∞ t =
γ t R ( s t , a t ) , where s
=
s , a 0 = a , s t +1 ∼ P ( ·| s t , a t ) , and a t ∼ π ( ·| s t ) .
The Q -value (state-action value) is deﬁned as Q π ( s, a ) =
E [ Z π ( s, a )] .
The optimal Q -value, Q ∗ ( s, a ) = max π Q π ( s, a ) , is the unique ﬁxed point of the Bellman optimality operator T (Bellman, 1966)
[Equation]
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
[Equation]
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)
[Equation] [Equation]
with S (cid:48) ∼ P ( ·| s, a ) , A (cid:48) ∼ π ( ·| S (cid:48) ) , 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
[Equation]
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. Q UANTILE R EGRESSION FOR D ISTRIBUTIONAL RL
In QR-DQN (Dabney et al., 2018b), the random return is approximated by a uniform mixture of N
Diracs
[Equation]
with each θ i set to a ﬁxed quantile fraction, ˆ τ i =
τ i −
+ τ ifor 1 ≤ i ≤ N , and τ i = i/N .
The quantile estimation is performed by minimizing the quantile Huber loss, with threshold η
[Equation]
on the pairwise TD error δ ij = r + γθ j ( s (cid:48) , a (cid:48) )
− θ i ( s, a ) , where
[Equation]
, with
[Equation]
Based on QR-DQN, Dabney et al.
(2018a) proposed to sample quantile fractions from a base distribution, e.g. τ ∈ U ([0 , 1]) rather than ﬁxing 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 (cid:88)
j = χ j,a , and χ j,a = softmax(Ω( s ))
j,a
[Equation] [Equation]
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 ﬁrst learning a baseline value and then adding non-negative increments.
2. O THER D ISTRIBUTIONAL M ETHODS
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:
[Equation]
where { υ i } Ni =1 ∼ Z ( s, a ) and { ψ i } Mi =1 ∼ R
( s, a ) +
γZ ( S (cid:48) , A (cid:48) ) .
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 efﬁciency.
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.
O UR A LGORITHM
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 inﬁnite fractions.
But in practice, one can only use ﬁnite 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.
3. Q UANTILE A PPROXIMATION WITH M ONOTONIC R ATIONAL Q UADRATIC S PLINES
Monotonic splines produce a monotonic interpolant to a set of monotonic data points (called knots) { ( x k , y k ) }
Kk =0 , which has been recently used as a transformation function in normalizing ﬂows (Durkan et al., 2019; Dolatabadi et al., 2020).
Furthermore, denote { d k } Kk =0 , a set of positive numbers, as the derivative of each knot.
The monotonic rational-quadratic splines aim to ﬁnd rational-quadratic functions with form f k ( x ) =
O k ( x ) P k ( x ) to ﬁt the points and derivatives in each interval (called bin)
[ x k , x k +1 ] , where O k and P k are quadratic functions (with form ax 2 + bx + c ).
Gregory & Delbourgo (1982) suggested to construct O k and P k as follows.
Denote g k =
( y k +1 − y k ) /
( x k +1 − x k ) and h k ( x ) =
( x − x k ) /
( x k +1 − x k ) for x ∈ [ x k , x k +1 ] .
The expressions for the quadratic O k ( h k ( x )) and P k ( h k ( x )) for the k th bin is deﬁned by (use h k for short of h k ( x ) ):
[Equation]
Then, the rational-quadratic function for the k th bin is computed by the quotient of O k and P k
[Equation]
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 ﬁts the non-decreasing property of F − 1 .
Let F − 1 Z ( 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 − 1 Z ( s,a ) are threefold.
First, propose a partition for the domain of deﬁnition [0 , 1] with τ 0 < ... <
τ k < ... <
τ K .
Here τ 0 =and τ K = 1 .
Second, estimate the corresponding quantile values q 0 < ... <
q k < ... <
q K .
Third, assess the derivatives at those points with d 0 , ..., d k , ..., d K .
We give a small ﬁxed positive value for d 0 and d K 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 x k by τ k and y k by q k ).
3. M ODEL I MPLEMENTATION
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 modiﬁed 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 ) ∈ R d of state s .
Then the Knots Logit Network W maps F ( s ) to unconstrained logits v with dimension |A| × (3 K − 1) using a fully-connected layer.
The vector v a for each action a is partitioned as v a =
[ v Wa , v Ha , v Da ] , where v Wa and v Ha have length K , and v Da has length K − 1 .
Instead of directly learning τ k,a and q k,a associated with each monotonic knot, we propose to learn the normalized width and height of each bin.
Here, vectors v Wa and v Ha are each passed through a softmax function and are interpreted as the normalized widths and heights.
Vector v Da 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 q k,a of each knot can be easily calculated by a cumulative sum.
Since the values of v Wa and v Ha fall into [0 , 1] , each τ k,a ( k > 0 ) is computed by
[Equation] [Equation]
without rescaling as the domain of a quantile function is [0 , 1] ( τ 0 ,a = 0 by deﬁnition).
To compute each q k,a , another transformation is required to rescale v Ha 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 : R d → R |A|× 2 to the state embedding F ( s ) .
We compute the exponential of α to ensure the total bin height is positive.
Then q 0 ,a = β a and for k > 0 , q k,a is computed by
[Equation]
3. A PPROXIMATE W ASSERSTEIN M ETRIC M INIMIZATION
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 ﬁxed 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 ﬁxed), we project the monotonic spline quantile function f to a quantile distribution space Z Q by computing
[Equation]
where each ˆ q i ( s, a ) is the corresponding quantile value at the quantile fraction ˆ τ i
=
˜
τ i − 1 +˜ τ igiven by f (ˆ τ i ) with 1 ≤ i ≤ N .
To compute f (ˆ τ i ) ,
we ﬁrst search which bin ˆ τ
i lies in.
Then the value is returned by the corresponding spline function given h k
(ˆ τ i ) as input.
In this case, the optimal value distribution Z is achieved by minimizing the 1-Wasserstein metric with Z q
[Equation]
which is equivalent to ﬁnding a projection operator Π W 1 such that
[Equation]
Furthermore, Dabney et al.
(2018b) shows that the unique minimizer of this operator is given by
[Equation]
Proposition 1.
[Proposition 2 in Dabney et al.
(2018b)]
Let Π W 1 be the quantile projector deﬁned above.
When applied to value distributions, it gives a projection for each state-value distribution.
For any two value distributions Z 1 , Z 2 ∈ Z for an MDP with countable state and action spaces,
[Equation]
where d p ( Z 1 , Z 2 )
= sup s,a W p ( Z 1 ( s, a ) , Z 2 ( s, a )) and Z is the space of action-value distributions with ﬁnite moments.
Proposition 1 suggests that after projecting f to Z q , the operator Π W 1 T π is a γ -contraction under the measure d ∞ and the repetition of this operator converges to a ﬁxed point in space Z Q .
Based on Proposition 1, the ultimate goal is to estimate quantile values in Equation 17 for F − 1 Z ( s,a ) using quantile regression in each training batch.
In our implementation, we uniformly ﬁx ˜ τ
=
(˜ τ 0 , ..., ˜ τ N ) to be consistent with QR-DQN and NC-QR-DQN, which leads to the same quantile
Huber loss as shown in Equations 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 NCQR-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 ﬁxed 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., q 0 , ..., q K , 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 ﬂag.
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 Figures 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 coefﬁcient α to ensure it is non-negative.
However, this often sets α to zero and the quantile distribution will only depend on the shift parameter β (which 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.
Though still exhibiting estimation errors, the quantile functions learned by SPL-DQN are often the closest to the ground truth.
E XPERIMENTS
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 unrectiﬁed 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 unrectiﬁed 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.
4. D ISCRETE CONTROL IN C ARTPOLE
We begin our experimental results in a stochastic environment with a discrete action space modiﬁed 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 signiﬁcantly due to stochasticity, to better reﬂect the training process, we deﬁne the running score as a soft update of episode rewards:
[Equation]
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 the 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 and N = 24 ) for NDQFN, whose training curve is labeled by ’NDQFN1’ in Figure 7 in the appendix.
Although its training performance improves, SPL with N = 24 is still better.
4. C ONTINUOUS CONTROL IN P Y B ULLET G YM
PyBulletGym provides RoboSchool 1 , which is a free port of MuJoCo 2 .
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 noise N ( µ, σ ) 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 of Q 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 clariﬁes that the quantile crossing issue can distort policy learning as pointed out by Zhou et al.
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 to N (0 , 0 . 008) (labeled by HalfCheetah1) and N (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 ( µ (cid:48) , σ (cid:48) )
(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 ﬁrst two environments, although the training performances vary signiﬁcantly 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 signiﬁcantly better than other methods.
Appendix C provides 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-distributional 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.
C ONCLUSION
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.
A CKNOWLEDGMENTS
We acknowledge the funding from the Canada CIFAR AI Chairs program, the Natural Sciences and Engineering Research Council of Canada (NSERC) and support from SportLogiq.
Resources used in this work were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute https://vectorinstitute.
ai/partners/ .
E THICS S TATEMENT
We conﬁrm that our work has no ethics issue.
R EPRODUCIBILITY S TATEMENT
To ensure reproducibility we have discussed the detailed experiment settings in Appendix A.
The code for the main experiments is released in the supplementary material.
A E XPERIMENT D ETAILS A. W INDY G RIDWORLD
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 inputs are the coordinates of the state.
The (cid:15) -greedy parameter decreases by a half every two thousand episodes.
To make the results comparable, we use the same Feature Extractor for the methods.
For SPL-DQN and NDQFN, the number of bins is K = 30 .
[Table]
Additional results The learned quantile functions for the green square state are 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 quantile 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. square state
A. C ARTPOLE
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 (cid:15) -greedy parameter decreases by 0 .
00005 every time step.
To make the results comparable, we use the same Feature Extractor for the methods.
Speciﬁcally, 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 .
5 e − 9 , and the optimizer is RMSProp (Yang et al., 2019).
[Table]
Additional results As NDQFN also learns continuous quantile functions, the number of quantile samples can be enlarged when computing the 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.
[Figure]
Noise setting for different environments We introduce different noise levels in PyBulletGym’s environments while ensuring that the robots won’t exhibit unrealistic motion.
The noise settings are shown in Table 3.
[Table]
InvertedPendulum
N (0 , 0 . 02)
InvertedDoublePendulum
N (0 , 0 . 01)
InvertedPendulumSwingup
N (0 , 0 . 05)
Reacher
N (0 , 0 . 01)
Walker2D
N (0 , 0 . 005)
HalfCheetah
N (0 , 0 . 005)
HalfCheetah
N (0 , 0 . 008)
HalfCheetah
N (0 , 0 . 01)
Ant
N (0 , 0 . 01)
Hopper
N (0 , 0 . 003)
Humanoid
N (0 , 0 . 003)
Training hyperparameters Hyperparameters for DDPG and DDPG based models are summarized in Table 4.
The critic also uses an L 2 weight decay of 10 − 2 .
The soft target update coefﬁcient is 0.001.
Ornstein-Uhlenheck noise ( OU ( µ (cid:48) , σ (cid:48) ) )
(Uhlenbeck & Ornstein, 1930) is combined with actions for exploration in DDPG, where we use µ (cid:48)
= 0 and σ (cid:48)
= 0 .
Hyperparameters for SAC and SAC based models are summarized in Table 5.
The soft target update coefﬁcient 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 .
5 e − [Table] and the corresponding optimizer is RMSProp.
[Table]
[Table]
Additional results We test the best models given by different methods with 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
[Equation]
= method raw test
[Equation]
B C OMPUTING QR LOSS
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 ( { q 1 i } , i = 1 , ..., N ) and N (cid:48) target quantile samples ( { q 2 i } , i = 1 , ..., N (cid:48) ) corresponding to two quantile fraction sets ( { τ 1 i } , i = 1 , ..., N ) and ( { τ 2 i } , i = 1 , ..., N (cid:48) ).
Without loss of generalization, we consider N = N (cid:48) .
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 [Table]
state. For QR-DQN and NC-QR-DQN, { τ 1 i } and { τ 2 i } are assumed to be uniformly spaced.
For
IQN, { τ 1 i } and { τ 2 i } are independently drawn from a uniform distribution U ([0 , 1]) .
For FQF, { τ 1 i } and { τ 2 i } 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 into K bins.
For SPL-DQN, it leanrs 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 ﬁxes { τ 1 i } and { τ 2 i } , and { q 1 i } and { q 2 i } are obtained by querying the closed form with { τ 1 i } and { τ 2 i } 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 { τ 1 i } and { τ 2 i } from a uniform distribution U ([0 , 1]) independently, and { q 1 i } and { q 2 i } are obtained by querying the linear functions in each bin.
C C OMPARISON OF SPL AND NDQFN WITH TWO TRAINING REGIMES
Figure 8 shows additional results that compare SPL and NDQFN when trained with uniformly spaced points versus points sampled uniformly at random.
In this ﬁgure, 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.
D DDPG AND SAC WITH DISTRIBUTIONAL CRITIC
We summarize the algorithm when 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 uses 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.
Require:
Initialize critic network Z ( s, a | ψ ) with weights ψ ; Initialize actor network µ ( s | θ ) with weights θ ; Initialize target network Z (cid:48) and µ (cid:48) with weights ψ (cid:48) ← ψ , θ (cid:48) ← θ ; 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: a t ∼ µ ( s t | θ ) + OU t 5: s t +1 ∼ p ( s t +1 | s t , a t ) 6:
D ← D ∪ { ( s t , a t , r ( s t , a t ) , s t +1 ) } 7:
Sample minibatch of N transitions ( s i , a i , r i , s i +1 ) from D 8:
Choose current quantile fractions { τ } according to critic’s strategy 9:
Compute corresponding current quantiles { q i } ← Z ( s i , a i | ψ ) 10:
Choose target quantile fractions { ¯ τ } according to critic’s strategy 11:
Compute corresponding target quantiles { ¯ q i +1 } ← r i + γZ (cid:48) ( s i +1 , µ (cid:48)
( s i +1 | θ (cid:48) )
| ψ (cid:48) ) 12:
Update critic by minimizing QR loss 13:
Compute expectation of quantiles Q s,a ← E [ Z ( s, a | ψ )] 14:
Update actor by 15: ∇
[Equation]
∇
[Equation]
s
[Equation]
µ ) | s i 16:
Update target networks 17: ψ (cid:48) ← σψ + (1 − σ ) ψ (cid:48) 18: θ (cid:48)
← σθ + (1 − σ ) θ (cid:48) 19: end for 20: end for
Require:
Initialize critic value network
Z ( s, a | ψ ) with weights ψ ; Initialize actor network µ ( s | θ ) with weights θ ; Initialize target network Z (cid:48) and µ (cid:48) with weights ψ (cid:48)
← ψ , θ (cid:48)
← θ ; Initialize reply buffer D 1: Initialize critic fraction proposal network P ( s, a | ω ) with weights ω and its target P (cid:48) ( s, a | ω (cid:48) ) with weights ω (cid:48) 2: for each episode do 3:
Initialize random process OU for action exploration 4: for each time step t do 5: a t ∼ µ ( s t | θ ) + OU t 6: s t +1 ∼ p ( s t +1 | s t , a t ) 7:
D ← D ∪ { ( s t , a t , r ( s t , a t ) , s t +1 ) } 8:
Sample minibatch of N transitions ( s i , a i , r i , s i +1 ) from D 9:
Compute current quantile fractions { τ } ← P ( s i , a i | ω ) 10:
Compute corresponding current quantiles { q i } ← Z ( s i , a i | ψ ) 11:
Choose target quantile fractions { ¯ τ } ← P ( s i +1 , µ (cid:48)
( s i +1 | θ (cid:48) ) | ω ) 12:
Compute corresponding target quantiles { ¯ q i +1 } ← r i + γZ (cid:48) ( s i +1 , µ (cid:48)
( s i +1 | θ (cid:48) )
| ψ (cid:48) ) 13:
Update critic fraction proposal network by 14:
[Equation]
Update critic value network by minimizing QR loss 16:
Compute expectation of quantiles Q s,a ← E [ Z ( s, a | ψ )] 17:
Update actor by 18: ∇ θ J ≈ 1 N (cid:80) i ∇ a Q s,a | s = s i ,a = µ ( s i ) ∇ θ µ ( s | µ ) | s i 19:
Update target networks 20: ψ (cid:48) ← σψ + (1 − σ ) ψ (cid:48) 21: θ (cid:48) ← σθ + (1 − σ ) θ (cid:48) 22: ω (cid:48) ← σω + (1 − σ ) ω (cid:48) 23: end for 24: end for
Algorithm
Require:
Initialize critic network Z ( s, a | ψ ) with weights ψ ; Initialize actor network µ ( s | θ ) with weights θ ; Initialize target network Z (cid:48) and µ (cid:48) with weights ψ (cid:48) ← ψ , θ (cid:48) ← θ ; 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: a t ∼ µ ( s t | θ ) + OU t 5: s t +1 ∼ p ( s t +1 | s t , a t ) 6:
D ← D ∪ { ( s t , a t , r ( s t , a t ) , s t +1 ) } 7:
Sample minibatch of N transitions ( s i , a i , r i , s i +1 ) from D 8:
Compute current Q samples { q i } ← Z
( s i , a i | ψ ) 9:
Compute target Q samples { ¯ q i +1 } ← r i + γZ (cid:48)
( s i +1 , µ (cid:48) ( s i +1 | θ (cid:48) )
| ψ (cid:48) ) 10:
Update critic by minimizing MMD loss 11:
Compute expectation of Q values Q s,a ← E [ Z ( s, a | ψ )] 12:
Update actor by 13:
[Equation]
Update target networks 15:
[Equation]
[Equation]
17: end for 18: end for
Require: The learning rates λ π , λ Z , and λ V for functions π θ , Z w , and V ψ ; Initialize parameters θ , w , ψ , ¯ ψ ; Initialize reply buffer D 1: for each iteration do 2: for each time step t do 3: a t ∼ π θ ( s t ) 4: s t +1 ∼ p
( s t +1 | s t , a t ) 5:
D ← D ∪ { ( s t , a t , r ( s t , a t ) , s t +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 J Z ) 10: ψ ← ψ − λ V ∇ ψ J V ( ψ ) 11: w i ← w i − λ Z ∇ w i J Z ( w i ) for i ∈ { 1 , 2 } ( J Z is QR loss) 12: θ ← θ − λ π ∇ θ J π ( θ ) 13: ¯ ψ ← σψ + (1 − σ ) ¯ ψ 14: end for 15: end for
Require: The learning rates λ π , λ Z , and λ V for functions π θ , Z w , 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: a t ∼ π θ ( s t ) 4: s t +1 ∼ p ( s t +1 | s t , a t ) 5:
D ← D ∪ { ( s t , a t , r ( s t , a t ) , s t +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 J Z ) 10: ψ ← ψ − λ V ∇ ψ J V ( ψ ) 11: φ ← φ − λ φ ∇ φ J P ( φ ) ( J P is quantile fraction loss) 12: w i ← w i − λ Z ∇ w i J Z ( w i ) for i ∈ { 1 , 2 } ( J Z 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 π θ , Z w , and V ψ ; Initialize parameters θ , w , ψ , ¯ ψ ; Initialize reply buffer D 1: for each iteration do 2: for each time step t do 3: a t ∼ π θ ( s t ) 4: s t +1 ∼ p
( s t +1 | s t , a t ) 5:
D ← D ∪ { ( s t , a t , r ( s t , a t ) , s t +1 ) } 6: end for 7: for each gradient update step do 8:
Compute current and target Q value samples (for computing J Z ) 9: ψ ← ψ − λ V ∇ ψ J V ( ψ ) 10: w i ← w i − λ Z ∇ w
i J Z ( w i ) for i ∈ { 1 , 2 } ( J Z is MMD loss) 11: θ ← θ − λ π ∇ θ J π ( θ ) 12: ¯ ψ ← σψ + (1 − σ ) ¯ ψ 13: end for 14: end for