Correlated Proxies: A New Definition and Improved Mitigation for Reward Hacking

We thank the reviewer for their thorough review and many helpful suggestions. We are glad that they found the paper to be "well-written", and that it "provides an interesting perspective on reward hacking." Below, we have addressed individual points in the review:

• Comparison to preference-based RL: The reviewer suggests that in some environments preference-based RL (PBRL) may be able to outperform regularized policy optimization. We believe that a comparison to PBRL is out-of-scope for our problem setting, since we are interested in how to optimize a misspecified proxy reward without further input. In contrast, PBRL generally requires multiple rounds of feedback from a human annotator during learning. However, regularized policy optimization is often a key subcomponent of PBRL, since once a batch of preference data is collected, it is often optimized using a regularized objective, as in RLHF. Thus, we believe that our results may be useful in improving PBRL algorithms.
• Comparison to inverse reward learning: The reviewer also suggests comparing our algorithm to inverse reinforcement learning (IRL). However, we note that IRL is generally used for producing a reward function that can be used to imitate demonstrations, and thus, if we applied IRL to the base policy, we would simply recover equal performance to it. In contrast, we aim to improve performance over the base policy using a reward function which is misspecified but correlates with the true objective.
• OM-based regularization for RLHF: We appreciate the reviewer bringing this important reference to our attention. Nika et al. [1] do consider an objective with log-linear occupancy measure regularization. However, their analysis is limited to deterministic MDPs, and they do not evaluate their approach empirically. In contrast, we show both in theory and in our experiments, that our approach to preventing reward hacking can succeed in general stochastic environments. We have added a reference to [1] in our related work.
• Estimating occupancy measures in high-dimensional state spaces: Our approach to estimating occupancy measure divergence in ORPO is based on using a learned discriminator, which has succeeded in many similar applications (e.g., GANs [2] and GAIL [3]). We test ORPO in environments with state space dimensions ranging from 1 (glucose monitoring) to 50 (traffic control), showing that it is effective in high-dimensional spaces. Furthermore, we have added Figure 6 to Appendix B showing that the learned discriminator for the glucose monitoring environment is close to the exact log ratio of the occupancy measures, validating that our approximation of occupancy measure divergence is accurate.
• Tuning of hyperparameters: We found that in some environments, if the RL hyperparameters were not tuned, PPO would simply fail to optimize the proxy reward. In these cases, reward hacking did not occur, while PPO also failed to improve the true reward. To study various techniques to prevent reward hacking, we tuned PPO such that it can consistently optimize the proxy reward, leading to reward hacking in all environments.

[1] Nika et al. Reward Model Learning vs. Direct Policy Optimization: A Comparative Analysis of Learning from Human Preferences. ICML 2024.

[2] Goodfellow et al. Generative Adversarial Nets. NIPS 2014.

[3] Ho and Ermon. Generative Adversarial Imitation Learning. NIPS 2016.

We thank the reviewer for their thorough review. We are glad they found our paper "provides some very interesting insight." We have responded to individual points in the review below:

• Can the lower bound in Theorem 5.1 be optimized? We agree that it would be ideal if we could guarantee that the lower bound on the true reward in Theorem 5.1 can be increased above 0. However, it turns out to be quite difficult to prove this in general; we have added Lemma A.3 in Appendix A.1.3 showing that there are MDPs for any correlation coefficient $r$ where it is impossible to increase the true reward by optimizing the lower bound. This holds true even when there is a policy that improves both true and proxy reward. However, the advantage of the lower bound in Theorem 5.1 is that it always allows safe optimization of the proxy reward: even if it is not possible to increase the lower bound above zero, optimizing the objective will at least prevent reward hacking. If we are able to increase the objective, then we know that the true reward must have increased too. Our empirical results suggest that in many realistic environments it is possible to increase the lower bound in Theorem 5.1.
• "There is no claim about whether or not reward hacking can (at least in principle) be prevented with action distribution regularization": We have added Theorem A.5 to Appendix A.2 (and referenced it from the main text) showing that action distribution regularization cannot achieve the same guarantees as occupancy measure regularization. Theorem A.5 is very general: it shows that regularizing using any $f$-divergence between action distributions scaled using any function cannot prevent reward hacking.
• EPIC distance: The EPIC distance could serve as a formalization of what makes a "good" proxy. However, there are multiple barriers in applying it to the environments we study. First, it only applies to reward functions of the form $R(s, a, s')$ that depend on the next state as well as the current state and action; none of the environments we explore have this type of reward function. Furthermore, it requires independent state and action distributions to use for computing the canonically shaped reward, whereas we use occupancy measures where the state and action distributions are not independent. For these reasons, we do not use it in our study. However, we have added a reference to the related work.
• RLHF as a contextual bandit problem: We agree that RLHF can also be formulated with each generated token representing a separate action. We have added Lemma A.7 in Appendix A.3 showing that in this case, action distribution and occupancy measure KL divergence are still equivalent.
• Table 1 cell for traffic control with state occupancy KL: This is not a typo; RL is quite unstable in the traffic environment with state-only KL occupancy regularization. This is likely because the reward function in the traffic environment depends on both the state and action, so our theory suggests state-only regularization will not work--we only include it for completeness. As shown in Table 4 in the Appendix, a slightly larger regularization coefficient achieves true reward of -1.47 ± 0.20, which is a slightly lower median but with much less variance.
• First example of KL regularization to an "offline policy": We agree that maximum entropy RL can be understood as KL regularization to the static uniform policy. However, to our knowledge this is not a common tool used to prevent reward hacking. Instead, it usually aids in exploration or is used to model human behavior. We have clarified in the related work that Stiennon et al. [1] are the first to consider KL regularization to a fixed policy in the context of reward hacking.
• Related work in offline RL: We have added the Fujimoto et al. reference to our related work, and we additionally reference several other works on offline RL there.
• Definition 4.2 ("reward hacking"): We agree that this may not be the best term for the phenomenon. Based on the feedback in this review and that of Reviewer qUra, we have updated the definition to instead refer to the reward function as "hackable" if the condition is met.
• Narrative of the paper: We appreciate the reviewer's suggestions for improving the narrative of the paper. Our choice of narrative was based on previous rounds of reviewing in which we received comments like "it is not entirely clear what the authors mean by the term reward hacking" and "this work lacks a global view and theory on the problem." Because of these comments, we believe that the extended introduction and discussion of the definition of reward hacking is helpful for a broad audience to understand and appreciate the paper.
• Other notes: We have fixed minor issues in an updated version of the paper that we uploaded to OpenReview.

[1] Stiennon et al. Learning to summarize from human feedback. NeurIPS 2020.

Thank you for responding to our rebuttal! Regarding whether the lower bound in Theorem 5.1 can be positive, we have added Lemma A.4 to Appendix A.1.3. This lemma shows that for any correlation coefficient $r \in (0, 1)$, there is an MDP where the lower bound in Theorem 5.1 can be positive, and additionally, in this MDP optimizing the lower bound leads to a policy with the optimal true reward. We hope this addresses your concern.

We thank the reviewer for their very thorough review and are glad that they believe the paper is “nicely presented and a nice idea.” Below, we have addressed some of their questions and concerns:

• Is the state-occupancy distribution induced by a "natural" base policy close enough to optimal? The reviewer notes that "the optimal policy wrt the true reward might be quite different from what we think of as the 'natural' base policy." We agree that even if the proxy reward is $r$-correlated with the true reward, it could be that the base policy is still quite far from the optimal policy. In this case, it may not be possible to achieve optimal performance under the true reward. However, we argue that this is often not the goal: if all we have is a proxy objective, then it may be unreasonable to expect to achieve optimal (or even near-optimal) performance by optimizing the proxy. Instead, as we describe at the beginning of Section 5, we are interested in whether we can "optimize a proxy reward function and have it translate into an improvement in true reward over the base policy." Our Theorem 5.1 shows that this goal of increasing true reward relative to the base policy is possible as long as the proxy and true rewards are correlated under the base policy.

  If one cares about finding a near-optimal policy under the true reward, one could combine Theorem 5.1 with an assumption that the base policy achieves reward that is not too far from optimal. We have added such a result as Corollary A.1 in Appendix A.1.1.

• Why do you only consider the optimal policy wrt the proxy in Definition 4.2? We agree that this is a restrictive definition and that in practice the policy is not fully optimized with respect to the proxy reward function. The reason we originally referred to the optimal policy is because it is mathematically precise, while notions of "partially optimizing with respect to a reward function" are harder to formally define. We have updated the definition of "reward hacking" to describe situations when optimizing the proxy (even partially) leads to a decrease in true reward compared to the base policy. Furthermore, a "hackable" proxy reward is one whose optimal policy exhibits reward hacking. Thus, this revised definition now covers cases where the proxy reward is only partially optimized.

• The paper doesn't say how to know whether, and what to do if the proxy reward function is not r-correlated with the true one. We agree that this is a limitation of our paper. While we showed that in several realistic environments that the $r$-correlated proxy assumption is satisfied, this may not be the case in all environments. However, in the case of learned reward functions in particular, we have added Lemma A.8 to Appendix A.4, where we show that under basic assumptions, a learned reward function will be $r$-correlated with the true reward. This gives further theoretical evidence, in addition to our empirical evidence, that many proxies are $r$-correlated with the true reward.

• How do you expect your technique to scale? Since our theoretical results in Theorem 5.1 apply to any method of optimizing the regularized proxy reward on the right side of equation (4), we believe that our technique for preventing reward hacking should scale well to larger models. Furthermore, our experimental results cover both very small networks (e.g., a 2-layer MLP in the pandemic environment) and large language models (Pythia-1.4B for our RLHF experiments). This also suggests that our technique is quite scalable.

• KL divergence penalty in PPO: The reviewer writes that "KL divergence, to my understanding, is primarily used to avoid instabilities in the training procedure [of PPO]." It is certainly true that in PPO there is usually a KL penalty between the current and previous iterations' policies to prevent instability. However, RLHF uses an additional KL penalty to avoid reward hacking. While the first KL penalty is between the current and previous PPO policy (which changes each iteration), the second KL penalty is between the current PPO policy and the SFT policy (which does not change). In Stiennon et al. [1], the authors find that removing this second KL penalty leads to overoptimization, producing nonsensical results (see Appendix H.2 of their paper), even though training is still stable (see Figure 5 of their paper).

• Other suggestions: We thank the reviewer for their extensive suggestions about how to clarify our writing. We have incorporated these suggestions in an updated version of the paper uploaded to OpenReview.

[1] Stiennon et al. Learning to summarize from human feedback. NeurIPS 2020.

Thank you for responding to our rebuttal. We had already addressed a few of the comments you mention, and have now uploaded a new version of the paper with fixes for the remaining ones:

• Use of the word "superior": we used this word as a synonym for "better." However, we understand that it has a stronger connotation, so we have now changed uses of "superior" to "better" or to saying that $\chi^2$ occupancy measure regularization "outperforms" KL action distribution regularization rather than saying it is "superior."

• SEIR and glucose diagrams are too small: we have now changed these diagrams in Figure 2 to be more readable.

• Use of the word "misaligned": we thought that this word was standard in the reward hacking literature [1, 2], but have removed it from the paper.

• Reference to the SEIR model: we removed the reference to Mwalili et al. (2020). We could not find a good original source for the SEIR model—for example, the Mwalili et al. paper does not cite a source for SEIR. It appears that the model has its origins going back more than a century and in most SEIR-related papers there is no given citation. Let us know if you believe there is a better way of handling this reference.

• RHS acronym: we had already changed this to define the acronym at its first use: "the right-hand side (RHS) consists of two terms."

• Purpose of KL divergence in PPO: we are glad that the reviewer finds the updated wording better. It's true that Stiennon et al. note that the KL to the SFT policy also serves as an entropy bonus. We do think the point that "it ensures the policy doesn’t learn to produce outputs that are too different from those that the reward model has seen during training" is addressing reward hacking, particularly since without the penalty they get clear reward hacking behavior (see their Appendix H.2). Also, follow-up work to their paper [3] has more clearly shown how KL divergence prevents reward hacking in RLHF. We definitely agree that there is not a clear delineation between "overfitting" to the learned reward function versus "reward hacking" and that stability is closely related as well. Thank you for bringing up this interesting discussion!

[1] Zhuang and Hadfield-Menell. Consequences of Misaligned AI. NeurIPS 2020.

[2] Pan et al. The Effects of Reward Misspecification: Mapping and Mitigating Misaligned Models. ICLR 2022.

[3] Gao et al. Scaling Laws for Reward Model Overoptimization. ICML 2023.

We thank the reviewer for their review and are glad that they found the paper "very well written" with "a wide range of experiments that show both the soundness and robustness of the claims." Below, we have addressed the weaknesses and questions raised in the review:

Assumption that the state-occupancy measure of the policy is absolutely continuous with respect to the base policy: Theoretically, we need this assumption because otherwise the $\chi^2$ divergence is undefined (there is a division by zero in the definition). However, in practice, we believe it is not a significant problem for a few reasons. First, in all experiments we initialize training with the base policy. Thus, at the beginning of training, the occupancy measure of the RL policy is clearly absolutely continuous w.r.t. that of the base policy. Second, during training, the $\chi^2$ penalty strongly disincentivizes reaching state-action pairs that have very low probability under the base policy's occupancy measure, since the $\chi^2$ divergence increases with $1 / \mu_{\pi\_\text{base}}(s, a)$. Therefore, the penalty is likely to push policy optimization away from areas of the state-action space where $\mu_{\pi\_\text{base}}(s, a)$ is very small, making it unlikely that the RL policy $\pi$ will reach states where $\mu_{\pi\_\text{base}}(s, a) = 0$.

Third, in practice we do not exactly calculate the $\chi^2$ divergence but approximate it with a discriminator. Suppose the RL policy $\pi$ does reach a state with zero probability under the base policy. In this case, while the exact $\chi^2$ divergence would be undefined, in practice the discriminator will assign an extremely high (but not infinite) score to the state, strongly penalizing it and causing RL to avoid it.

Thus, since the RL policy begins with an occupancy measure that is identical to that of the base policy, and it is strongly penalized for reaching state-action pairs with low or zero probability under $\mu_{\pi\_\text{base}}$, we believe that is it is likely that $\mu_\pi$ will remain absolutely continuous w.r.t. $\mu_{\pi\_\text{base}}$ during training.

The definition of correlated rewards allows for the correlated reward to deviate arbitrarily in states with zero occupancy by the base policy. We agree that our definition allows for arbitrarily large differences between the proxy and true reward functions at states that are unreachable by the base policy. However, as we argue above, those states are in practice almost infinitely penalized during training, so we believe that our method will still avoid reward hacking in such cases.

We thank the reviewer for their review and appreciate that they found that our method "demonstrates impressive effectiveness." Below we have responded to individual questions.

• Does the paper assume that both the base policy and the optimized policy are able to acess the same observation states? Yes, we assume that both policies receive the same inputs.

• Handling of unobserved variables in the traffic domain: We agree with the reviewer that unobserved variables in driving policies trained via imitation learning can cause problems. However, the traffic domain we consider is primarily focused on modeling traffic—the high-level interactions of large groups of vehicles—rather than individual drivers. In this case, models of individual driving do not need to be as high-fidelity. Furthermore, the traffic simulator we use is based on SUMO [1], which has been validated for large-scale studies of traffic, such as the VABENE project [2] that was used to inform traffic control during the FIFA World Cup in 2006.

• During the training phase, does the agent have access to the ground truth reward value? The agents do not have access to the true reward $R$ during training, only the proxy reward $\widetilde{R}$. Our setting is challenging because we evaluate based on the unobserved true reward but only allow policy optimization using a specified proxy reward.

• What are major assumptions behind your framework? The main assumption that enables our approach to preventing reward hacking is Definition 4.1 in the paper: that the proxy reward is correlated with the true reward over state-action pairs sampled from the base policy. We show that this assumptions holds across four realistic environments and a representative gridworld (see Figure 2).

• Unknown true reward: We agree with the reviewer that it is difficult to tune hyperparameters, such as the regularization strength $\lambda$, when the true reward is unknown. One option for evaluation in this setting is to learn or specify two different proxy rewards, one for training and one for evaluation. This approach was used for evaluating RLHF-fine-tuned models by Gao et al. [3]. In fact, since the true reward function for RLHF is unknown, in our experiments we leverage a more robust learned reward function as the true reward $R$ and a less robust learned reward function as the proxy reward $\tilde{R}$.

• Alternative metrics for RLHF: We have additionally computed win rates for both KL and $\chi^2$ regularization in RLHF:

  Divergence	Median win rate
  $\\chi^2$	52.83
  KL	51.50

  To compute the win rates, we compared model outputs from the RLHF-trained models to the SFT baseline using GPT 4o-mini and reported the median rate across five seeds. We see that $\\chi^2$ regularization results in a higher win rate, reflecting our other results. We have included these results in Appendix E.1 of the paper as well.

• Similarity to prior work: In comparison to the work of Laidlaw et al. [4], we present a more theoretically complete framework and additionally experiment with $\\chi^2$ regularization. In particular, [4] does not have a result like our Theorem 5.1 showing that regularized optimization of the proxy reward can lower bound the true reward. Furthermore, [4] only explores the use of KL divergence for regularization, while we additionally experiment with $\\chi^2$ divergence, which our theoretical results suggest is more principled. We find in practice that regularization with $\\chi^2$ divergence generally outperforms KL regularization.

• $\\chi^2$ vs. KL divergence: Empirically, we found that across all our experiments, $\\chi^2$ regularization performed similarly to or better than KL divergence. Even when KL divergence performed similarly to $\\chi^2$ divergence, we often found that $\\chi^2$ divergence was more stable across regularization coefficients. Combined with our theoretical results that provide guarantees regarding $\\chi^2$ divergence, this suggests that $\\chi^2$ divergence should be used instead of KL divergence when regularizing to prevent reward hacking.

[1] Krajzewicz, D. (2010). Traffic simulation with SUMO–simulation of urban mobility. Fundamentals of traffic simulation, 269-293.

[2] Krajzewicz et al. (2012). Recent development and applications of SUMO-Simulation of Urban MObility. International journal on advances in systems and measurements.

[3] Gao et al. Scaling Laws for Reward Model Overoptimization. ICML 2023.

[4] Laidlaw et al. Preventing Reward Hacking with Occupancy Measure Regularization. ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems.

While the initial version of the paper is not perfect, I believe the updated version is better. Especially with the newly added theoretical results, comparing action distribution regularization vs. occupancy measure regularization; and the learned reward can be r-correlated. However, I am uncertain whether adding ten new pages is standard practice for such revisions.

Still, regarding the newly added proofs, does it mean that assumption "the proxy reward is correlated with the true reward over state-action pairs sampled from the base policy" is required for Theorem 5.1, Corollary A.1, and Lemma A.2? If this assumption is indeed required, it would be beneficial to include experimental evidence or a practical demonstration to substantiate these theoretical claims.

Another concern: Could you elaborate on how these results differ fundamentally from those in [3]. The current proof seems to be easily extended from [3] with Cauchy-Schwartz inequality and commonly known entropy bounds.

Additionally, many common divergences, such as KL-divergence, Symmetric KL divergence, chi-square divergence, Hellinger distance, and total variation distance, are special cases of f-divergence. Why can't we directly use f-divergence here?

There are numerous other $f$-divergences that could be used for regularization but they do not have theoretical justification and have not been used in prior work.

The claim that these lack theoretical justification or have not been explored in prior work might be somewhat misleading. For reference, here are some important papers on $f$-divergences on RLHF [4][5]. For theoretical justification, see Theorem 1 of [4].

[4] Go, Dongyoung, et al. "Aligning language models with preferences through f-divergence minimization." arXiv preprint arXiv:2302.08215 (2023).

[5] Wang, Chaoqi, et al. "Beyond reverse kl: Generalizing direct preference optimization with diverse divergence constraints." arXiv preprint arXiv:2309.16240 (2023).