Distributional Reinforcement Learning for Multi-Dimensional Reward Functions

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.


Reviews and author responses:

 We are glad to see that our rebuttal addressed your questions. We will make further revisions on analyzing Atari experiments as well as comparing the MMDQN baseline in the next version of the manuscript. We believe that our paper will be strengthened after that. Thank you again for your valuable and generous suggestions!

Best,

Authors

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.   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. 

Quality: 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. 

Clarity: 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. 

Significance: 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. 

Misc: I am curious if there is any intuition for selecting the kernel bandwiths? The authors do not address limitations and societal impact is not applicable. See "significance" in main review. 

 Thanks for your response. I've read the rebuttal and discussion. I will raise my score given the authors addressed my questions and with an additional experiment highlighting the importance of learning the joint distribution. I also agree with reviewer (3) that this is an important and open question in RL and this seems to be a good direction. 

However, I am still concerned now about the inclusion of the Atari experiments. As stated in their response to both me and reviewer (2), "Atari games might not be related to the joint distribution" and even that "MMDQN outperforms" their algorithm. I would strongly suggest some additional clarification in the revision around this, perhaps even adding the MMDQN baseline for transparency.

 We are happy to hear that our rebuttal addressed your concerns well. The additional results in the rebuttal will be added in the revision. Thank you again for your valuable suggestions and comments, which we believe strengthen our paper. 

Best,

Authors

This paper builds on two distinct yet important streams of RL literature:

1. 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.

2. 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).

The 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.   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. 

quality: 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.

clarity: 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:

1. why is r_t bold in equations (1) and line 89?

2. 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.

3. "Crame loss" -> "Cramer loss" in line 169

significance: 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. 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. 

The authors do not provide any discussion about the limitations of the proposed algorithm, and when it could possibly fail. For instance:

2. 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? 

3. 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?

 Thanks a lot for your response, which has addressed my main concerns. I have looked at the added result and it is indeed convincing to me that modelling the full joint distribution yields a tangible benefit beyond just the marginals. I believe this paper addresses an important, and still open, problem in RL, e.g. how to do distRL in a principled way with high dimensional value functions, and I can easily foresee many applications of this work in other areas of RL in addition to HRA, e.g. successor features. I am willing to raise my score for all these reasons. 

 We are happy to hear that our rebuttal addressed your concerns well. The additional results will be added in the revision. Thank you again for the valuable suggestions and comments, which we believe strengthen our paper!

If you have any remaining suggestions or concerns, please let us know!

Best,

Authors

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.   Strength:
- The paper is well-written and easy to follow. The authors moved almost all proof to the appendix and made the paper accessible. 
- 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. 

Weakness:
- The experiment part is not strong enough:
 - 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)
 - 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. 
 - 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. 

Minor Remarks:
1. In eq(2), should it be $\theta_{i-1}$ or $\theta_i$?
2. 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. 
3. In eq(17), an $\infty$ sign is missing.
4. 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. 



 yes

 Thank you for the response. 
The rebuttal has addressed my major concerns, and I hope the additional results can be added in the revision to make the results more convincing. 
I thereby raise my score to 6. 


 Dear reviewers,

Thank you for your time and efforts in reviewing our paper.

This is a kind reminder that we are one week into the discussion period. We believe that we have sincerely and successfully addressed your concerns and questions in your reviews, with the added supporting experimental results.

If you have any further questions or concerns after reading our comments, please do not hesitate to let us know.

Authors

 Thank you for the thoughtful and constructive suggestions! We have taken all the comments into consideration and summarized the responses as follows:

1) *Quality: Action selection in Atari games do not consider reward tradeoffs. A baseline learns 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.*

**Response:**
Thanks for your constructive advice! We agree that, as mentioned in your comments, the performance of Atari games might not be directly related to the joint distribution. Taking your advice in the comment (3) (significance), we further design an extra experiment to simulate a real-world scenario and compare it with the baseline that learns N marginal distributions. We illustrate why only learning marginal distributions may fail from theory and experiment results. For detailed information and results, please refer to the response to (3).

-------------------------------------------------------------------------------------

2) *Clarity: The clarity around the claim about not having to use the same metric to prove convergence and to train is somewhat confusing.*

**Response:**
Thanks for your advice! Unfortunately, this point was not clear in our original manuscript, and we apologize for all the confusion. We are eager to clarify it here:

First of all, we would like to restate our claim: (1) we can use a metric A to prove that the (projected) distributional Bellman operator is a contraction, and (2) we can use another distribution measure B as training loss (here B is not necessarily a metric), as long as global minimization of this loss leads to the same result as the application of the (projected) Bellman operator. 

To be specific, for the distributional Bellman operator, the metric A used for proof and measure B used as loss are not necessarily the same. In our case, we use MMD loss as B, and the above claim holds when the number of samples goes to infinity. For the projected distributional Bellman operator, there exists another metric used for projection, and this is not necessarily the same as B, too.

The analysis for Bellman operator and projected Bellman operator is the same: Since (projected) Bellman operator with metric A is a contraction, and global minimization of loss B has the same result (projected) Bellman operator with metric A, it is natural to conclude that it is okay to use B as training loss with the theoretical guarantee proved by A not broken.

-------------------------------------------------------------------------------------

3) *Significance: 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.*

**Response:**
Thanks for your advice! To further illustrate the necessity of learning a full joint distribution, we add another experiment during the rebuttal, and explain why only modeling marginal distributions may fail in some settings (e.g., in environments with multiple correlated constraints)  from theory and experiment results.

There are a number of real-world scenarios where multiple constraints should be met simultaneously. For instance, in autonomous driving, we need to balance the safety (distance from other cars), the speed, the comfort (the acceleration, etc.), and many other factors to make sure the car functions normally. Those statistics can be viewed as sub-rewards in our settings. For instance, we would like to make the agent be aware of the joint distributions of distance from other cars and the speed, and add a constraint on this distribution (e.g., d > 1m, v < 50km/h).

From the theoretical perspective, only the algorithm which is capable of modeling the joint distribution can find this optimal solution. If the algorithm can only model the marginal distributions, we can correctly compute the probability of simultaneously meeting multiple constraints only if they are independent. However, this is not common in real-world scenarios.

We use the same Maze environment as in Figure 1 of our paper and modify the layout as a simplified analogues (see [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/constraint_maze.png ]). We set multiple constraints on the total return (in our experiment, we have three constraints: for the initial state, total red return > 0.6, total green return > 0.6, and total blue return > 0.6, where the agent gets a reward = 1.0 of the specific color after collecting a specific colored block), and the goal is to find a policy which has the highest probability to satisfy all constraints. 

We use the joint distribution by MD3QN to achieve this: specifically, given the modeled joint distribution $\textbf{Z}(s,a)$, the agent can compute the probability to satisfy all the constraint in the joint distribution and take action by $\arg \max_a \mathbb{P}(\textbf{Z}(s,a) \text{ satisfy all three constraint})$. We also modify the joint distributional Bellman optimality operator to maximize the probability to satisfy all constraints.

For baseline methods, we extend the MMDQN by multiple heads $Z_{1:N}$ to reflect $N$ sources of rewards. This method is different from our algorithm in that only marginal distributions of each dimension are learned. We test two ways to approximately compute the probability to satisfy all constraints by marginal distributions: the Marginal_SUM baseline maximizes $\sum_{i=1}^N \mathbb{P}(Z_i(s,a) \text{ satisfy i-th constraint})$, the sum of probabilities for each reward to satisfy constraint, while the Marginal_PROD baseline maximizes $\Pi_{i=1}^N \mathbb{P}(Z_i(s,a) \text{ satisfy i-th constraint})$, the product of probabilities for each reward to satisfy constraint. 

The results can be found here [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/multiple_constraint.png ]. It can be seen that maximizing the probability based on joint distribution can significantly outperform two baseline methods which both use marginal distribution information. 

-------------------------------------------------------------------------------------

4) *Misc: I am curious if there is any intuition for selecting the kernel bandwidths?*

**Response:**
To conclude, we believe that the kernel bandwidths should cover a wide range. Here are the experiment results during the hyperparameter tuning, and we believe this can support this intuition. We tested 3 sets of kernel bandwidths: bandwidths $B_1 = [2^{-8}, 2^{-7}, 2^{-6}, \cdots, 2^8]$, bandwidths $B_2 = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.48, 0.64, 0.80, 0.96, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0]$, bandwidths $B_3 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]$) and perform the same experiment as in Section 5.1 and Figure 2 in our paper to test how MD3QN models the joint distribution under different bandwidth settings. The result can be found in  [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/bandwidth.png ]. For $B_1$, the MMD between prediction and ground truth distribution is 0.026, where for $B_2$, MMD = 0.036, and for $B_3$, MMD = 1.348. We found $B_1$ and $B_2$ can capture the joint distribution better than $B_3$, and $B_1$ is slightly better than $B_2$. Finally we use Bandwidth $B_1$ in our experiments.

 Thank you for the thoughtful and constructive suggestions! We have taken all the comments into consideration and summarized the responses as follows:

1) *Concerns on the experiments:*

    a) *Multiple seeds?*

    b) *To compare with MMDQN?*

    c) *To compare with RD2?*

**Response:**
Thanks for your constructive advice! We add some experiments according to your suggestions.

a) & c):  First, we add two additional Atari games, and we also add an RD2 for comparison, and MD3QN also outperforms RD2 in control settings on Atari games. We show the training curve on those games with multiple random seeds here [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/Atari_HRA_RD2.png ]. We further tune the kernel bandwidths and use $B = [2^{-8}, 2^{-7}, 2^{-6}, \cdots, 2^8]$.

b) We would like to highlight that our algorithm is capable of modeling the joint distribution of return. First of all, our algorithm can outperform HRA on all the Atari games, which shows the validity of modeling the joint distribution compared to the scalar.

However, the reason why we cannot outperform MMDQN might lie in the fact that the performance of Atari games might not be related to the joint distribution. In our algorithm, we choose the argmax of the expected total return as the desired action and the reward tradeoffs are not under consideration. Therefore, it is hard to directly interpret how the performance is explicitly related to the joint distribution. This might partially explain why MMDQN can outperform our method given it only models the total return distribution.

During the rebuttal, we add another experiment to show the necessity of modeling the joint distribution, and explain why only modeling marginal distributions may fail in some settings (e.g., in environments with multiple correlated constraints) from theory and experiment results.

There are a number of real-world scenarios where multiple constraints should be met simultaneously. For instance, in autonomous driving, we need to balance the safety (distance from other cars), the speed, the comfort (the acceleration, etc.), and many other factors to make sure the car functions normally. Those statistics can be viewed as sub-rewards in our settings. For instance, we would like to make the agent be aware of the joint distributions of distance from other cars and the speed, and add a constraint on this distribution (e.g., d > 1m, v < 50km/h).

From the theoretical perspective, only the algorithm which is capable of modeling the joint distribution can find this optimal solution. If the algorithm can only model the marginal distributions, we can correctly compute the probability of simultaneously meeting multiple constraints only if they are independent. However, this is not common in real-world scenarios.

We use the same Maze environment as in Figure 1 of our paper and modify the layout as a simplified analogues (see [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/constraint_maze.png ]). We set multiple constraints on the total return (in our experiment, we have three constraints: for the initial state, total red return > 0.6, total green return > 0.6, and total blue return > 0.6, where the agent gets a reward = 1.0 of the specific color after collecting a specific colored block), and the goal is to find a policy which has the highest probability to satisfy all constraints. 

We use the joint distribution by MD3QN to achieve this: specifically, given the modeled joint distribution $\textbf{Z}(s,a)$, the agent can compute the probability to satisfy all the constraints in the joint distribution and take action by $\arg \max_a \mathbb{P}(\textbf{Z}(s,a) \text{ satisfy all three constraints})$. We also modify the joint distributional Bellman optimality operator to maximize the probability to satisfy all constraints.

For baseline methods, we extend the MMDQN by multiple heads $Z_{1:N}$ to reflect $N$ sources of rewards. This method is different from our algorithm in that only the marginal distribution of each dimension is learned. We test two ways to approximately compute the probability to satisfy all constraints by marginal distributions: the Marginal_SUM baseline maximizes $\sum_{i=1}^N \mathbb{P}(Z_i(s,a) \text{ satisfy i-th constraint})$, the sum of probabilities for each reward to satisfy constraint, while the Marginal_PROD baseline maximizes $\Pi_{i=1}^N \mathbb{P}(Z_i(s,a) \text{ satisfy i-th constraint})$, the product of probabilities for each reward to satisfy constraint. 

The results can be found here [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/multiple_constraint.png ]. It can be seen that maximizing the probability based on joint distribution can significantly outperform two baseline methods which both use marginal distribution information. 

-------------------------------------------------------------------------------------

2) *In eq(2), should it be θi−1 or θi?*

**Response:**
It should be $\theta_{i-1}$, since the value iteration uses the parameter $\theta_{i-1}$ in the previous step to compute the target value, which is used to optimize the parameter $\theta_{i}$ in the next step. 

-------------------------------------------------------------------------------------

3) *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.*

**Response:**
$\mu^\pi$ in line 125 is a joint distribution since $\mathcal{P}(\mathbb{R}^N)$ denotes the joint distribution of an N-dimensional vector. 

-------------------------------------------------------------------------------------

4) *In eq(17), an $\infty$ sign is missing.*

**Response:**
We appreciate your comments and will correct the error in the revision.

-------------------------------------------------------------------------------------

5) *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.*

**Response:**
Thanks for your advice! First of all, we would like to restate our claim: (1) we can use a metric A to prove that the (projected) distributional Bellman operator is a contraction, and (2) we can use another distribution measure B as training loss, as long as global minimization of this loss leads to the same result as the application of the (projected) Bellman operator. 

To be specific, for the distributional Bellman operator, the metric A used for proof and measure B used as loss are not necessarily the same. In our case, we use MMD loss as B, and the above claim holds when the number of samples goes to infinity. For the projected distributional Bellman operator, there exists another metric used for projection, and this is not necessarily the same as B, too.

Since (projected) Bellman operator with metric, A is a contraction, and global minimization of loss B has the same result (projected) Bellman operator, it is natural to conclude that it is okay to use B as training loss with the theoretical guarantee proved by A not broken.


 Thank you for the thoughtful and constructive suggestions! We have taken all the comments into consideration and summarized the responses as follows:

1) *Notations and details:*

**Response:**
About notation issues in our paper: we really appreciate your advice and agree that these notations are not very precise. We will adopt your suggestion and update the notation in the revision. 

-------------------------------------------------------------------------------------

2) *Theorem 2’s presentation is different from Th. 1, in that the metric is changed. So it does not hold for the Wasserstein distance?*

**Response:**
Theorem 2 proves the contraction of the expectation of our proposed joint distributional Bellman optimality operator. However, the joint distributional Bellman optimality operator itself is not a contraction on the Wasserstein metric, which follows the conclusion of Proposition 1 from Bellemare [1]. Therefore, we only prove that the expectation of this operator is a contraction. 

[1] A Distributional Perspective on Reinforcement Learning. 

-------------------------------------------------------------------------------------

3) *The discussion on the misunderstanding about the metric is interesting. So the paper used the 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.*

**Response:**
Thanks! We use Wasserstein to prove, and the training loss is the MMD metric which is detailed in section 3.3. We will clarify this in the first paragraph of section 3 in the revision. 

-------------------------------------------------------------------------------------

4) *Can the authors explain the technical changes in the proofs because of the extension to hybrid reward setting?*

**Response:**
We believe the proofs are not a trivial extension to existing work. Because the Wasserstein metric has a closed-form in the 1-d situation, while in n-d, we can only use the coupling form in the proof. Therefore, we use new proof methods in Lemma 3 and Lemma 4 of the appendix which both use the coupling form of Wasserstein distance. In Lemma 3, we prove that if the Wasserstein distance of two joint distributions is $C$, then after applying the same linear transform to these two distributions with a scale factor of $\gamma$, the Wasserstein distance is at most $\gamma C$. In Lemma 4, we prove that the Wasserstein distance of two joint distributions under conditional random variable $A$ is no greater than the maximum Wasserstein distance overall specific conditions ($A=a$). Both the proof of Lemma 3 and Lemma 4 are not proposed in the existing work. By the results of Lemma 3 and Lemma 4, we are able to prove the contraction results in Theorem 1. 

-------------------------------------------------------------------------------------

5) *Four 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.*

**Response:**
We are sorry that this point was not clear in our original manuscript, and we will add the following details in the revision of our paper.

On the experiment to evaluate the performance of MD3QN, we use two criteria for selecting the Atari games:
* There are multiple sources of reward.
* Different sources of reward can be recognized from scalar score signals.
The way we compute the vectored reward is detailed in Appendix A.2. It is worth noticing that not all Atari games have multiple reward sources, and not all the rewards with different sources can be decomposed only from the scalar. 

For instance, in Gopher, the possible rewards for different events are 0 (for nothing), 20 (for filling the hole), 100 (for hitting the gopher). If r = 120, we can tell that it is composed of a reward of 20 and a reward of 100. However, in Seaquest, the agent will get the same reward for shooting the sharks and shooting the submarines. It is also hard to tell how many rewards are collected due to the leftover oxygen and how many rewards are collected due to saving the divers. This makes Sequest not meet the second criteria, and we omit Seaquest in our experiment. To conclude, we use the two aforementioned criteria to select the games that are of our interest.

For Pong, the agent will get a reward = +1 for winning a round and a reward = -1 for losing a round. In fact, there are a lot of games with the same pattern of reward: +1 for winning a round and -1 for losing a round. However, we consider the sources of reward to be the same, so we only choose Pong as a representative game in this type of games. 

During the rebuttal, we go over all the Atari games and locate another two that meet the criteria, and add them to the experiments. Finally, we choose AirRaid, Asteroid, Gopher, MsPacman, Pong, UpNDown as the evaluation set. We further tune the kernel bandwidths and use $B = [ 2^{-8}, 2^{-7}, 2^{-6}, \cdots, 2^8]$ here. The full results on Atari games with multiple random seeds are shown in [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/Atari_HRA.png ].


 Thank you for the thoughtful and constructive suggestions! We have taken all the comments into consideration and summarized the responses as follows:

1) *Originality: the algorithmic design and choice of optimization criterion seems to be rather trivial modifications of existing work to higher-dimensional distributions.*

**Response:**
Our work proposes an algorithm that is capable of capturing high-dimensional joint distribution, and we believe this is not merely a trivial modification. We list two major challenges and the differences below:

First of all, it is not trivial to extend the theory of C51 into multi-dimensional, although the proof is not (that) fancy. The Wasserstein metric has a closed form in the 1-d situation, while in n-d, we can only use the coupling form in the proof. Therefore, we use new proof methods in Lemma 3 and Lemma 4, which both use the coupling form of Wasserstein distance. In Lemma 3, we prove that if the Wasserstein distance of two joint distributions is $C$, then after applying the same linear transform to these two distributions with a scale factor of $\gamma$, the Wasserstein distance is at most $\gamma C$. In Lemma 4, we prove that the Wasserstein distance of two joint distributions under conditional random variable $A$ is no greater than the maximum Wasserstein distance over all specific conditions ($A=a$). Both the proof of Lemma 3 and Lemma 4 are not proposed in existing work. By the results of Lemma 3 and Lemma 4, we are able to prove the contraction results in Theorem 1.
To the best of our knowledge, we are the first to give a rigorous proof for the contraction of the multidimensional Bellman evaluation operator and the contraction of the expectation summation of the multidimensional Bellman optimality operator, although it seems to be intuitive.

We would also like to highlight that our algorithm is capable of learning the joint distribution. and this is of great use in downstream tasks. There are a number of real-world scenarios where multiple constraints should be met simultaneously. To further illustrate that, we add an experiment. For detailed information and results, please refer to the response to (3).

-------------------------------------------------------------------------------------

2) *clarity: However, some wording and notation can be significantly improved. in particular:*

    a) *Why is r_t bold in equations (1) and line 89?*

    b) *Details regarding how the joint distribution is represented architecturally.*

    c) *"Crame loss" -> "Cramer loss" in line 169*

**Response:**
In equation (1) and line 89, $r_t$ should not be bold. And we also have corrected the “Cramer loss” typo.

As for the architectural representation of the joint distribution, we also noticed that this point was not clear in our original manuscript. We would like to supplement it here:
First of all, as you mentioned, to avoid the curse of dimensionality, we use sample-based methods here instead of methods like C51 or QR-DQN that model the joint distribution directly. 

Furthermore, for the architecture using samples, we also tried several different designs at the beginning before fixing the architecture:

1. The architecture used in the current paper. To be specific, we replace the final layer of DQN to output some deterministic samples of n-d reward for each action(l.203 - l.206).
2. An IQN like architecture to model the joint distribution (as you mentioned in the comment).
3. A network architecture that concatenates a noise to the state feature, and outputs the deterministic samples.

We perform the same experiment as in Section 5.1 and Figure 2 in our paper to test how MD3QN models the joint distribution under these three network architectures. The results can be found in [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/network.png ]. For the first architecture, the MMD between prediction and ground truth distribution is 0.026, while for the second architecture, MMD is 1.541, and for the third architecture, MMD is 0.204. It can be concluded that the first one is the best in terms of modeling the joint distribution, and we adopt this architecture in our paper. 

-------------------------------------------------------------------------------------

3) *Whether or not a full joint distribution is really necessary for this setting.*

    *To compare with a baseline simply learn marginal distributions of the return associated with each reward.*

    *How the full distribution could improve HRA in other respects beyond aliasing, since it appears identical.*

**Response:**
Thanks for your advice! To further illustrate the necessity of learning a full joint distribution, we add another experiment during the rebuttal, and explain why only modeling marginal distributions may fail in some settings (e.g., in environments with multiple correlated constraints) from theory and experiment results.

There are a number of real-world scenarios where multiple constraints should be met simultaneously. For instance, in autonomous driving, we need to balance the safety (distance from other cars), the speed, the comfort (the acceleration, etc.), and many other factors to make sure the car functions normally. Those statistics can be viewed as sub-rewards in our settings. For instance, we would like to make the agent be aware of the joint distributions of distance from other cars and the speed, and add a constraint on this distribution (e.g., d > 1m, v < 50km/h).

From the theoretical perspective, only the algorithm which is capable of modeling the joint distribution can find this optimal solution. If the algorithm can only model the marginal distributions, we can correctly compute the probability of simultaneously meeting multiple constraints only if they are independent. However, this is not common in real-world scenarios.

We use the same Maze environment as in Figure 1 of our paper and modify the layout as a simplified analogues (see [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/constraint_maze.png ]). We set multiple constraints on the total return (in our experiment, we have three constraints: for the initial state, total red return > 0.6, total green return > 0.6, and total blue return > 0.6, where the agent gets a reward = 1.0 of the specific color after collecting a specific colored block), and the goal is to find a policy which has the highest probability to satisfy all constraints. 

We use the joint distribution by MD3QN to achieve this: specifically, given the modeled joint distribution $\textbf{Z}(s,a)$, the agent can compute the probability to satisfy all the constraint in the joint distribution and take action by $\arg \max_a \mathbb{P}(\textbf{Z}(s,a) \text{ satisfy all three constraint})$. We also modify the joint distributional Bellman optimality operator to maximize the probability to satisfy all constraints.

For baseline methods, we extend the MMDQN by multiple heads $Z_{1:N}$ to reflect $N$ sources of rewards. This method is different from our algorithm in that only marginal distributions of each dimension are learned. We test two ways to approximately compute the probability to satisfy all constraints by marginal distributions: the Marginal_SUM baseline maximizes $\sum_{i=1}^N \mathbb{P}(Z_i(s,a) \text{ satisfy i-th constraint})$, the sum of probabilities for each reward to satisfy constraint, while the Marginal_PROD baseline maximizes $\Pi_{i=1}^N \mathbb{P}(Z_i(s,a) \text{ satisfy i-th constraint})$, the product of probabilities for each reward to satisfy constraint. 

The results can be found here [ https://raw.githubusercontent.com/NeurIPS21-MD3QN/NeurIPS21-MD3QN-figures/main/multiple_constraint.png ]. It can be seen that maximizing the probability based on joint distribution can significantly outperform two baseline methods which both use marginal distribution information. 

-------------------------------------------------------------------------------------

4) *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?*

**Response:**
Thanks for pointing out the potential limitations! We agree that, given the network capacity, as the number of distinct reward signals increases, the agent may fail at some point. However, when the network capacity and learning process  is good enough, the upper bound of the numbers will increase. We think this is a tradeoff between network capacity and algorithm effectiveness. 


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). 
Experiments on four Atari games compared with HRA for multiple-reward functions. 

  
The definition of distributional RL is more complex than necessary. In eq 4, the a’ notation is overloaded. 

In 6) and 8), the bold fonts are not kept consistent. The reward notation in 8) is better consistent too. 

140: the function f has been defined. 

Theorem 1 and Prop. 1 are for policy evaluation. 

Theorem 2’s presentation is different from Th. 1, in that the metric is changed. So it does not hold for Wasserstein distance?

The 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. 

Can the authors explain the technical changes in the proofs because of the extension to hybrid reward setting? 

Using MMD for training distRL is straightforward and not new. 

Four 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/A. 

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