Non-decreasing Quantile Function Network with Efficient Exploration for Distributional Reinforcement Learning

1 Introduction

Distributional reinforcement learning (DRL) algorithms [Jaquette, 1973; White, 1988; Bellemare et al., 2017; Dabney et al., 2018b; Dabney et al., 2018a; Yang et al., 2019], different from the value based methods [Watkins, 1989; Mnih et al., 2013] which focus on the expectation of the return, characterize the cumulative reward as a random variable and attempt to approximate its whole distribution. Most existing DRL methods fall into two main classes according to their ways to model the return distribution. One class, including categorical DQN [Bellemare et al., 2017], C51, and Rainbow [Hessel et al., 2018], assumes that all the possible returns are bounded in a known range and learns the probability of each value through interacting with the environment. Another class, for example the QR-DQN [Dabney et al., 2018b], tries to obtain the quantile estimates at some fixed locations by minimizing the Huber loss [Huber, 1992] of quantile regression [Koenker, 2005]. To more precisely parameterize the entire distribution, some quantile value based methods, such as IQN [Dabney et al., 2018a] and FQF [Yang et al., 2019], are proposed to learn a continuous map from any quantile fraction τ ∈ (0, 1) to its estimate on the quantile curve.

∗Contact Author

The theoretical validity of QR-DQN [Dabney et al., 2018b], IQN [Dabney et al., 2018a] and FQF [Yang et al., 2019] heavily depends on a prerequisite that the approximated quantile curve is non-decreasing. Unfortunately, since no global constraint is imposed when simultaneously estimating the quantile values at multiple locations, the monotonicity can not be ensured using their network designs. At early training stage, the crossing issue is even more severe given limited training samples. Some existing studies try to solve this problem [Zhou et al., 2020; Tang Nguyen et al., 2020]. However, their main architecture is built on a set of fixed quantile locations and not applicable to quantile value based algorithms such as IQN and FQF. Some most recent work [Yue et al., 2020] naturally ensures the monotonicity of the learned quantile function. However, it requires extremely large computation cost and performs significantly worse in Atari games than the proposed method in this paper (a brief comparison is provided in Section H of the supplement).

Another problem to be solved is how to design an efficient exploration method for DRL. Most existing exploration techniques are originally designed for non-distributional RL [Bellemare et al., 2016; Ostrovski et al., 2017; Machado et al., 2018], and very few of them can work for DRL. Mavrin et al.[2019] proposes a DRL-based exploration method, DLTV, for QR-DQN by using the left truncated variance as the exploration bonus. However, this approach can not be directly applied to quantile value based algorithms since the original DLTV method requires all the quantile locations to be fixed while IQN or FQF resamples the quantile locations at each training iteration and the bonus term could be extremely unstable.

To address these two common issues in DRL studies, we propose a novel algorithm called Non-Decreasing Quantile Function Network (NDQFN), together with an efficient exploration strategy, distributional prediction error (DPE), designed for DRL. The NDQFN architecture allows us to approximate the quantile distribution of the return by using a non-decreasing piece-wise linear function. The monotonicity of N + 1 fixed quantile levels is ensured by an incremental structure, and the quantile value at any τ ∈ (0, 1) can be estimated as the weighted sum of its two nearest neighbors among the N + 1 locations. DPE uses the 1-Wasserstein distance between the quantile distributions estimated by the target network and the predictor network as an additional bonus when selecting the optimal action. We describe the implementation details of NDQFN and DPE and examine their performance on Atari 2600 games. We compare NDQFN and DPE with some baseline methods such as IQN and DLTV, and show that the combination of the two can consistently achieve the optimal performance especially in some hardexplored games such as Venture, Montezuma Revenge and Private Eye.

For the rest of this paper, we first go through some background knowledge of distributional RL in Section 2. Then in Sections 3, we introduce NDQFN, DPE and describe their implementation details. Section 4 presents the experiment results using the Atari benchmark, investigating the empirical performance of NDQFN, DPE and their combination by comparing with some baseline methods.

2 Background and RelatedWork

Following the standard reinforcement learning setting, the agent-environment interactions are modeled as a Markov Decision Process, or MDP, (X, A, R, P, γ) [Puterman, 2014]. X and A denote a finite set of states and a finite set of actions, respectively. R : X ×A → R is the reward function, and γ ∈ [0, 1) is a discounted factor. P : X ×A×X → [0, 1] is the transition kernel.

On the agent side, a stochastic policy π maps state x to a distribution over action a ∼ π(·|x) regardless of the time step t. The discounted cumulative rewards is denoted by a random variable Z π (x, a) = ∑∞t=0 γ tR(xt, at), where x0 = x, a0 = a, xt ∼ P(·|xt-1, at-1) and at ∼ π(·|xt). The objective is to find an optimal policy π∗ to maximize the expectation of Z π (x, a), EZ π (x, a), which is denoted by the state-action value function Qπ (x, a). A common way to obtain π∗ is to find the unique fixed point Q∗ = Qπ∗ of the Bellman optimality operator T [Bellman, 1966]:

The state-action value function Q can be approximated by a parameterized function Qθ (e.g. a neural network). Q - learning [Watkins, 1989] iteratively updates the network parameters by minimizing the squared temporal difference (TD) error

on a sampled transition (xt, at, rt, xt+1), collected by running an ϵ-greedy policy over Qθ.

Distributional RL algorithms, instead of directly estimating the mean EZ π (x, a), focus on the distribution Z π (x, a) to sufficiently capture the intrinsic randomness. The distributional Bellman operator, which has a similar structure with (1), is defined as follows,

where X′ ∼ P (·|x, a) and A′ ∼ π(·|X′ ). A D= B denotes the equality in probability laws.

3 Non-monotonicity in Distirbutional Reinforcement Learning

In Distributional RL studies, people usually pay attention to the quantile function F -1 Z (τ) = inf{z ∈ R : τ ≤ FZ(z)} of total return Z, which is the the inverse of the cumulative distribution function FZ(z) = Pr(Z < z). In practice, with limited observations at each quantile level, it is much likely that the obtained quantile estimates F -1 Z (τ) at multiple given locations {τ1, τ2, . . . , τN} for some stateaction pair (x, a) are non-monotonic as Figure 1 (a) illustrates. The key reason behind this phenomenon is that the quantile values at multiple quantile fractions are separately estimated without applying any global constraint to ensure the monotonicity. Ignoring the non-decreasing property of the learned quantile function leaves a theory-practice gap, which results in the potentially non-optimal action selection in practice. Although the crossing issue has been broadly studied by the statistics community [He, 1997; Dette and Volgushev, 2008], how to ensure the monotonicity of the approximated quantile function in DRL still remains challenging, especially to some quantile value based algorithms such as IQN and FQF, which do not focus on fixed quantile locations during the training process.

4 Non-Decreasing Quantile Function Network

To address the crossing issue, we introduce a novel NonDecreasing Quantile Function Network (NDQFN) with two main components: (1). an incremental structure to estimate the increase of quantile values between two pre-defined nearby supporting points pi-1 ∈ (0, 1) and pi ∈ (0, 1), i.e., ∆i(x, a; p) = F -1 Z (pi) - F -1 Z (pi-1), and subsequently obtain each F -1 Z (pi) for i ∈ {0, . . . ,N} as the cumulative sum of ∆i's; (2). a piece-wise linear function which connects the N + 1 supporting points to represent the quantile estimate F -1 Z(x,a) (τ) at any given fraction τ ∈ (0, 1). Figure 2 describes the whole architecture of NDQFN.

We first describe the incremental structure. Let p = {p0, · · · , pN} be a set of N + 1 supporting points, satisfying 0 < pi ≤ pi+1 < 1 for each i ∈ {1, 2, . . . , N -1}. Thus, Z(x, a) can be parameterized by a mixture of N Diracs, such that

where each θi denotes the quantile estimation at ˆ pi = pi+pi+1 2 , and θ = {θ0, . . . , θN-1}. Let ψ : X → Rd and φ : [0, 1] → Rd represent the embeddings of state x and quantile fraction τ, respectively. The baseline value ∆0(x, a; p) := F -1 Z(x,a) (p0) at p0 and the following N nonnegative increments ∆i(x, a; p) for i ∈ {1, · · · ,N} can be represented by

where i = 1, · · · ,N, f : Rd → R|A| and g : R3d → [0,∞)|A| are two functions to be learned. ω includes all the parameters of the network.

Then, we use Πp,∆ω to denote a projection operator that projects the quantile function onto a piece-wise linear function supported by p and the incremental structure ∆ω = {∆0,ω, · · · ,∆N,ω} above. For any quantile level τ ∈ (0, 1), its projection-based quantile estimate is given by

where Gi,∆ω could be regarded as a linear combination of quantile estimates at two nearby supporting points satisfying Gi,∆ω (τ, x, a; p) = { ∑ij=1 ∆j,ω(x, a; p)+ τ-pi pi+1-pi ∆i+1,ω (x, a; p)}I(pi ≤ τ < pi+1). By limiting the output range of ga in (3) to be [0,∞), the obtained N + 1 quantile estimates are non-decreasing. Under the NDQFN framework, the expected future return starting from (x, a), also known as the Q-function can be empirically approximated by

For notational simplicity, let Pτ,ω(x, a) = F -1,p,∆ω Z(x,a) (τ), given the support p. Following the idea of IQN, two random sets of quantile fractions τ = {τ1, · · · , τN1 }, τ ′ = {τ ′1 , · · · , τ ′N2 } are independently drawn from a uniform distribution U(0, 1) at each training iteration. In this case, for each i ∈ {1, . . . ,N1} and each j ∈ {1, . . . ,N2}, the corresponding temporal difference (TD) error [Dabney et al., 2018a] with n-step updates on (xt, at, rt, · · · , rt+n-1, xt+n) is computed as follows,

where ω and ω′ denote the online network and the target network, respectively. Thus, we can train the whole network by minimizing the Huber quantile regression loss [Huber, 1992] as follows,

where

κ is a pre-defined positive constant, and | · | denotes the absolute value of a scalar.

Remark: As shown by Figure 1 (c), the piece-wise linear structure ensures the monotonicity within the union of any two quantile sets τ and τ ′ from two different training iterations regardless of whether they are included in the p or not. However, as Figure 1 (b) demonstrates, directly applying a incremental structure similar to [Zhou et al., 2020] onto IQN without using the piece-wise linear approximation may result in the non-monotonicity of τ ∪τ ′ although each of their own monotonicity is not violated.

Now we discuss the implementation details of NQDFN. For the support p used in this work, we let pi = i/N for i ∈ {1, · · · ,N - 1}, p0 = 0.001 and pN = 0.999. NDQFN models the state embedding ψ(x) and the τ-embedding φ(τ) in the same way as IQN. The baseline function f in (3) consists of two fully-connected layers and a sigmoid activation function, which takes ψ(x) as the input and returns an unconstrained value for ∆0(x, a; p). The incremental function g shares the same structure with f but uses the ReLU activation instead to ensure the non-negativity of all the N increments ∆i(x, a; p)'s. Let ⊙ denote the element-wise product and g take ψ(x) ⊙ φ(pi) and φ(pi) - φ(pi-1) as input, such that

which to some extent captures the interactions among ψ(x), φ(pi) and φ(pi-1). Although ψ(x) ⊙ φ(pi) ⊙ φ(pi-1) may be another potential input form, the empirical results show that the combination of ψ(x) ⊙ φ(pi) and φ(pi) - φ(pi-1) is more preferred in practice considering its outperformance in Atari games. More details are provided in Section C of the Supplement file.

5 Exploration using Distributional Prediction Error

To further improve the learning efficiency of NDQFN, we introduce a novel exploration approach called distributional prediction error (DPE), motivated by the Random Network Distillation (RND) method [Burda et al., 2018], which can be extended to most existing DRL frameworks.

The proposed method, DPE, involves three networks, online network, target network and predictor network, which have the same architecture but different network parameters. The online network ω and the target network ω′ use the same initialization, while the predictor network ω∗ is separately initialized. To be more specific, the predictor network are trained on the sampled data using the quantile Huber loss defined as follows:

where δ∗ i,j = Pτ ∗ j ,ω′(xt, at)-P˜ τi,ω∗ (xt, at) and ρκτ (·) are defined in (6). On the other hand, following Hasselt[2010] and Pathak et al.[2019], the target network of DPE is periodically synchronized by the online network. DPE employs the 1Wasserstein metric W1(·, ·) to measure the distance between the two quantile distributions associated with the predictor network and the target network. And its empirical approximation based on NDQFN is defined as follows,

Since the minimum of L(xt, at) is a contraction of the 1Wasserstein distance [Bellemare et al., 2017], the prediction error would be higher for an unobserved state-action pair than for a frequently visited one. Therefore, i(xt, at) can be treated as the exploration bonus when selecting the optimal action, which encourages the agent to explore unknown states. With i(xt, at) being the exploration bonus, the optimal action is determined by

where ct is the bonus rate. As might be expected, a more reasonable quantile estimate would highly enhance the exploration efficiency according to our exploration setting. To more clearly demonstrate this point, we compare the performance of DPE based on NDQFN and IQN in Section 6.3.

On the other hand, we compared the exploration efficiency of distributional prediction error and value-based prediction error in Section D of supplement, where value-based prediction error is defined as i′ (xt, at) = |Qω∗(xt, at) - Qω′(xt, at)|. As might be expected, DPE enhances the exploration efficiency by utilizing more distributional information.

6 Experiments

In this section, we evaluate the empirical performance of the proposed NDQFN together with DPE on 55 Atari games using the Arcade Learning Environment (ALE) [Bellemare et al., 2013]. The baseline algorithms we compare include IQN [Dabney et al., 2018a], QR-DQN [Dabney et al., 2018b], C51 [Bellemare et al., 2017], prioritized experience replay [Schaul et al., 2015] and Rainbow [Hessel et al., 2018]. The whole architecture of our method is implemented by using the IQN baseline in the Dopamine framework [Castro et al., 2018]. Note that the whole method could be built upon other quantile value based baselines such as FQF.

Our hyper-parameter setting is aligned with IQN for fair comparison. Furthermore, we incorporate n-step updates [Sutton, 1988], double Q learning [Hasselt, 2010], quantile regression [Koenker, 2005] and distributional Bellman update [Dabney et al., 2018b] into the training for all compared methods. The size of p and the number of sampled quantile levels ,τ and τ ′ , are 32. We use ϵ-greedy policy with ϵ = 0.01 for training. At the evaluation stage, we test the agent for every 0.125 million frames with ϵ = 0.001.

For the DPE exploration, we fix ct = 1 without decay (which is the optimal setting based on experiment results). As a comparison, we extend DLTV [Mavrin et al., 2019], an exploration approach designed for DRL algorithms based on fixed quantile locations such as QR-DQN, to quantile value based methods by modifying the original left truncated variance used in their paper to σ2+ = ∫ 1 12 [F -1 Z (τ)-F -1 Z ( 1 2 )]2dτ . All training curves in this paper are smoothed with a moving average of 10 to improve the readability. With the same hyper-parameter setting, NDQFN with DPE is about 25% slower than IQN at training stage due to the added predictor network.

6.1 Full Atari Results

In this part, we evaluate the performance of NDQFN and the whole network, which incorporates NDQFN with DPE in all the 55 Atari games. Following the IQN setting [Dabney et al., 2018a], all evaluations start with 30 non-optimal actions to align with previous distributional RL works. 'IQN' in the table tries to reproduce the results presented in the original IQN paper, while 'IQN*' corresponds to the improved version of IQN by using n-step updates and some other tricks mentioned above.

Table 1 presents the mean and median human normalized scores by the seven compared methods across 55 Atari games, which shows that NDQFN with DPE significantly outperforms the IQN baseline and other state-of-the-art algorithms such as Rainbow.

The detailed comparison results including the improvement of DPE over ϵ-greedy, NDQFN+DPE over IQN, the training curves of NDQFN and the whole network for all the 55 games and the raw scores are provided in Sections E, F and G of the supplement.

6.2 NDQFN vs IQN

In this part, we compare NDQFN with its baseline IQN [Dabney et al., 2018a] to verify the improvement it achieves by employing the non-decreasing structure. The two methods are fairly compared without employing any additional exploration strategy.

Figure 3 visualizes the training curves of eleven selected easy explored Atari games. NDQFN significantly outperforms IQN in most cases, including Berzerk, Hero, Ms. Pac-Man, Qbert, Up and Down, and Video Pinball. Although NDQFN achieves similar final scores with IQN in Zaxxon and Tennis, it learns much faster. This tells that the crossing issue, which usually happens in the early training stage with limited training samples, can be sufficiently addressed by the incremental design in NDQFN to ensure the monotonicity of the quantile estimates. In particular, Figure 4 visualizes the relative improvement of NDQFN over IQN ((NDQFNrandom)/(IQN-random)-1) in testing score for all the 55 Atari games. We observe that NDQFN consistently outperforms the IQN baseline in most situations, especially for some hard explored games with extremely sparse reward spaces, such as Private Eye and Montezuma Revenge.

6.3 Effects of NDQFN on DPE

In this section, we aim to show whether the NDQFN design can help improve the exploration efficiency of DPE compared to using the IQN baseline. We pick three hard explored games, Montezuma Revenge, Private Eye, Venture, and three easy explored games, Ms. Pac-Man, River Raid and Up and Down. The main results are summarized in Figure 5.

As Figure 5 illustrates, DPE is an effective exploration method which highly increases the training perfogrmance of both IQN and NDQFN in three hard explored games. In the three easy-explored games, DPE still performs slightly better than the baseline methods, which to some extent demonstrates the stability and robustness of DPE. On the other hand, we find that NDQFN + DPE significantly outperforms IQN + DPE especially in the hard-explored games. This agrees with the conclusion we make in Section 5 that NDQFN obtains a more reasonable quantile estimate by adding the nondecreasing constraint which helps to increase the exploration efficiency.

6.4 DPE VS DLTV

We do some more experiments to show the advantage of DPE over DLTV when applying both of the two exploration approaches to quantile value based DRL algorithms. To be specific, we evaluate the performance of NDQFN+DPE and NDQFN+DLTV on the six games examined in the previous subsection. Figure 6 shows that DPE achieves a much better performance than DLTV especially in the three hard explored games. DLTV does not achieve significant performance gain over the baseline NDQFN without doing exploration.

This result tells that DLTV, which computes the exploration bonus based on some discrete quantile locations, performs extremely unstable for quantile value based methods such as IQN and NDQFN since the quantile locations used to train the model are changed each time. As a contrast, DPE focuses on the entire distribution of the quantile function, which relies less on the selection of the quantile fractions, and thus performs better in this case.

7 Conclusion and Discussion

In this paper, we make some attempts to address two important issues people care about for distributional reinforcement learning studies. We first propose a Non-decreasing Quantile Function Network (NDQFN) to ensure the validity of the learned quantile function by adding the monotonicity constraint to the quantile estimates at different locations via a piece-wise linear incremental structure. We also introduce DPE, a general exploration method for all kinds of distributional RL frameworks, which utilizes the information of the entire distribution and performs much more efficiently than the existing methods. Experiment results on more than 50+ Atari games illustrate NDQFN can more precisely model the quantile distribution than the baseline methods. On the other hand, DPE perform better in both hard-explored Atari games and easy ones than DLTV. Some ablation studies show that the combination of NDQFN with DPE achieves the best performance than all the others.

There are some questions remain unsolved. First, since lots of non-negative functions are available, is Relu a good choice for function g? Second, will a better approximated distribution affect agent's policy? If so, how does it affect the training process as NDQFN can provide a potentially better distribution approximation. More generally, how to analytically compare the optimal policies NDQFN and IQN converge to. Third, how much does the periodically synchronized target network affect the efficiency of the DPE exploration.

For the future work, we first want to apply NDQFN to other quantile value based methods such as FQF to see if the incremental structure can help improve the performance of all the existing DRL algorithms. Second, as there exists some other kind of DRL methods which naturally ensures the monotonicity of the learned quantile function [Yue et al., 2020], it may be interesting to integrate our method with these methods to obtain an improved empirical performance and training efficiency.

Acknowledgments

The paper is supported by National Natural Science Foundation of China (11601313, 12001356), Program for Innovative Research Team of SUFE, Shanghai Sailing Program (20YF1412300), and Fundamental Research Funds for the Central Universities.