Paper: The Distributional Reward Critic Architecture for Reinforcement Learning Under Confusion Matrix Reward Perturbations (ICLR-2024)

Reviewer comment (Theory):
**Potential concept misuse**. The reward is a naturally random variable and considering its distribution in the corrected reward setting is reasonable. However, it does not mean it is strongly correlated with distributional RL, which focuses on the return distribution instead of its expectation: value function. I agree some density estimation techniques, such as categorical parameterization from C51 paper, can also be used here from distributional RL, but most of the technical issues, including the theoretical parts, between these two scenarios are very different. Emphasizing too much about the so-called distributional reward seems not natural.

Author response:
A1: We do not think there is a concept misuse. As the reviewer says, distributional RL focuses on the return distribution. In our methods, our distributional reward critic focuses on reward distribution. "Reward Critic" is to predict the reward instead of value functions. "Distributional" means we apply the idea of distributional RL to predict the distribution instead of a single value as RE does. We would be happy to provide clarification if we have misunderstood the comment.

Reviewer comment (Theory):
**Technical questions.** I am very confused by selecting the most probable reward label in step 7 of algorithm 1 instead of doing a reweighted sampling from the distribution. Note that if we only choose a single one, the optimal label tends to be deterministic rather than truly learning reward distributions. In the current version, along with the computation of the predicted reward value step, I think it is very similar to just doing the regression instead of a real multi-class classification since the $\arg\max$ does not consider the other probabilities when the reward belongs to other bins.

Author response:
A8: First, mode-preserving is the assumption of our methods. With mode-preserving, we can extract the most probable label as the predicted reward label. If we compute the expectation instead, it becomes a distributional version of RE, meaning we cannot recover the perturbations breaking the optimal policy. For the significance of the mode-preserving assumption, it is used by SR as well as shown in Table 1. Besides, mode-preserving is a weak assumption because our methods still work well under large noise ratios without a guarantee of mode-preserving under GCM perturbations and all the continuous perturbations people care about in RE necessitate mode-preserving shown in the Perturbations paragraph in Section 5.1. Moreover, GDRC even wins over RE under their targetted perturbations, showing mode-preserving is weak and GDRC is generally applicable. We clarify the motivation of mode-preserving in the last paragraph of Section 3 of the revised paper in blue.

Second, we show the performance of DRC by computing the expectation instead of extracting the mode in Figure 12 in Appendix G for reference, from which we can tell DRC almost always performs much better than DRC_EX.

Third, our method trains the distributional reward critic with observed rewards, learning the distribution of reward labels. When using rewards to train the critic and actor, we extract the reward label from the distribution using argmax and convert it into predicted reward values. Therefore, reward label distribution as the output from the distributional reward critic is learned.

Reviewer comment (Experiments):
**Computation issues.** Learning a confusion matrix mapping is OK in generative models, but learning a reward critic (similarly viewed as a confusion matrix mapping) can be costly in RL as RL algorithms typically lack sample efficiency. Based on this knowledge, determining $n_o$ via the training loss is normally intractable, especially in a large environment with huge state and action spaces. Last but not least, using the training loss is a typical cross-validation strategy, which is straightforward and even trivial for me.

Author response:
A11: Actually, lacking sample efficiency in RL contributes to the correct prediction from the distributional reward critic. Our methods for training the distributional reward critic use the same number of samples of standard RL training processes. According to Figure 8, the numbers of the samples used to train the distributional reward critic in P2 and used to train the critic and/or actor are the same. Indeed, the lack of sample efficiency in training an RL agent is really helpful for the training of our distributional reward critic. This is because many samples collected for training the agent are not that helpful in converging to the optimal policy quickly, but all the collected samples almost have the same value for training the distributional reward critic. To tell it from another perspective, the agent might be stuck in a local area of the state-action space for a long time, which is not good for policy improvement but good for the correct prediction of the distributional reward critic.

A12: As mentioned in A7, we believe the feasibility of our methods has been sufficiently demonstrated. One thing to notice is that we do use the CE, but we use it wisely and differently from what people always do. Using the CE loss to choose an appropriate $n_o$ has been proved through four steps as discussed in A7. To address the limitation of recovering the confusion matrix for SR in large environment spaces, we incorporate the distributional reward critic, a neural network, to resolve this issue. There is nothing to do with cross-validation regarding how we use the CE. We want the reviewer to clarify the relation between using the CE loss for cross-validation and our work.

Reviewer comment (Theory):
For continuous perturbations, the discretization technique should always induce a reconstruction error, which seems to have conflicts with some statements in the paper.

Author response:
A13: All of the analysis of our methods analysis occurs under the GCM perturbations as clearly stated in the method section. Continuous perturbations are not our targeted ones in designing the methods, where we want to demonstrate the broad applicability of our approaches and the weakness of the mode-preserving assumption by showing their general application even under untargeted perturbations.

Reviewer comment (Writing):
The writing can also be improved as there are many grammar errors and typos.

Author response:
A14: We will revise the paper carefully as suggested. We welcome feedback from the reviewer on areas requiring corrections, allowing us to make targeted improvements.

Reviewer comment (Theory):
The major potential weakness of the paper is that it has a restrictive type of perturbation. The paper assumes the rewards in different intervals are shuffled and the true reward is still the most frequently observed (mode-preserving property). However, this assumption appears too artificial and not natural, especially given that while the paper claims the assumption of RE (Romoff et al. 2018) to be stringent, the RE assumption seems more natural. Although RE can’t properly handle non-order-preserving perturbation, further justification for the mode-preserving assumption is needed.

Author response:
A2: Mode-preserving is the assumption of the surrogate reward (SR) method instead of being created by us as shown in Table 1. When recovering the confusion matrix, they use majority voting to decide the true reward of the samples for a state-action pair assuming mode-preserving. Besides, the forms of $\tilde{R}$ of the continuous perturbation functions (e.g. Gaussian perturbations) people pay attention to in RE necessitate mode-preserving although the space is continuous, which shows mode-preserving is a weaker assumption in the cases people are interested in. Following this idea, we propose a method that combines the advantages of both SR and RE under the assumption of mode-preserving.

One thing to notice is that we do not guarantee mode-preserving in the generation of GCM we experiment with. Nevertheless, our methods perform effectively even under high noise ratios (e.g., $\omega =0.7$), demonstrating that our assumption of mode-preserving is weak and our methods are robust. Furthermore, our method outperforms baseline methods even under the perturbation (continuous perturbation) tailored for RE shown in Section 5.4, indicating our method's broad applicability and mode-preserving is a weak assumption again.

In a word, with a weak mode-preserving assumption, our methods are more broadly applicable under many kinds of perturbations. We clarify the motivation of mode-preserving in the last paragraph of Section 3 of the revised paper in blue.

Reviewer comment (Writing):
Clarification is needed on what SR_W is and how it differs from SR.

Author response:
A3: SR can learn the confusion matrix within a simple environment space (e.g. Pendulum), but SR_W needs to know the whole confusion matrix before being applied, which is the last sentence of the first paragraph of Page 3. We uniform the use of the reference style and modify the notation as suggested.

Reviewer comment (Presentation):
The reference style is a bit messy and inconsistent throughout the paper.

Author response:
A3: SR can learn the confusion matrix within a simple environment space (e.g. Pendulum), but SR_W needs to know the whole confusion matrix before being applied, which is the last sentence of the first paragraph of Page 3. We uniform the use of the reference style and modify the notation as suggested.

Reviewer comment (Theory):
One desirable property of any approach to this problem is that it should never decrease performance on tasks where the rewards are *unperturbed*. However, this doesn't hold for the proposed approach. For example, consider a simple betting game where the reward of some bet is -1 with probability 0.9 and +10 with probability 0.1. Romoff et al.'s approach will (in theory) learn that the expected reward is 0.1, and thus the agent will learn that the bet is worth taking. On the other hand, the proposed approach will learn that the mode reward is -1, and hence the agent will learn to avoid the bet. Generalising from this, if we want the approach to be "safe" then it ought to preserve the expected reward.

Author response:
A1: It is overly stringent to expect a method designed for dealing with perturbations to not suffer from any performance loss in a clean environment. If there is always no degradation in performance for handling perturbations, there is no need for the original method.

In the definition of MDP in the first paragraph of Section 3, the reward $r_t$ is from states and actions $(s_t, a_t)$ expressed as $r_t = R(s_t, a_t)$, but the reward $r_t$ is also allowed to depend on the transition, meaning $r_t = R(s_t, a_t, s_t^\prime)$. To make the betting game MDP, we need to include $s_t^\prime$ as the input to the reward function to tell whether we win or lose if we choose to play the betting machine. Otherwise, it even does not accord with the definition of MDP as the same state-action pairs can generate different rewards. To handle the betting game, there is only a small adaption needed – input $s_t^\prime$ to the distributional reward critic other than $(s_t, a_t)$. Therefore, our methods also work in the betting game.

Reviewer comment (Experiments):
The difference between the proposed approach and Romoff et al.'s can be broken into two parts: (1) Learning a categorical estimate, rather than a scalar one; (2) The agent is trained from the mode reward, rather than the mean. I'd like to see a comparison versus an approach that uses (1) but not (2), i.e., the reward distribution is learned, but the agent is trained from the mean of the distribution. This would disengangle how much of the performance impact comes from assuming a categorical distribution and how much comes from assuming that the perturbation is mode-preserving.

Author response:
A3: This is a good question. If we calculate the expectation, it essentially becomes a distributional version of RE, making it impossible to recover the perturbation where the assumption of a positive affine transformation does not hold. We include the experiment results in Figure 12 in Appendix G for verification, from which we can tell DRC almost always performs much better than DRC_EX. Therefore, we conclude extracting the mode is the most important factor contributing to the performance improvement, instead of a categorical distribution fitting.

Reviewer comment (Theory):
Page 8 states: "The agents in Hopper and Walker2d are encouraged to live longer because of the positive expectation of perturbed rewards". However, this contradicts the statement on page 2 that affine transformations preserve optimal policies. The statement on page 2 should be refined to say that this only holds for non-episodic tasks.

Author response:
A4: There is no contradiction. The two statements are under different perturbations. Under GCM perturbations, the statement on Page 8 shows the agent can still get some information from the expectation of fully perturbed reward signals because the perturbed reward range stays unchanged after the perturbation. It is to explain the not-too-bad performance of the baseline methods under GCM perturbations. Under continuous perturbations, the statement on Page 2 is about the assumption of RE expressed as $\mathbb{E}(\tilde{r})=\omega_0\cdot r +\omega_1(\omega_0 > 0)$, which ensures the unchanged optimal policy for both episodic and non-episodic tasks after the perturbation.

Reviewer comment (Theory):
In my opinion, the proof of Theorem 1 is unnecessary and just overcomplicates the paper. All it really says is "under perfect conditions, the cross entropy loss will go to zero, and hence the approach learns what it's supposed to". We'd have AGI already if we had "sufficiently expressive" neutral networks + unlimited training samples + the necessary compute :p

Author response:
A5: The reviewer’s understanding is partially incorrect. In typical classification problems, an input (e.g., an image) has a definite label. However, samples of the same state-action pair can have varying labels in our perturbed environment. It is difficult to prove the minimum of the loss function for all the samples directly so we wisely pick the samples of a state-action pair and prove the minimum of the loss function for these samples. Our goal is to demonstrate that the output of the trained distributional reward critic corresponds to the distribution of perturbed reward labels by using Cross-Entropy (CE) loss. The reviewer misunderstand the CE will go to zero as the CE does not converge to zero because of perturbation, evident from the proof of Theorem 1, Figures 2, and Figure 4, where the finally converged output is the distribution of perturbed reward labels. It might cause such misunderstandings if we do not emphasize it enough.

Reviewer comment (Theory):
I was slightly confused by the statement on page 3: "we assume that the distribution of the perturbed reward depends only on the true reward, i.e., $\tilde{r}_t \sim \tilde{R}(R(s_t, a_t))$". Couldn't the proposed approach handle state- and action-dependent perturbations too?

Author response:
A6: Yes, our method can address state- and action-dependent perturbations, provided that the mode-preserving assumption holds. The reward-dependent (state-action dependent) perturbation is just the perturbation we experiment with for evaluation. How the perturbation is constructed does not impact our method at all. We experiment with this kind of perturbation because it is commonly used and the general confusion matrix (GCM) perturbation is general.

Reviewer comment (Theory):
Though the proposed method avoids some of the drawbacks of previous works, such as directly estimating the confusion matrix or the true reward function, or assuming order-preservation of the perturbation, the mode-preserving condition seems quite restrictive. Mode-preserving perturbations are not a generalization of order-preserving perturbations, since the latter may not be order-preserving. Also, methods that learn the confusion matrix may be more widely applicable (though more computationally expensive), since the perturbations can be quite general. Without further motivation of the mode-preserving condition and its advantages over other conditions imposed in the literature, the broader impact of this work is difficult to judge.

Author response:
>Q1: Order-preserving perturbations might not be mode-preserving and what the motivation of the mode-preserving assumption is.

A1: Order-preserving is necessary for RE to work but is not sufficient. After analyzing all the methods in detail, we find the assumption of RE is $\mathbb{E}(\tilde{r})=\omega_0\cdot r +\omega_1(\omega_0 > 0)$, a positive affine transformation. That means RE might not work if order-preserving holds although mode-preserving might not hold as well. Here we clarify the difference between mode-preserving and order-preserving. Mode-preserving and order-preserving are two kinds of concepts although they look similar. Mode-preserving is the assumption of our methods and order-preserving is a property assumptions might necessitate. The assumption of RE is positive affine transformation instead of order-preserving. It is not that clear which method is more broadly applicable with different assumptions. Then we find positive affine transformation necessitates order-preserving, making RE not work under non-order-preserving cases as we discuss in the last paragraph of Page 2. However, mode-preserving does not necessitate order-preserving, which means our methods can work in cases where the reward order is totally broken, and then the optimal policy gets changed as long as mode-preserving holds, demonstrating the larger applicability and capability of our methods. That is the reason we have a separate column for order-preserving in Table 1 to help understand the applicable scenarios of different methods. To ensure there is no confusion, we change the term from "order-preserving" to "optimal-policy-unchanged", which is comparable to the positive affine transformation and is easy to tell the applicable scenarios as well. Besides, mode-preserving is not created by us, which is the assumption of SR as shown in Table 1. When recovering the confusion matrix, they use majority voting to decide the true reward of the samples for a state-action chunk while assuming mode-preserving.

One more thing to notice is that the forms of $\tilde{R}$ of the continuous perturbation functions (e.g. Gaussian perturbations) people pay attention to in RE necessitate mode-preserving although the space is continuous, which shows mode-preserving is a weak assumption in the cases people are interested in. With mode-preserving, our methods win over the baseline methods under both GCM perturbations and continuous perturbations, which is the strongest evidence that mode-preserving is a weak assumption and our methods are broadly applicable.

We clarify the motivation of mode-preserving in the last paragraph of Section 3 of the revised paper in blue.

>Q2: Methods that learn the confusion matrix may be more widely applicable since the perturbations can be quite general.

A2: Although the general confusion matrix (GCM) perturbation is general, the method (SR) learning the confusion matrix is not widely applicable.

As we analyze in Table 1 and the last paragraph of Page 4, SR requires some unrealistic information about GCM settings, making it inapplicable to perturbations without the concept of a confusion matrix like continuous ones. Besides, although a correctly recovered confusion matrix tells us how perturbations are applied, it is not the best solution to have the fixed substitute reward for a corresponding observed reward using $\hat{R}=C^{-1}\cdot R$ because they might be perturbed from different clean rewards, which ignores the valuable state-action information and could even cause significant variance compared with the one caused by perturbations. Last but not least, SR discretizes the state-action space and regards the clean rewards of a given state-action chunk as the same without any assumption of the smoothness / Lipschitz-ness of rewards, making itself incapable of dealing with the environments with large state-action space and complicated reward functions.

Reviewer comment (Experiments):
The experimental comparisons are extensive, but there are some issues regarding the fairness of comparisons with previous methods as well as their interpretability. Due to the lack of order-preservation mentioned at the end of the first paragraph of Sec. 5.3, the results in Fig. 5 and 6 may be unfair to the RE method. Comparisons on problems that enjoy conditions to which all methods are applicable would provide a fairer comparison.

Author response:
>Q3: The results in Figures 5 and 6 may be unfair to the RE method. Comparisons on problems that enjoy conditions to which all methods are applicable would provide a fairer comparison.

A3: As mentioned in A1, a positive affine transformation is the assumption of RE instead of order-preserving and mode-preserving is what we need to work under all the perturbations people are interested in of both RE and SR. For Figures 5 and 6, the generation of GCM perturbations does not ensure any condition including mode-preserving as we say in the Perturbations paragraph of Page 7, meaning they are fair comparisons for all the methods. Additionally, we have evaluated the methods under continuous perturbations of RE with a positive affine transformation and found that our method still performs better although they are even not our targets, as discussed in Section 5.4.

Reviewer comment (Evaluation):
Though the proposed DRC and GDRC methods show better performance than previous methods on a number of problem instances, they decidedly underperform on many others. Further discussion of this discrepancy is warranted to clarify when the proposed methods should be preferred.

Author response:
>Q4: While the proposed DRC and GDRC methods show improved performance on several problem instances, they underperform on many others.

A4: Our methods generally win over the baseline methods in all kinds of perturbation settings (e.g., environments, noise ratios, and reward discretizations) except for HalfCheetah with low noise ratios under GCM perturbations. Even under non-GCM perturbations (continuous ones), GDRC still has an edge over RE.
For HalfCheetah, we analyze the underperformance of our methods and provide some feasible solutions. In the last paragraph of Section 5.3, we have detailed discussions by analyzing the metrics of DRC in Hopper, DRC in HalfCheetah, and GDRC in HalfCheetah shown in Figure 10. Correspondingly, we have proposed potential solutions in the first paragraph of our Future Work section. As this is the first work introducing the distributional reward critic methods (DRC and GDRC) to address perturbed rewards and because of space constraints, we leave further improvement in the future work.

Reviewer comment (Experiments):
The impact of \(n_o\) studied in Sec 5.2 is an important issue, but the meaning of Fig. 3 is difficult to interpret. This leaves the question of how to choose \(n_r\) open, which is a key component to the applicability of GDRC, in particular.

Author response:
>Q5: The impact of $n_o$ studied in Section 5.2 is crucial, but the meaning of Figure 3 is difficult to interpret.

A5: Yes, choosing $n_o$ is an important problem to solve, which is the transition step from DRC to GDRC. There are two forces in Figure 3: one is the reconstruction error turns zero when $n_o$ is a multiple of $n_r$; the other one is overfitting becomes worse as $n_o$ increases because samples are not infinite.
For small $n_r$ like 5 or 10, the performance of DRC reaches the best when $n_o=n_o$, but it is almost always worse when $n_o$ is a multiple of $n_r$. For large $n_r$ like 20, the performance of DRC decreases even when $n_o$ is smaller than $n_r$. Both of the two phenomena above from Figure 3 support what we proposed by two forces before.

In the method section before any experiments, we discuss the tradeoff between the two forces in the first paragraph of Page 6 and decide to choose when $n_o=n_r$, which is further proved correct in Section 5.2. Moreover, we discuss some potential solutions to further weaken the assumption of unlimited samples in the Future Work section.

We clarify what we want to convey in the second paragraph of Section 5.2 of the revised paper in blue.

Reviewer comment (Evaluation):
The proposed method replaces the computational effort of estimating the confusion matrix with the computational cost of training a classifier estimating the bin to which the reward at a given state-action pair belongs; can you comment on how these two costs compare?

Author response:
>Q6: Comment on how the costs of estimating the confusion matrix and the reward distribution compare.

A6: For time complexity, the cost difference between the two methods is negligible compared to the inherent training expenses of the critic and actor. For sample complexity, the samples we need to make the two methods work increase as $n_r$ increases. In the case of the known reward range and $n_r$, the sample complexity of SR and DRC is the same as the two methods both need enough samples to extract the mode reward. In the case of unknown discretization of the perturbation, the sample complexity of SR and GDRC is incomparable as SR cannot apply in the unknown cases.

Reviewer comment (Theory):
Does the classifier trained in line 7 of Alg. 1 -- which is called the "distributional reward critic" -- correspond to any specific action value function, like the standard critics used in actor-critic algorithms? If not, it is a little misleading to call it a critic.

Author response:
>Q7: Does the classifier trained in line 7 of Algorithm 1, referred to as the "distributional reward critic," correspond to any specific value function?

A7: No, it is a reward critic to predict rewards instead of values, but the inputs are also state-action pairs. As we turn the regression problem into a classification one, it is a distributional reward critic. If you go to Figure 8, the pipeline of our methods from a high perspective, you can tell the classifiers are processed in P2, which is a different process from Value Fitting.

Reviewer comment (Theory):
What is the main difficulty in proving Theorem 1? It seems like it should follow from a straightforward application of the law of large numbers.

Author response:
>Q8: What is the main difficulty in proving Theorem 1? It seems like it should follow from a straightforward application of the law of large numbers.

A8: The primary challenge lies in the fact that an input (e.g., an image) has a definite label in typical classification problems. However, in our perturbed environment, samples of the same state-action pair can have varying labels. It is difficult to prove the minimum of the loss function for all the samples directly so we wisely pick the samples of a state-action pair and prove the minimum of the loss function for these samples. Our goal is to demonstrate that the output of the trained distributional reward critic corresponds to the distribution of perturbed reward labels by using Cross-Entropy (CE) loss. Indeed, the law of large numbers applies to the assumption we say as infinite samples.

Reviewer comment (Presentation):
What should the reader take away from Fig. 3? It is currently tough to interpret.

Author response:
>Q9: What should the reader take away from Figure 3?

A9: As mentioned in A5, the tradeoff between the two forces, the relation between the reconstruction error and $n_o$ and overfitting issues as we increase $n_o$, we propose in the first paragraph of Page 6 under the method section is proved in Section 5.3, meaning our strategy to choose $n_o$ concluded from the simulation environments is correct in experimental environments.

Reviewer comment (Reproducibility):
At the end of the **Algorithms** subsection of Sec. 5.2, it says the rewards are averaged over "10 seeds and 20 random trials"; independent trials are usually each associated with one seed -- can you clarify what you mean?

Author response:
>Q10: What does the use of 10 seeds and 20 random trials imply?

A10: We train the agents using 10 different seeds. For each seed, we run the agent for 20 episodes after the training for evaluation.