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 taoqin@microsoffe
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 (Danney 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 mulei-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 (s 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 0z. 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(11, 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 μ4 and ju" decays as increases. We have the following corollary:
Corollary If Pi+1 TT then
o0, p4 p="
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 Eis 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 (Vnn 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 square 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": "TuLP7kkWVwV", "reviewer": "Reviewer_A2nn", "summary": "The paper suggests exploiting the awareness and available identification of multiple reward sources to factorise a distributional value function and exploit this explicit structure to learn and exploit the correlation structure of distinct reward sources. The overall architecture predicts representative samples that describe the joint (total) value distribution. The paper presents experimental results that confirm efficacy of the approach. The paper appears to be correct overall, although there are some technical issues that mar the overall impression. \n", "questions": "", "limitations": "", "rating": 5, "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": "First noticeable issue with the paper is Algorithm-1. Authors suggest that it is supposed to calculate the gradient of the MMD loss. However, it actually calculates an estimate of the MMD loss itself. This is in a sharp contrast to the preceding paragraph that speaks of the necessary simplifications. The appendix is of no help in this situation. Authors do not describe how the loss estimate will translate into it's gradient estimate. The text of the algorithm is also festered by repeated use of symbols at distinct semantic locations (\"a prime\" serves too many roles here, so is Z_{1:M}). It is difficult to imagine how this pseudo-code was translated into the actual code to function correctly. Notwithstanding the practical success of the code, this does not bode well for the alignment of the theory and practice of this paper. \n\nIt is also not clear why borrowing a few factorisation tricks from the MARL literature, where the correlation of independent executions is implicitly learned, can not be of use here. QR-MIX and MMD-MIX would be among the more recent structures that can be successfully deployed here. Authors cannot avoid a justification of not using these architectures (at least as a baseline) in their domain. \n\nAuthors also miss more important references that bear key similarities to their approach:\n\"Distributional Multivariate Policy Evaluation and Exploration with the Bellman GAN\" by Freirich et al\n\"Sample-based distributional policy gradient\" by Singh, Lee and Chen\nIt would be interesting if authors could compare their work to the above two (at least in the rebuttal).\n\nPeeves and typos:\n\nThe paper is also overly reliant on the Appendix. Though the claims to appear to be correct, there are intuitions that may be gathered here, that do not appear in the main text. \n\nLine 169 \"Crame_r\" \nLine 170 \"Operator\" \nLine 277, what is the exact nature of an \"iteration\" here? \n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Work is sufficiently distant from real world applications to avoid social impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "~5h (with extra noise)"}, "extra_text": {}}, {"review_id": "GGmxKzcNXRX", "reviewer": "Reviewer_X5DH", "summary": "The authors present MD3QN, which leverages the distributional approximation capacity of C51 and its likes as well as the reward multiplicity of some reinforcement learning tasks. The line or research centered around C51 has the advantage of learning an entire distribution of the reward (instead of just the expectation), which is often considered a better representation of the RL task at hand. Meanwhile, the line of research around HRA separately model the value function for each reward source to better the algorithm's sample efficiency. MD3QN combines the said advantages in a correlated multiple-reward setting where the agent learns the joint reward distribution for a better task representation and a better sample efficiency. The proposed algorithm is shown to outperform one previous work (HRA) in a number of Atari games, and it seems to learn the reward correlations from custom-designed mazes. ", "questions": "", "limitations": "", "rating": 3, "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:\nI think the presentation could do a much better job at highlighting the novelty of the algorithm design. The current paper does not pinpoint the challenges encountered in developing MD3QN. I struggled to identify the unique contribution of MD3QN in the transition from Section 2 to Section 3, where MD3QN is illustrated almost like a mere combination of C51 and HRA when that is certainly not the case according to the introduction. The strength of MD3QN must then lie in capturing the correlated randomness in source-specific returns, but the experiments and the figures do not point to that strength. Figure 2 is especially obscure in that I wasn't convinced by its effectiveness or its message. How better does MD3QN do in identifying the correlations than do previous works? What quantitative measures could you utilize to gauge the performance in capturing the correlations? Also, why have the maze-multireward rewards been split into blue-orange and red-green pairs?\n\nOn a separate note, the paper claims multiple times that it \"corrects the misunderstanding\" of existing works, but I would like to raise a few points on this claim. The first mention of this claim is in Section 1, where \"We correct a misunderstanding in existing works on distributional RL, and argue that we can use other metrics than the Wasserstein metric over the joint distribution as training losses to approximate the Bellman operator.\" The second and third mentions of the claim both appear in Section 3: \"... misunderstanding that we must use the same metric for proving convergence and as training loss.\" Now, those two understandings are very different from each other, and the rest of the paper does not clarify which \"misunderstanding\" that the authors are \"correcting\". Is it that only the Wasserstein metric must be used? Or is it that the same metric must be used in convergence and in training loss? Furthermore, the first mention of the \"misunderstanding in existing works\" requires references. The second and third mentions of the \"misunderstanding\" do contain references to existing works, but as far as my understanding of Bellemare2017 and Nguyen2020 goes, those two papers in no way claim  that \"only the Wasserstein metric must be used\" or \"the same metric must be used in convergence and in training loss\". In fact, the MD3QN authors only list some other works that employ different metrics for showing convergence and in training loss. Consequently, I am lost as to which misunderstanding in which previous work the authors are referring to. If the \"corrects the misunderstanding\" statement is not addressed properly, it will put the submitted paper in a very bad light. You would be claiming to have corrected a misunderstanding that never even existed to begin with. Far worse than that would be to wrongly attribute that non-existent misunderstanding to a well-established prior art. It is very much needed that the authors cite exactly which sources make the said misunderstanding, but even that would not be enough to credit the MD3QN authors with \"correcting a common misunderstanding\" because many studies that the authors themselves have cited already use different metrics for convergence and in training loss. I am left puzzled as to why such a statement is needed in this paper after all.\n\n\nQuality:\nThe paper was easy to read. There is considerable effort in founding a mathematical establishment before delving into the experiments, notations are consistent, as far as I can tell.\nSome minor issues:\nSection 2.1 is are sampled from the replay memory --> are sampled from the replay memory\nSection 3.2 Crame loss --> Cramer loss\nSection 5.1 The move has no effect --> The attempted move does not succeed (because, it's not that there is a move and then it has no effect. It's that there is no move in the first place.)\n\n\nClarity:\nThe sectioning feels a bit awkward. For example, what was meant to be the distinction between Section 2.1 and Section 3.1? At first, it looked like Section 2.1 is generally applicable to related research while Section 3.1 is specific to MD3QN, but that does not seem to be the case after continuing on to Section 3.2, where a concept (Bellman operator) much broader than this specific paper is introduced. The scopes of Sections 2 and 3 did not help in clarifying the notations, background, or the settings.\n\nAgain, referring to the \"common misunderstanding\" claim that the authors are making: authors write that they \"clarify that the metric for proving convergence and the metric for training loss play different roles, and are not necessarily the same metric\". No, they have not clarified that statement. They have cited a number of works that do not adhere to the \"common misunderstanding\", which then calls to question whether the said misunderstanding is indeed \"common\". It remains unclear i) who said that convergence metric and the training metric have to be the same, as well as ii) what mathematical tools the MD3QN authors have wielded to disprove the misconception.\n\nThe term \"exclusive\" is used without definition in Section 5.1. Is Figure 2 (a) supposed to show negative correlation?\n\nSignificance:\nThe key element of the paper is the correlatedness of the multiple rewards, and the two theorems do not show how the correlation plays a pivotal role in the design philosophy of MD3QN. Theorem 1 being the contraction mapping theorem for two joint distributions and Theorem 2 being the unbiasedness of the joint distribution estimator, neither of the two are, strictly speaking, contributions that are attributable to MD3QN, even if they show the \"interesting\" part when it comes to the correlation of the multiple rewards.\n\nThe experiments do not particularly outshine other submitted papers. For instance, the maze setting looks a bit artificial, and it would have been much more effective if accompanied with a real-world scenario, despite the fact that only some simplified analogues are tested.\n\nAlso, reward multiplicity poses some intriguing opportunities for rewards defined \"outside the game\". Intrinsic motivation papers geared towards better exploration or more diversity are examples of reward structure tuning. It would be enlightening to observe how MD3QN performs under those circumstances and whether the distributional RL advantage and/or the sample efficiency stand out in an ablation study.\n\nThe effectiveness of Figure 2 is unclear. The figures do not quantify how well MD3QN does, and they do not compare MD3QN to previous works. More importantly, why have the Atari Games been limited to just four games? For example, Atari Seaquest game grants points for destroying the sharks as well as some points for rescuing the divers. Why has this game been excluded?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors acknowledged that, despite MD3QN being a general framework, its actual application to goal-conditioned RL and successor states are left to future works. I believe the same kind of limitation acknowledgement is in order for i) testing with only four multiple-reward Atari games, ii) not clarifying a \"common misunderstanding\" in existing works despite claiming to have done so, and iii) the rather factitious nature of the maze experiments.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "EXTv7EKkGkO", "reviewer": "Reviewer_TpSB", "summary": "This paper proposes a method for distributional reinforcement in environments with multiple channels of reward. The joint distributional Bellman operator is introduced and its properties analysed, and an implementation based on MMD distances is described, with experiments carried on a small-scale gridworld and several Atari games.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "__Post-discussion__\n\nI thank the authors for their detailed rebuttal. I like the fundamental idea of the paper, although I think that the updates required to the paper are substantial enough that further review may be beneficial, particularly around the unbiasedness claims for MMD estimation, \"common misunderstanding\" remarks, and experimental details. Based on this, I would not recommend acceptance at this time, but I think that with further updates along these lines, this will be a strong paper.\n\n__Main points__\n\nThere is an interesting central contribution in this paper, that I think should be of interest broadly to the RL community. Both distributional reinforcement learning and multi-reward channel RL are existing ideas, but their combination is an interesting idea, and the perspective the authors take in this paper is novel as far as I am aware. The exposition is generally clear, and the experimental results show interesting results from the newly-proposed approach. I would therefore like to see this paper accepted.\n\nHowever, there are several parts of the paper that I think require clarification before the paper is ready for acceptance:\n\n1. I found the discussion in the paragraphs at Lines 107 & 160 very unclear, for several reasons:\n      * First, it is known that metrics for analysis and for projection play different roles in distributional reinforcement learning; see for example the analysis of Dabney et al. (2018b), in which a W_1 projection is used, while the contraction of the projected Bellman operator is performed under W_infty.\n      * Second, the paragraph at Line 160 is unclear when mentioning \"convergence\" as to what it is referring to.  Repeated application of the distributional Bellman operator (or joint distributional Bellman operator) can be shown to be convergent using a variety of metrics, but this is distinct from understanding convergence of approximate algorithms. The example of C51 mentioned in this paragraph is useful to consider: global minimization of the expected KL loss leads to the same result as the application of the projected Bellman operator, this does not imply that using the KL loss leads to a convergent algorithm. Analysis of the projected operator, for example, would be required to show this. On this basis, I don't think the sentence \"We illustrate that, once the convergence of the Bellman operator is proved under a certain metric, we can use other metrics than the metric for proving convergence, as training loss to approximate the Bellman operator.\" is true.\n2. I think currently there is not enough discussion of the fact that multiple independent samples from T mu (s, a; \\theta_i) are required to unbiasedly estimate gradients of the MMD^2 loss. This is an important fact, since multiple independent samples from a single given state-action pair are typically not available in online/replay-based reinforcement learning (e.g. Atari), which makes the sentence \"the temporal difference loss MMD^2 can be optimized without bias by transition samples\" potentially misleading.\n3. Experiments: Overall these experiments nicely demonstrate the methods introduced by the paper, but there are some missing details that should be added. I couldn't find values for some hyperparameters such as the number of particles used, and in other cases I couldn't find details of how hyperparameter sweeps (e.g. on kernel bandwidth, learning rates, etc.) were carried out, or on how the Atari environments were selected, or their multichannel reward functions were designed (i.e. were other alternative rewards/games tried, did these work as well, etc.)\n\n__Minor points__\n\nProposition 1: The metric of convergence should be specified; presumably it is \\overline{d}_p.\n\nThere is some previous work on multi-channel distributional reinforcement learning that it would be good to comment on in the related work section, such as \"Distributional Reward Decomposition for Reinforcement Learning\" (NeurIPS, 2019).\n\nThe paper would naturally be stronger with an evaluation carried out on the full set of Atari games. However, I don't think this should be necessary for acceptance. The experiments that are provided in the paper serve as a useful proof-of-concept that the approach can learn meaningful joint distributions over rewards, and that this approach can have benefits in large-scale RL agents.\n\n__Notation/typos__\n\nE_\\Sigma vs \\mathbb{E}_\\Sigma\n\"Guassian\" -> \"Gaussian\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See main review for comments on limitations. I did not see any discussion around societal impact, although this work is concerned with a core reinforcement learning algorithm, rather than applications.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "8nDoU6UnK9l", "reviewer": "Reviewer_fJbd", "summary": "The paper proposes a new learning algorithm for modeling the distribution of multi-dimensional reward in the reinforcement learning (RL) setting. It proposes a bellman evaluation operator for the distribution of multi-dimensional reward and shows its optimality. Moreover, a MMD based metric is proposed to measure the difference between two distributions and this metric is fit for the typical RL learning settings. The new learning method is tested and performs the best on one designed maze environment and multiple 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": "Originality:\nThe idea of the learning distribution of multi-dimensional reward in RL is straightforward. But building a learning framework upon that is complex and significant. \n\nQuality:\n\nThe learning framework is designed based on value-based reinforcement learning framework. The proof of the bellman operator is sound. Moreover, it uses different losses for proving and training. The MMD based training loss shows good performance in the experimental section.\n\nClarity:\n\nThe paper is well-written. And I do like the paper's structure.  It frames each main theorems and proposition, then explains the relation between them. Moreover, the authors also use many intuitions among them. Though the proof may not be fancy, however, the structure does make the understanding much easier.\n\nSignificance:\n\nThe problem to model the distribution of multi-dimensional reward is vital to push the performance of RL algorithms.\n\nQuestions:\n\n1. I am wondering can you explain more about the kernel selection? Is any kernel that can make MMD a metric is a good kernel? \n2. It would be better if you can compare MMD with some other distance metric and show the superiority of MMD.\n\nMinor:\nThe $r_t$ of equation (2) should be bolded.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A.", "review_text": "", "extra_scores": {"ethical_concerns": "NO", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers generally like the fundamental ideas of this paper.  However, there was a lot of confusion regarding the notation and the supposedly common misunderstandings of the community. The reviewers also raised some concerns about the experiments.  While the rebuttal clarified many concerns, the revisions required for the paper to be accepted are significant and therefore another round of reviews will be needed.  Hence, the authors are encouraged to revise the paper and to resubmit it at another venue.
Decision: Accept (Poster)