Distributional Reinforcement Learning for Multi-Dimensional Reward Functions
Authors: Pushi Zhang
Xiaoyu Chen
Li Zhao
Wei Xiong
Tao Qin
Tie-Yan Liu
Venue: NeurIPS 2021 Poster (2021)
Keywords: Reinforcement Learning
Deep Reinforcement Learning
Distributional Reinforcement Learning

Abstract: A growing trend for value-based reinforcement learning (RL) algorithms is to capture more information than scalar value functions in the value network. One of the most well-known methods in this branch is distributional RL, which models return distribution instead of scalar value. In another line of work, hybrid reward architectures (HRA) in RL have studied to model source-specific value functions for each source of reward, which is also shown to be beneficial in performance. To fully inherit the benefits of distributional RL and hybrid reward architectures, we introduce Multi-Dimensional Distributional DQN (MD3QN), which extends distributional RL to model the joint return distribution from multiple reward sources. As a by-product of joint distribution modeling, MD3QN can capture not only the randomness in returns for each source of reward, but also the rich reward correlation between the randomness of different sources. We prove the convergence for the joint distributional Bellman operator and build our empirical algorithm by minimizing the Maximum Mean Discrepancy between joint return distribution and its Bellman target. In experiments, our method accurately models the joint return distribution in environments with richly correlated reward functions, and outperforms previous RL methods utilizing multi-dimensional reward functions in the control setting.

Distributional Reinforcement Learning for Multi-Dimensional Reward Functions
Pushi Zhang*t
Xiaoyu Chen*t Tsinghua University
Li Zhao'
Tsinghua University
Microsoft Research Asia lizo@microsoft com
zpschang@gmail.com com
chen foon-xy21000 @mails tsinghua, edu. cn
Wei Xiong
Tao Qin
Tie- Yan Liu
The Hong Kong University of Science and Technology wxiongae&cconect. ust.hk
Microsoft Research Asia taoqingmicrosoffe
Microsoft Research Asia sylliuficroooot. com
Abstract
A growing trend for value-based reinforcement learning (RL) algorithms is to capture more information than scalar value functions the value network. One of the most well-known methods in this branch is distributional RL, which models return distribution instead of scalar value. In another line of work, hybrid reward architectures (HRA) in RL have studied to model source-specific value functions for each source of reward, which also shown to beneficial in performance. To fully inherit the benefits of distributional RL and hybrid reward architectures, we introduce Multi-Dinn Dimensional Distributional DQN (MD3QN), which extends distributional RL to model the joint return distribution from multiple reward sources. As a by-product of joint distribution modeling, MD3QN can capture not only the randomness in returns for each source of reward, but also the rich reward correlation between the randomness of different sources. We prove the convergence for the joint distributional Bellman operator and build our empirical algorithm by minimizing the Maximum Mean Discrepancy between joint return distribution and its Bellman target. In experiments, our method accurately models the joint return distribution in environments with richly correlated reward functions, and outperforms previous RL methods utilizing multi-dimensional reward functions in the control setting.
Introduction
Making value network capture more information than scalar value functions is a growing trend in value-based reinforcement learning, which helps the agent gain more knowledge about the environ- ment and has great potentials to improve the sample efficiency of RL agents. In the early stage of deep reinforeement learning, DQN (Mnih al. 2013) uses the scalar output of the neural network to represent value functions. As value- based RL algorithms evolve, distributional RI. algorithms start to use neural networks approximate return distributions for each state-action pair, and take action based the expectation of return distributions. By capturing the randomness in return as auxiliary tasks, distributional RL agents can gain more knowledge about the environment and learn better representations to avoid state aliasing (Bellemare al., 2017). Distributional RL algorithms including C51 (Bellemare et al. (017). QR-DQN (Dabney 20T8b), IQN (Dabney al. 2018a)
'Equal contribution.
Work done during an internship at Microsoft Research Asia.
I Corresponding author.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

FQF YYng et al. 2019) and MMDQN Nguyen et al. 2020) achieve substantial performance gain compared to DQN.
In another line of work, HRA Van Seijen et al. 2017 and RD (Lin et al. 2020) consider the setting where multiple sources of reward exist in the environment and modify the value network to model source-specific value function for each source of reward. In these works, the outputs of value networks can be interpreted as multiple source-specific value functions, and the agent takes action based on the sum of all source-specific value functions. Similar to return distribution, estimating the source-specific value functions can be seen as auxiliary tasks, which serve as additional supervision for how the total reward is composed and enable the agent to learn better representations, Several previous works support that it beneficial to model source-specific value functions (Boutilier et al., 1995: Sutton et al. 2011a Van Seijen et al. 2017 [in et al. (020).
Towards providing more supervision signals and enabling agents to gain more knowledge about the environment, we propose to capture the correlated randomness in source-specific returns. Specifically, we consider the source-specific returns from all sources of rewards as multi-dimensional random variable, and capture its joint distribution to model the randomness of returns from different sources. This provides an informative learning target for our agent. The framework is general and can be extended to capture the correlated randomness of other types of random variables than rewards. For example, we can capture the correlation between achieving different goals in goal-conditioned reinforcement learning (Schaul et al.) 2015), or visiting different states in successor representa- tion (Kulkarni et al. 2016). In this paper, we focus on the method for learning joint return distribution of given source-specific rewards and leave the extension to more general settings for future work.
Following existing works distributional RL., we study the convergence of the Bellman operator and propose an empirical algorithm to approximate the Bellman operator. First, we define the joint distributional Bellman operator, and prove its convergence under the Wasserstein metric. To derive an empirical algorithm, our proposed method (MD3QN) approximates the joint distributional Bellman operator by minimizing the Maximum Mean Discrepancy (MMD) loss over joint return distributions and its Bellman target. MMD holds desirable properties that, it is a metric over joint distribution and its square can be unbiasedly optimized with batch samples. This enables our algorithm to approximate the Bellman operator accurately when the number of samples goes to infinity and the loss minimized to zero. In experiments on Atari games and other environments with pixel inputs, our method accurately models the multi multi-dimensional joint distribution from multiple sources of reward. Moreover, our algorithm outperforms previous work HRA which also separately models multiple sources of reward on Atari game environments.
Our contributions can be summarized as follows:
We propose distributional RL algorithm, MD3QN, that extends distributional RL algo- rithms to model the joint return distribution from multiple sources of reward.
We establish convergence results for the joint distributional Bellman operator, and our proposed algorithm MD3QN approximates this Bellman operator by minimizing MMD loss over joint return distribution and its Bellman target.
Empirically, our method outperforms previous RL algorithms utilizing multiple sources of reward, and accurately models the joint return distribution from all sources of rewards.
2 Background
2.1 Notations and Problem Setting
Since work considers multiple sources of reward exist in the environment, our problem setting is slightly different from the traditional RL setting in reward function. We consider a Markov Decision Process with multiple sources of reward defined by (S, A P. R. po; 7), where S is the of states, A is the finite set of actions, P : S Sx.A A -> P(S) denotes the transition probability, R S x A - P(R denotes the reward function for total of N sources of reward, Po PP(S) denotes the distribution of initial state So. 7 E (0. 1) is the discount factor. Given a policy π : S + P(A), a trajectory is generated by Sy ~ Po. ag ~ π (st), $1+1 ~ P(st. and r. = [rr,1. Ft, .... Ft,N] ~ R(st, az).
We use Fi- EN Te. to denote the total reward received at time t. The goal for reinforcement

---

learning algorithms is to find the optimal policy и which maximizes the expected total return from all sources, given by (TT) E -===00 Σ Tt,n)]-
Next we describe value-based reinforcement learning algorithms general framework. In DQN, the value network (ss, a; f) captures the scalar value function, where @ is the parameters of the value network. fg the initial value of 8. i-th iteration, the training objective for 0g i
[(Q I(O+)) B.. (l 04; oz;0,) )=)
where y; rt + 7 max (';
(2)
and Sty ats Tri S are sampled from the replay memory.
2.2 Distributional RL
In distributional RI. algorithms such as C51, IQN and MMD-DQN, ju(8, 0) captures the return distribution for each state action pair (s, In i-th iteration, the training objective for is
L(O,) at, [d(u(s), az 0s),
(3)
where 1/: = (fren) # p(s++1, d'; 01-1), a' arg max Ez (:0
(4)
Here (x) - Given a probability distribution & E P(RR), fav E P(R) is the pushforward measure defined by f (A)) for all Borel sets C I and a measurable function f: + R (Rowland et al. 2018). d is some distributional metric different in each method. In C51 d KL divergence, MMDQN d is Maximum Mean Discrepancy, while in QR-DQN and IQN d is Wasserstein metric, which is optimized approximately via quantile regression loss.
2.3 Hybrid Reward Architecture
HRA (Van Seijen et al. 2017) proposes the hybrid architecture to separately model the value function for each source of reward. With multiple sources of reward, we use Q (I=V)SXA and Qn SxA to denote the vectored value function and the n-th value function respectively.
The loss used by HRA is given by:
L(0;) = Es (I (1 (st, ax: 02) - y:
(5)
where Vi,n = Ti, + Q d' 1,,1), arg max (s41,, 0-1))
(6)
RD2 (Lin al. 2020) uses the similar hybrid architecture, which also learns the value function separately for each source of reward, but with slight differences with HRA in how the next action a' computed. The loss used by RD is given by
L(O.) Es, ae. (1 ar; (,; vill2).
where y; = r 7 Q(1+111 d';)))))''' arg max Σ Q1(
(8)
x=1
3 Distributional Reinforcement Learning for Multi-Dimensional Reward Functions
In this paper, we propose capture the correlated randomness from multiple sources of reward, forcing the agent to gain more knowledge about the environment and learn better representations. Specifically, we consider the joint distribution of returns from different sources of reward. First, we introduce the joint distributional Bellman operator and establish its convergence. Second, we
'In practical implementations. the iterative process is approximated by substituting 0.-1 the parameter of the target network which tracks the exponential moving average the ouline network parameters 0s. In equation [B) ys is an N N-dimensional vector. We use the bold font to denote that the variable is for the multi-di molti-dimensional reward function.
3

---

introduce empirical algorithm, which approximates the joint distributional Bellman operator through stochastic optimization with deep neural networks. Inspired by MMDQN (Nguyen al.) 2020), we introduce temporal difference loss based on Maximum Mean Discrepancy over joint return distributions. Finally, we present full implementation details of our method, including network architectures and algorithms.
3.1 Definitions and Settings
We consider the joint return under policy и as random vector (*(s, composed of N random variables Ted ("sss a) = (Zf (s) (s,
Do
zT a) Σ 1rr
1=0
where 50 5, a0 a, Ti R(- at) se+1 P(-|se, az), ae+1 ~ ((1s(+1))
(10)
Here (8, a) denotes i-th dimension (s, a), which is the random variable of the i-th source of discounted return. Here the random variables of discounted return in different sources can be correlated. We denote the distribution Z as μx Law (Z*), and the distribution of i-th random variable z 1 Law (2,), where μt' E P(R and nu E PxA) We let E PPRRNSSX to denote arbitrary joint distribution over all state-action pairs, and the goal of our algorithm is let the joint distribution μ to be close to the joint distribution μ1 as possible in policy evaluation setting, i.e. to model the real joint return distribution.
To measure how close N-dimensional joint distributions are, we adopt the Wasserstein metric W,. For two joint distributions vI, U2 E P(RR"), the p- Wasserstein metric Wp v2) on Euclidean distance d is given by:
W Wp U2) = ( inf
/
1/p
d(r,u)² v) xd" y
(11)
(Er((e))))) XB#
where [(v), 12) is the collection of all distributions with marginal distributions vs and the first and second N random variables respectively. The Wasserstein distance can be interpreted as the minimum moving distance to move all the mass from distribution 1/1 to distribution 1/2
-rr-suuu----sssetteennn We further define the supremum-pe Wasserstein distance on P P PSSA by
(14, U2) = sup W (m (s. 0), (2)) s.a
(12)
where μl, f2 E P(RN)SSXA
We will use dp to establish the convergence the joint distributional Bellman operator.
3.2 Convergence of the Joint Distributional Bellman Operator
We define the joint distributional Bellman evaluation operator T as
T p((0, on) D Ja
(13)
where (x) r ya, where æ,r N and f is the pushforward measure which N-dimensional extension of the pushforward measure defined in Rowland et al. (018). The below theorem shows that T c-contractionn operator on dp-
Theorem For two joint distributions μ4 and μ2, we have
d,(TT p1, T (2) S (d1.
(14)
The proof for Theorem is provided in Appendix A. which can be briefly described as follows. In Lemma 3, we prove that if the Wasserstein distance of two joint distributions is C, then after applying the same linear transformation to these two distributions with scale factor of y, the Wasserstein

---

distance is at most yC. In Lemma 4, we prove that the Wasserstein distance of two joint distributions under the same conditional random variable no greater than the maximum Wasserstein distance over all possible conditions A = a. By Lemma 3 and Lemma 4, we are able to prove the contraction results in Theorem
Following the result of Theorem 1, we consider the following scenario: initially we have joint distribution 140 P(R over all state-action pairs, and by iteratively applying the Bellman evaluation operator, we get Mi+1 ' p4. According to Banach's fixed point theorem, operator 7 has unique fixed point, which is μ² by definition. The distance between μe and ju" decays as increases. We have the following corollary:
Corollary If Pi+1 TT then
o0, p4 pa"
We also provide contraction proof on the expectation of the optimality operator. The joint distribu- tional Bellman optimality operator T defined as
Tp(si, at) -2 . Jex (Imn) d') R')R(drelsss 0r) P(dsssssss
(15)
where a' E arg Zxp Σ Zn and Zn is the n-th element in Z.
(16)
Theorem For two joint distributions fe and p2z, we have
(E (B S <1l(B=)m (E
(17)
where the operator TEx- is defined by (ES) (a(s, a) E= Ex= Zn for any (s, (s, and Zn is the ni-hh element in Z. The operator Ey converts the joint distribution all state-action pairs to the corresponding total expected returns over all state-action pairs.
Corollary If M441 pbg then as i co, Ez~pi(ssa) Em Zn (", for all (8,a)
Corollary 2 can be interpreted follows: as we iteratively apply the joint distributional Bellman optimality operator, the expected total return of the joint distribution will converge to optimal value function. We refer the readers Appendix A.1 for detailed proofs for all the lemmas, theorems, and corollaries.
Next, establish the empirical algorithm which simulates the iterative process in Corollary I and Corollary with a neural network approximating the joint distribution μ4; for all state-action pairs,
3.3 Optimizing MMD as Approximation for Bellman Operator
In practical algorithms, the joint distribution μ4; [s; a) is represented by deep neural network pa(s, a; 0.) which outputs joint distribution for each state-action pair with parameter 0,. In the (i+1)-th iteration, we can only adapt the parameter 0.+1 of the model, and cannot directly apply the Bellman operator as in the tabular case. Moreover, we only have sampled transitions, rather the environment dynamics model P. It is desirable to choose a loss that is compatible with stochastic optimization, an approximation for the Bellman operator.
We achieve this by minimizing the Maximum Mean Discrepancy (MMD) between joint distribution μi+1 and T u,, which defined as the equation below.
MMD² (p. qi k) = Ep,x'>pha,, x') - ((0-psngk(sr222 y) + kyy.
(18)
where k(.,-) is some kernel function and MMD defined upon k and each pair of random variable (xr, x"), 3, (I,, y), (v.y/ y') are independent. In scenario,
MMD (M a). (4i(s,a), k) as the temporal difference loss. The reason why we choose MMD as the objective function is that MMD is both metric distribution and easy to optimize with sampled transitions. The original paper of MMD (Gretton et al. (0122) proves that MMD is metric on distribution when kernel is characteristic, from which we can prove that
pit1 TH + V(s,a) S ( au) k)
0,
(19)
5

---

when the kernel k is characteristic. If the MMD metric is optimized to zero and the appropriate kernel selected, Mi+1 is the exact result of applying the joint distributional Bellman operator.
In the algorithmic aspect, [8, a) serves as the target distribution, [4441 (s, a) serves as the pre- diction from online network, and we use MMD' AMMD(uu++(0 a), T [ a); k) as temporal difference loss to match the online prediction and target distribution. Next, establish the method to optimize MMD* (s, a). T [ a); k) with transition samples. Note that we use pu(s, a; 8s) to represent pe(s, a), so the temporal difference loss can also be written as MMD2 (u(s, a; 0c+1), ), 7 p1 a; 0,); k). We use equation 18 to estimate the gradient of this squared MMD loss with respect to parameter 0.41. Specifically, the first term in equation (18) can be unbiasedly estimated by drawing multiple independent samples from the online network po(s, a; 8(+1); the second term of equation (8)) can be unbiasedly estimated by drawing independent samples from the online network pl(s, a: 8,+1) and from the target distribution T Tu(s, a; 0, | Batch of samples from the target joint distribution T a; 0,) can be collected from only one transition sample' (s, a, r, s'), and Lemma 6 in Appendix A.1 provides detailed proofs. For the third term in equation (8) it has no gradient with respect to the parameter 0 and we can safely ignore it during optimization. In summary, the gradient of the temporal difference loss MMD² can be estimated without bias by transition samples. Algorithm 1 summarizes the above analysis and shows how our method computes the gradient estimates of the temporal difference loss MMD?
Algorithm 1 Gradient estimation of MMD loss by transition samples
Require: Number of samples M, kernel k; discount factor y E (0. 1)
Require: Joint distribution network fa(s, a: 0)
Input:
Transition sample (s, a, r, s')
Input:
Online network parameter 0, target network parameter O'
Output: Gradient estimation of MMD with respect 0
1: a' < arg max, Em= Em 1 (Zm)r. where ZI:M i.idd uls', 0*)
2: Sample ZI:M iidd. a)) (samples from online network)
a(ss',';'0) 4: Y 4 r + + for every 1 S iis M (samples from target distribution)
3: Sample Zrort i.ddd.
5: MMD2 + EEsisn Σ Eisignayi [k(Z, Zj) 2k Z Yj) + k(Y., X,)]
6: return VgMMD
3.4 Network Architecture and Implementation Details
For given state-action pair (s, a), we use the value network of DQN (Mnih et al. 2013) to compute the low-dimensional (512d) embedding of state 8, and replace the final layer of DQN network to output a vector with M x N x |AA] dimensions, representing M samples from the joint return distribution from sources of reward for every action.
We follow the kernel selection in MMDQN (Nguyen et al. 2020 and apply the Gaussian kernel function with mixed b
M >
Mx,x') =
ki (x,x'), x'),
(20)
where kw (x,, x²)
(21)
Each wi can be seen as the square of bandwidth in a Gaussian kernel. We choose the value of squared bandwidth w; with diverse ranges, to make sure that all scales of reward don't 't suffer from gradient vanishing. In Appendix A.3.4 and A.3.5, we test the performance of different network architectures and different bandwidth sets to find the optimal network architecture and bandwidth configurations.
"We only require each sample from the online network 11s, a; 0++1) be independent of every sample from the target distribution T (uss, a; 0:), and multiple samples from the target distribution T ju(s, a; 0.) can be dependent each other.
"Gaussian kernel is proved to be characteristic (Grettnn et al. 2012 so the MMD distance with Gaussian kernel is metric. It follows that when use Gaussian kernels with mixed bandwidth, equation (99 still holds.

---

4 Related Work
Distributional RL algorithms propose to model the entire distribution, rather than the expectation, of the random variable return. C51 (Bellemare et al.) 2017) is the first distributional RL algorithm, and establishes the convergence of distributional RL algorithms based on the contraction property of the distributional Bellman operator. To approximate such an operator with sample transitions, different distributional algorithms employ different losses over distribution, along with different parameterization for distribution. C51 1 uses KL-divergence loss and categorical distribution as parameterization, QR-DQN (Dabney et al. 2018b) takes quantile regression as surrogate loss for the Wasserstein metric, and approximates set of quantile values with fixed probabilities. IQN (Dabney and FQF (Yang et al. 2019) extend QR-DQN to learn with sampled probabilities and self-adjusted probabilities. The current SOTA method in distributional RL is MMDQN (Nguyen et al. 2020), which takes MMD loss with a non- non-parametric approach by modeling deterministic samples from return distribution. Our method extends the theoretical work of C51, and the empirical algorithm of MMDQN to multi-dimensional returns. While distributional RL algorithms focus on capturing the randomness in return, our method first proposes capture the correlation between different randomness from different reward sources.
et al. 8a)
HRA (Van Seijen et al. 2017 proposes hybrid architecture to separately model the value functions for different sources of rewards. Their work provides empirical justification that learning with hybrid rewards can improve sample efficiency. HRA is built upon the Horde architecture (Sutton al., 2011b), which trains a separate general value function (GVF) for each pseudo-reward function. The Horde architecture uses a large number of GVFs to model general knowledge about the environment. Previous works on reward decomposition also demonstrate that it is beneficial to learn with multiple reward functions (Lin al., 2020). Based on HRA, our method further considers the correlated randomness in hybrid-source returns with a joint distribution model.
Bellman GAN (Freirich et al. 2019) shows that distributional RL can be used for multi- dimeiiiaiinni rewards. Compared to this work, we additionally provide the theoretical convergence results for the joint distributional Bellman operator, algorithm where the gradient of objective function (MMD) can be unbiasedly estimated, and experimental results that directly compare the modeled joint distribution to the true distribution, demonstrating the accuracy of our method under diverse reward correlations. DRDRL. (Lin et al.) (019) focuses on decomposing multiple types of rewards from one source reward, while in our paper, we are directly provided with multiple types of rewards as inputs, and model the joint distribution of different sources of rewards.
By modeling the joint return distribution, we provide an informative target to learn, which can be seen as auxiliary tasks. Previous works in RL have constructed various auxiliary tasks to learn better representations, such methods based temporal structures {Aytar et al. 2018) and local spatial structures (Anand et al.) 2019). Compared with these methods, our method does not entirely focus learning better state representations. The learned joint distribution is also beneficial to risk-sensitive tasks (Zhang et al. 2020). Besides, our method for learning the joint distribution of multi-dimensions random variables is general, and can be further combined with goal-conditioned
RL (Schaul et al. 2015) or successor representation (Kulkarni et al. 2016) to capture correlated randomness in achieving different goals or visiting different states.
Experimental Results
In our experiment provide empirical results to answer the following questions:
On policy evaluation settings, can MD3QN accurately model the joint distribution of multiple sources of reward?
On policy optimization settings, can MD3QN learn better policy compared to HRA and distributional RL algorithm on environments with multiple sources of reward?
To answer the first question, we design a maze environment with rich way to generate correlated sources of rewards, and compare the joint distribution predicted by our method with the true joint distribution ,². The detailed setting of the environment and results is shown in Section 5.
*Our code for the experiments is available in the supplementary material.
7

---

144
(a)
(b)
(c)
rewards Figure Observation of initial states in maze environments. (a): maze environment "maze-exclusive" with two exclusive sources of rewards, (b): maze environment "maze-identical" with two positively correlated sources of rewards. (c): maze environment "maze-multireward" with four correlated
sources of rewards.
To answer the second question, we use several Atari games with multiple sources of reward, and provide training curve results of MD3QN compared to previous work HRA which deals with multiple sources of reward. The detailed setting of the environment and results is shown in Section 5.2.
5.1 Modeling Joint Return Distribution for Multi-Dimensional Reward Functions
In this experiment, we validate the capacity MD3QN to model the joint return distribution for multiple sources of rewards the policy evaluation setting with pixel inputs. The experiment is performed on the maze environments with rich reward correlation signals described as follows.
Figure [ iillustrates three maze environments with differrnnt reward correlation properties. The agent (represented by the triangle) is initially located at a specific position in the maze, and the policy π is a fixed policy to uniformly choose one direction random and try to move. If the agent is blocked by the walls that direction or the agent trying to return to a previous position, the move has no effect, otherwise, the agent moves in that direction for one block. Each color of the square in the maze represents one source reward. When the agent reaches the position with reward, it receives a reward of the source aligned with the color of the squard The rules described above general enough to enable the design of diversely correlated sources of rewards detailed below.
Figure I(a) shows maze environment with two exclusive sources of rewards. Once the agent obtains reward from one side, cannot obtaín the reward from the other side. Figure [(b] shows maze environment with two positively correlated sources of rewards, where the agent either gets no reward or gets both of the red and green rewards. Figure [(c) shows a maze environment with four correlated sources of rewards. In these three environments, the agent needs to capture the negative correlation, positive correlation, and the complex correlation in respectively to precisely model the joint return distribution.
each of the three mazes, we train the MD3QN agent for iterations (1.25M frames) to model the joint distribution of all sources of reward, and compare the prediction of joint distribution by the model to the samples from the true distribution μ on the initial state of the maze. The results are shown in Figure 2 where each point in the figure is sample representing one possibility of future discounted return in each source of reward. The result shows that MD3QN accurately models the joint return distribution from different sources of reward.
In Appendix 1. visualize the modeled joint distribution throughout the training process, where the joint distribution modeled by MD3QN gradually match the true distribution.
It is also worth noting that the observation for the agent is based on pixels, computed by downsampling the image in Figure to the size of 84 x 84, which requires the agent to extract information on high-dimensional inputs.
Coolred walls can only be passed one way, where yellow walls can be passed from the top or from left, and walls can be passed from bottom or from the right.
More detailed description of how reward set in these three environments is provided in Appendix A.2.

---

(a)
(b)
(c)
(d)
Return in red
Return blue
Return red
Return
1.
Return in green
Return in green
Return in orange
Return in green
MD3QN
ground truth
Figure 2: Comparison between samples from joint distribution by MD3QN and samples from the true distribution μ² in three maze environments. The number of samples is set to be 200. (a): the result of "mazce-xclusivee" environment, (b): the result of "maze-identical" environment. (c): the result of muazzemultireward" environment of orange and blue reward sources. (d): the result of "maze- muazze-uutireward"" environment of green and red reward sources.
5.2 Performance on Atari Games
In this experiment, compare MD3QN algorithm with Hybrid Reward Architecture (HRA) Van Sei jen et al., 2017) and MMDQN (Nguyen et al., 2020) on Atari games from Arcade Learning En- vironment (Bellemare al. 2013). Our implementation of MD3QN is built upon the Dopamine framework (Castro et al. 2018). The hyper-parameter settings used by MD3QN and the environment settings are detailed in Appendix A.2.
We use six Atari games with multiple sources of rewards: AirRaid, Asteroids, Gopher, MsPacman, UpNDown, and Pong. In all of these six games, the primitive rewards can be decomposed into multiple sources of rewards as follows:
In AirRaid and Asteroids, the agent gets different value of reward for killing different types monsters.
In Gopher, the agent gets +80 reward for killing monster and gets +20 reward for removing holes on the ground.
In MsPacman, the agent gets [+200, +400, +800, + +1600) reward for killing a monster and gets +10 reward for eating beans.
In UpNDown, the agent gets +400 reward for killing an enemy car, +100 reward for reaching a flag, and +10 reward for being alive.
In Pong, the agent gets +1 reward for winning a round, and gets -1 reward for losing a round.
We decompose the primitive rewards into munti-ddimensionll rewards according to the above re- ward structure, and meanwhile keeping the total reward not changed. The details for the reward decomposition method is provided Appendix A.2. The training curves of MD3QN compared to Van Seijen et al. 2017) and MMDQN (Nguyen et al.) 2020) are shown in Figure 3
HRA
Figure shows that MD3QN does not outperform MMDQN in performance under several Atari games. We suspect the reason why MD3QN cannot outperform MMDQN is that the performance of Atari games might not be related the joint distribution. We look forward to further studies on how the additional information captured in distributional RI. algorithms helps to improve the final performance, better interpret the experimental results. Despite that, with the distributional information, MD3QN outperforms or matches the performance of HRA under all six games.
On Atari games, also show the joint distributions modeled by MD3QN, which is detailed in Appendix A.3.2. Overall, both the reward correlation and the reward scale in the return distribution reasonably captured by MD3QN. The joint distribution captured by MD3QN show different and reasonable reward correlation across different Atari games.
To show the necessity of modeling the joint distribution, we add another experiment in Appendix A.3.3 showing that only modeling marginal distributions may fail in environments with multiple correlated constraints, from theory and experiment results.

---

AirRaid
Asteroids
Gopher
Average Score
12400
10:00
200
Be Dr 358 Training Frames (millions)
2 10
up Training Frames (millions)
Training Frames (millions)
MsPacman
Pong
UpNDown
Average Score
10030
18080
2. Training Frames (millions)
Training Frames (millions)
Training Frames (millions)
HRA
MD3QN
MMDQN
Figure 3: Performance of MD3QN Atari games compared to HRA and MMDQN. The shaded
areas show the standard deviation across 3 seeds.
6 Conclusions
In this work, we have proposed Multi-Dimension Distributional DQN (MD3QN), distributional RL method that learns dimensioalnii joint return distribution for multiple sources of rewards. The effectiveness of our method is verified on pixel-input environments in terms of both the quality of modeled joint distribution, and the final performance of learnt policies.
In the future, possible to extend our framework to model the correlated randomness in achieving different goals with goal-condi RI, and in visiting different states with successor representation. We will leverage such informative modeling of the environment in terms of goals and successor states develop novel RL algorithms in our future work.
Acknowledgments and Disclosure of Funding
We thank the editor and each reviewer for their constructive comments and suggestions. We thank Tengyu Ma for the valuable suggestion on the MD3QN algorithm.
Funding in direct support of this work: Microsoft.
References
Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M.-A,, and Hjelm, R. D. (2019). Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, pages 8769-8788
Aytar, Y., Pfaff, T., Budden, D., Paine, T., Wang, Z., and de Freitas, N. (2018). Playing hard exploration games by watching youtube. In Advances in Neural Information Processing Systems, pages 2930-2941.
Bellemare, M. G., Dabney, W., and Munos, R. (2017). A distributional perspective on reinforcement learning. International Conference on Machine Learning, pages 459. PMLR.
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. (2013). The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279,
10

---

Boutilier, C., Dearden, R., Goldszmidt, M., et al. (1995). Exploiting structure in policy construction. In IJCAI, volume 14, pages 1104-1113.
Castro, P. S., Moitra, S., Gelada, C., Kumar, S., and Bellemare, M. G. (2018), Dopamine: A Research Framework for Deep Reinforcement Learning.
Dabney, W., Ostrovski, G., Silver, D., and Munos, R. (2018a). Implicit quantile networks for distributional reinforcement learning. In International conference on machine learning, pages 1096-1105. PMLR.
Dabney, W. Rowland, M., Bellemare, M., and Munos, (2018b). Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32.
Freirich, D., Shimkin, T., Meir, R., and Tamar, A. (2019). Distributional multivariate policy evaluation and exploration with the bellman gan. In International Conference on Machine Learning, pages 1992- PMLR.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. (2012). A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773.
Kulkarni, T. D., Saeedi, A., Gautam, S., and Gershman, S. (2016). Deep successor reinforcement learning. arXiv preprint 1606v.2660..022
Lin, Z., Yang, D., Zhao, L., Qin, T., Yang, G., and Liu, T.-Y. (2020). RD2: Reward decomposition with representation decomposition. Advances Neural Information Processing Systems, 33.
Lin, Z., Zhao, L., Yang, D., Qin, T., Yang, G., and Liu, T.-Y. (2019). Distributional reward decomposition for reinforcement learning. arXiv preprint 19/10/2
Mnih, V., Kavukcuoglu, K., Silver, D.,, Graves, A., Antonoglou, L., Wierstra, D., and Riedmiller, M.
(2013). Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.560
Nguyen, T. T., Gupta, S., and Venkatesh, S. (2020). Distributional reinforcement learning via moment matching. arXiv preprint arXiv:2007.1233:
Rowland, M., Bellemare, M. G., Dabney, W., Munos, R., and Teh, Y. W. (2018). An analysis of categorical distributional reinforcement learning.
Schaul, T., Horgan, D., Gregor, K., and Silver, D. (2015). Universal value function approximators. In International conference on machine learning, pages 1312-1328 1320. PMLR.
Sutton, R. S., Modayil, J., Delp, M., Degris, T., Pilarski, P. M., White, A., and Precup. D. (201 Ia).
Horde: A scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction. The 10th International Conference on Autonomous Agents and Multiagent Systems- Volume 2, pages 761-768.
Sutton, R. S., Modayil, J., Delp, M., Degris, T., Pilarski, P. M., White, A., and Precup, D. (201 1b). Horde: A scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems- Volume 2, AAMAS '11, page 761-768, Richland, SC. International Foundation for Autonomous Agents and Multiagent Systems.
Van Seijen, H., Fatemi, M., Romoff, J., Laroche, R., Barnes, T., and Tsang, (2017). Hybrid reward architecture for reinforcement learning. arXiv preprint arXiv: /706.04200
Yang, D., Zhao, L., Lin, Z., Qin, T., Bian, J., and Liu, T.-Y. (2019). Fully parameterized quantile function for distributional reinforcement learning. In Wallach, H., Larochelle, H., Beygelzimer, A., d' d'Alché- -ucc F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
Zhang. J., Bedi, A. S., Wang, M., and Koppel, A. (2020). Cautious reinforcement learning via distributional risk in the dual domain.


Reviews:
[{"review_id": "vmkk5SJvGf-", "reviewer": "Reviewer_oRgH", "summary": "The authors introduce Multi-Dimensional Distributional DQN (MD3QN) by combining distributional reinforcement learning and hybrid reward architecture. From the theoretical view, the authors prove the convergence of the joint distributional Bellman operator; empirically, the algorithm minimizes the Maximum Mean Discrepancy between joint return distribution and its Bellman target to make an update. The effectiveness of the MD3QN is demonstrated on a toy example and a set of Atari games. ", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n- The paper is well-written and easy to follow. The authors moved almost all proof to the appendix and made the paper accessible. \n- The motivation of the paper is well-stated, and MD3QN can capture the correlations between the dimensions across the reward function based on the toy demonstration. \n\nWeakness:\n- The experiment part is not strong enough:\n - The number of experiments is limited, and it's not clear if there are multiple runs or only one run for each task. The plot only has one trajectory for each algorithm without uncertainty which makes the comparisons not convincing due to the potential uncertainty (especially in the first 3 Atari tasks in Fig3)\n - I checked the experiment results from MMDQN(Nguyen et al., 2020) Fig 8 in Appendix, and MD3QN seems not outperform MMDQN on Gopher, MsPacman and Pong. If so, it will be more appropriate to run on the tasks that can truly reflect the value of HSA. \n - In Algorithm1, the action is the argmax across all reward dimensions, which is more similar to RD2 (eq7-eq8) than HRA (eq5-eq6). In this case, it's more fair to compare MD3QN with RD2. \n\nMinor Remarks:\n1. In eq(2), should it be $\\theta_{i-1}$ or $\\theta_i$?\n2. In l.125, is $\\mu^{\\pi}$ a joint distribution or marginal distribution? From its dimension, $\\mu$ looks like a marginal distribution, however in l.126 it represents a joint distribution. \n3. In eq(17), an $\\infty$ sign is missing.\n4. l.160-l.169, it is better to elaborate more about this part. It is still unclear why it's rigorous to use different metrics for proof and training aside from the fact that they lead to similar empirical performances. \n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "pClkiXhhoP", "reviewer": "Reviewer_ZZ5A", "summary": "This paper builds on two distinct yet important streams of RL literature:\n\n1. The idea of the hybrid reward RL is to train separate Q-value function heads on different reward sources, and combine them during action selection to improve the policy over using a single crafted reward.\n\n2. The idea of distributional RL is to learn the distribution of returns, rather than the expected value point estimates, which is shown to improve learning stability and quality of representation/policy (which I understood as aliasing in this work).\n\nThe paper first proves the theoretical convergence of the distributional Bellman operator for joint distributions of returns. Since this is not tractable to optimize in practice, the paper borrows from the idea of maximum mean discrepancy (MMD). This objective is differentiable and have favorable properties that can leverage the above analysis, based on combining: 1. being unbiased w.r.t. sample approximations of the Bellman operator, and 2. being exactly zero when the two measures are equal. A practical algorithm for implicitly learning the joint distribution of returns is proposed in which the samples from the joint distribution across all reward sources are the outputs of the network. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "originality: the main novelty of the work is the theoretical analysis of convergence of the joint distributional Bellman operators by generalizing contraction mappings to the n-dimensional setting. upon reading the related work (Nguyen et al., 2020), the algorithmic design and choice of optimization criterion seems to be rather trivial modifications of existing work to higher-dimensional distributions. \n\nquality: the work is technically sound. I have not checked all proof details, but they appear to be sound. the claimed results are quite intuitive and expected.\n\nclarity: the theoretical framework is easy to understand and main results are summarized well, along with their importance. the choice of MMD is also well-motivated as providing unbiased estimates for samples. finally, the experiments are well described. However, some wording and notation can be significantly improved. in particular:\n\n1. why is r_t bold in equations (1) and line 89?\n\n2. the details regarding how the joint distribution is represented architecturally is not well discussed. in particular, the authors mention that they replace the last layer of DQN with the samples from the joint distribution for each action. however, looking at the code briefly, and connecting the work to (Nguyen 2020), I am led to belief that you use something like IQN, where the samples from a base distribution are also provided to the network in addition to (s, a). This way, the joint distribution is implicitly modeled rather than facing the curse of dimensionality with C51 or QR-DQN where it is modeled directly. This then leaves the choice of how to model mu(s, a) architecturally to only IQR, correct? If so, then it would be better to discuss this in the main paper around the algorithm, as the current writeup may be confusing to readers.\n\n3. \"Crame loss\" -> \"Cramer loss\" in line 169\n\nsignificance: the work addresses an important problem in distributional RL that has not been studied despite much attention on the subject; how to learn joint distributions. I can envision this work being useful in many applications where Distributional RL has strong potential, and one particular use case is to learn distributions over successor features. however, the authors provide several other examples which could be interesting future work building on this paper. I therefore expect this paper to have significant impact in the RL community in the future.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "1. One major concern I have after re-reading the experiments and the hybrid reward discussion is whether or not a full joint distribution is really necessary for this setting. One particularly naive baseline is to simply learn a separate \"head\" for each reward as in HRA, that models only the marginal distribution of the return associated with each reward. Given that the joint distribution is never used directly in action selection (only the marginals appear to be), and given the added complexity of actually learning a full joint distribution on a complex task, I am not convinced that the naive approach will do much worse than the proposed method in this work, since it is still able to leverage distributional RL to avoid the aliasing problem albeit in a simpler manner. In either case, it would be good to provide such a baseline for comparison, since you could clearly show the advantage of learning a joint distribution over simple marginals. \n\nThe authors do not provide any discussion about the limitations of the proposed algorithm, and when it could possibly fail. For instance:\n\n2. how does the agent perform when the number of distinct reward signals increases? does it break at some point due to the burden of learning high-dimensional joint distributions? \n\n3. the authors don't mention how the full distribution could improve HRA in other respects beyond aliasing, since it appears identical. given the \"cost\" of learning a full joint distribution, what advantages does this bring? for example, when reward functions are not scaled equally , can the learned statistics from the distributions provide a scaling factor for free that can normalize them? are other complex transformations of distributions also possible here? could this improve robustness of HRA in which actions are selected by averaging the Q-values?", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "iGUacEYNUPS", "reviewer": "Reviewer_Soeq", "summary": "This paper extends distributional RL to multiple reward functions. Two contributions of the paper are 1) convergence of the joint distributional Bellman operator (under Wasserstein distance), 2) an algorithm that approximate the joint distributional Bellman operator by Maximum Mean Discrepancy (MMD). \nExperiments on four Atari games compared with HRA for multiple-reward functions. \n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThe definition of distributional RL is more complex than necessary. In eq 4, the a’ notation is overloaded. \n\nIn 6) and 8), the bold fonts are not kept consistent. The reward notation in 8) is better consistent too. \n\n140: the function f has been defined. \n\nTheorem 1 and Prop. 1 are for policy evaluation. \n\nTheorem 2’s presentation is different from Th. 1, in that the metric is changed. So it does not hold for Wasserstein distance?\n\nThe discussion on the misunderstanding about the metric is interesting. So the paper used Wasserstein distance to prove. — it is helpful to clarify this. And the training loss used the metric of ? This wasn’t discussed but perhaps later in the paper. It would be good to add here the discussion as it reads like incomplete. \n\nCan the authors explain the technical changes in the proofs because of the extension to hybrid reward setting? \n\nUsing MMD for training distRL is straightforward and not new. \n\nFour Atari games were selected (and unexplained). I’m not sure the experiments carries importance and whether they can support the claim. It is just too few and hard to interpret. \n \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "n/A. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "0.8"}, "extra_text": {}}, {"review_id": "8CxWhz0jsQv", "reviewer": "Reviewer_SgBk", "summary": "This paper studies the combination/extension of hybrid reward RL (HRA) with distributional value functions. The authors provide a convergence proof for the joint/multivariate distributional Bellman operator as a theoretical contribution. The authors also propose a practical algorithm which builds upon minimizing the MMD loss between a parametrized value distribution and the Bellman target. Finally, they show experimental results of their new algorithm and its performance gains on HRA baselines. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The methods used, HRA and distributional RL are not novel, but this is the first work that combines these two and in a well justified manner. Especially given that HRA uses independent losses, it is sensible to consider the modeling approach of learning a joint distribution of the different reward sources. While the extension of the MMD loss to the multivariate case is not a technically novel, but there doesn't necessarily need to be one if it works as is. \n\nQuality: The submission is technically sound in its theoretical contributions (convergence proof). On the experimental side, capturing the correlation between the random reward sources in the joint distribution is nicely illustrated. It is clear this would not happen when learning separate value functions in HRA. I am less convinced by the 4 Atari experiments as the correlation between the reward signals are either positively correlated or independent as stated in the appendix. Since the actions taken do not need to consider reward tradeoffs (as might happen when rewards are negatively correlated), I would suspect the results seen can be obtained by using separate distributional value networks for each of the reward signals, i.e. IQN or C51 and then combining them in the same way as RD^2 or other HRA architectures. That is to say, the performance gains seen are more attributable to modeling the randomness of each independent reward vs modeling the joint distribution. \n\nClarity: The clarity around the claim about not having to use the same metric to prove convergence and to train is somewhat confusing. I interpret line 58-61 to mean that the practical loss function used, i.e. sample based MMD is an unbiased estimator and therefore it is still a metric, which allows it to be used as the practical loss for this task. However, MMD is still different that the Wasserstein metric. Is it the case that if convergence is proven with one metric, that is sufficient to show that the operator converges for any metric? If so, this should simply be stated directly. In my opinion, the writing obfuscates this a little. Lines 109-114 and the paragraph starting on line 160 reads as if there will be results and/or a proof to demonstrate *why* the metrics used can be different, and that \"correcting this misunderstanding\" seems to be a contribution of the paper. Finally, the exposition on line 160 is still a bit confusing even under the above assumption, because the KL-divergence is not a metric. \n\nSignificance: I think this is a strong contribution if the authors could show that the joint distribution learned which captures the correlation between rewards could be useful in a more complex task/setting. Instead of the Atari experiments, maybe some experiments demonstrating the applications of the joint return learning as outlined in lines 237-245 could be useful. I see one current limitation as not empirically demonstrating a situation where modeling the joint distribution is particularly needed. \n\nMisc: I am curious if there is any intuition for selecting the kernel bandwiths?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors do not address limitations and societal impact is not applicable. See \"significance\" in main review. ", "review_text": "", "extra_scores": {"ethical_concerns": "There are no ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers were in favor of acceptance, and from reading the paper myself I find it to be a solid contribution and potentially impactful. I definitely want to second ZZ5A's callout for how this work could be combined with successor features as an interesting direction for future work (thinking of Barreto et al.'s transfer learning work in particular). As reviewers noted, there is work in the literature that overlaps and thus reduces novelty. One that I did not see mentioned by reviewers or authors is "Distributional Multivariate Policy Evaluation and Exploration with the Bellman GAN" by Freirich et al., and in an ideal world the final version of this paper would provide a comparison with this method, but certainly should be discussed due to the overlap in problems being solved. Another improvement that I would really appreciate in a final version, although I understand that computational costs can make this difficult, would be an evaluation on the rest of (or reasonable subset of) the other Atari games. My reasoning here is that the current experiments tell me how well this method works in the type of problem it was intended for, but I also want to know how well it works out of that scope or more generally. 
Decision: Accept (Poster)